@@ -147,4 +147,219 @@ describe('GenericBlockHandler', () => {
147147 'Block execution of Some Custom Tool failed with no error message'
148148 )
149149 } )
150+
151+ describe ( 'Knowledge block cost tracking' , ( ) => {
152+ beforeEach ( ( ) => {
153+ // Set up knowledge block mock
154+ mockBlock = {
155+ ...mockBlock ,
156+ config : { tool : 'knowledge_search' , params : { } } ,
157+ }
158+
159+ mockTool = {
160+ ...mockTool ,
161+ id : 'knowledge_search' ,
162+ name : 'Knowledge Search' ,
163+ }
164+
165+ mockGetTool . mockImplementation ( ( toolId ) => {
166+ if ( toolId === 'knowledge_search' ) {
167+ return mockTool
168+ }
169+ return undefined
170+ } )
171+ } )
172+
173+ it . concurrent (
174+ 'should extract and restructure cost information from knowledge tools' ,
175+ async ( ) => {
176+ const inputs = { query : 'test query' }
177+ const mockToolResponse = {
178+ success : true ,
179+ output : {
180+ results : [ ] ,
181+ query : 'test query' ,
182+ totalResults : 0 ,
183+ cost : {
184+ input : 0.00001042 ,
185+ output : 0 ,
186+ total : 0.00001042 ,
187+ tokens : {
188+ input : 521 ,
189+ output : 0 ,
190+ total : 521 ,
191+ } ,
192+ model : 'text-embedding-3-small' ,
193+ pricing : {
194+ input : 0.02 ,
195+ output : 0 ,
196+ updatedAt : '2025-07-10' ,
197+ } ,
198+ } ,
199+ } ,
200+ }
201+
202+ mockExecuteTool . mockResolvedValue ( mockToolResponse )
203+
204+ const result = await handler . execute ( mockContext , mockBlock , inputs )
205+
206+ // Verify cost information is restructured correctly for enhanced logging
207+ expect ( result ) . toEqual ( {
208+ results : [ ] ,
209+ query : 'test query' ,
210+ totalResults : 0 ,
211+ cost : {
212+ input : 0.00001042 ,
213+ output : 0 ,
214+ total : 0.00001042 ,
215+ } ,
216+ tokens : {
217+ input : 521 ,
218+ output : 0 ,
219+ total : 521 ,
220+ } ,
221+ model : 'text-embedding-3-small' ,
222+ } )
223+ }
224+ )
225+
226+ it . concurrent ( 'should handle knowledge_upload_chunk cost information' , async ( ) => {
227+ // Update to upload_chunk tool
228+ mockBlock . config . tool = 'knowledge_upload_chunk'
229+ mockTool . id = 'knowledge_upload_chunk'
230+ mockTool . name = 'Knowledge Upload Chunk'
231+
232+ mockGetTool . mockImplementation ( ( toolId ) => {
233+ if ( toolId === 'knowledge_upload_chunk' ) {
234+ return mockTool
235+ }
236+ return undefined
237+ } )
238+
239+ const inputs = { content : 'test content' }
240+ const mockToolResponse = {
241+ success : true ,
242+ output : {
243+ data : {
244+ id : 'chunk-123' ,
245+ content : 'test content' ,
246+ chunkIndex : 0 ,
247+ } ,
248+ message : 'Successfully uploaded chunk' ,
249+ documentId : 'doc-123' ,
250+ cost : {
251+ input : 0.00000521 ,
252+ output : 0 ,
253+ total : 0.00000521 ,
254+ tokens : {
255+ input : 260 ,
256+ output : 0 ,
257+ total : 260 ,
258+ } ,
259+ model : 'text-embedding-3-small' ,
260+ pricing : {
261+ input : 0.02 ,
262+ output : 0 ,
263+ updatedAt : '2025-07-10' ,
264+ } ,
265+ } ,
266+ } ,
267+ }
268+
269+ mockExecuteTool . mockResolvedValue ( mockToolResponse )
270+
271+ const result = await handler . execute ( mockContext , mockBlock , inputs )
272+
273+ // Verify cost information is restructured correctly
274+ expect ( result ) . toEqual ( {
275+ data : {
276+ id : 'chunk-123' ,
277+ content : 'test content' ,
278+ chunkIndex : 0 ,
279+ } ,
280+ message : 'Successfully uploaded chunk' ,
281+ documentId : 'doc-123' ,
282+ cost : {
283+ input : 0.00000521 ,
284+ output : 0 ,
285+ total : 0.00000521 ,
286+ } ,
287+ tokens : {
288+ input : 260 ,
289+ output : 0 ,
290+ total : 260 ,
291+ } ,
292+ model : 'text-embedding-3-small' ,
293+ } )
294+ } )
295+
296+ it ( 'should pass through output unchanged for knowledge tools without cost info' , async ( ) => {
297+ const inputs = { query : 'test query' }
298+ const mockToolResponse = {
299+ success : true ,
300+ output : {
301+ results : [ ] ,
302+ query : 'test query' ,
303+ totalResults : 0 ,
304+ // No cost information
305+ } ,
306+ }
307+
308+ mockExecuteTool . mockResolvedValue ( mockToolResponse )
309+
310+ const result = await handler . execute ( mockContext , mockBlock , inputs )
311+
312+ // Should return original output without cost transformation
313+ expect ( result ) . toEqual ( {
314+ results : [ ] ,
315+ query : 'test query' ,
316+ totalResults : 0 ,
317+ } )
318+ } )
319+
320+ it . concurrent (
321+ 'should process cost info for all tools (universal cost extraction)' ,
322+ async ( ) => {
323+ mockBlock . config . tool = 'some_other_tool'
324+ mockTool . id = 'some_other_tool'
325+
326+ mockGetTool . mockImplementation ( ( toolId ) => {
327+ if ( toolId === 'some_other_tool' ) {
328+ return mockTool
329+ }
330+ return undefined
331+ } )
332+
333+ const inputs = { param : 'value' }
334+ const mockToolResponse = {
335+ success : true ,
336+ output : {
337+ result : 'success' ,
338+ cost : {
339+ input : 0.001 ,
340+ output : 0.002 ,
341+ total : 0.003 ,
342+ tokens : { input : 100 , output : 50 , total : 150 } ,
343+ model : 'some-model' ,
344+ } ,
345+ } ,
346+ }
347+
348+ mockExecuteTool . mockResolvedValue ( mockToolResponse )
349+
350+ const result = await handler . execute ( mockContext , mockBlock , inputs )
351+
352+ expect ( result ) . toEqual ( {
353+ result : 'success' ,
354+ cost : {
355+ input : 0.001 ,
356+ output : 0.002 ,
357+ total : 0.003 ,
358+ } ,
359+ tokens : { input : 100 , output : 50 , total : 150 } ,
360+ model : 'some-model' ,
361+ } )
362+ }
363+ )
364+ } )
150365} )
0 commit comments