@@ -11,6 +11,7 @@ import {
1111 truncateConversationIfNeeded ,
1212} from "../index"
1313import { ApiMessage } from "../../task-persistence/apiMessages"
14+ import * as condenseModule from "../../condense"
1415
1516// Create a mock ApiHandler for testing
1617class MockApiHandler extends BaseProvider {
@@ -248,7 +249,14 @@ describe("truncateConversationIfNeeded", () => {
248249 maxTokens : modelInfo . maxTokens ,
249250 apiHandler : mockApiHandler ,
250251 } )
251- expect ( result ) . toEqual ( messagesWithSmallContent ) // No truncation occurs
252+
253+ // Check the new return type
254+ expect ( result ) . toEqual ( {
255+ messages : messagesWithSmallContent ,
256+ summary : "" ,
257+ cost : 0 ,
258+ prevContextTokens : totalTokens ,
259+ } )
252260 } )
253261
254262 it ( "should truncate if tokens are above max tokens threshold" , async ( ) => {
@@ -260,7 +268,7 @@ describe("truncateConversationIfNeeded", () => {
260268
261269 // When truncating, always uses 0.5 fraction
262270 // With 4 messages after the first, 0.5 fraction means remove 2 messages
263- const expectedResult = [ messagesWithSmallContent [ 0 ] , messagesWithSmallContent [ 3 ] , messagesWithSmallContent [ 4 ] ]
271+ const expectedMessages = [ messagesWithSmallContent [ 0 ] , messagesWithSmallContent [ 3 ] , messagesWithSmallContent [ 4 ] ]
264272
265273 const result = await truncateConversationIfNeeded ( {
266274 messages : messagesWithSmallContent ,
@@ -269,7 +277,13 @@ describe("truncateConversationIfNeeded", () => {
269277 maxTokens : modelInfo . maxTokens ,
270278 apiHandler : mockApiHandler ,
271279 } )
272- expect ( result ) . toEqual ( expectedResult )
280+
281+ expect ( result ) . toEqual ( {
282+ messages : expectedMessages ,
283+ summary : "" ,
284+ cost : 0 ,
285+ prevContextTokens : totalTokens ,
286+ } )
273287 } )
274288
275289 it ( "should work with non-prompt caching models the same as prompt caching models" , async ( ) => {
@@ -298,7 +312,10 @@ describe("truncateConversationIfNeeded", () => {
298312 apiHandler : mockApiHandler ,
299313 } )
300314
301- expect ( result1 ) . toEqual ( result2 )
315+ expect ( result1 . messages ) . toEqual ( result2 . messages )
316+ expect ( result1 . summary ) . toEqual ( result2 . summary )
317+ expect ( result1 . cost ) . toEqual ( result2 . cost )
318+ expect ( result1 . prevContextTokens ) . toEqual ( result2 . prevContextTokens )
302319
303320 // Test above threshold
304321 const aboveThreshold = 70001
@@ -318,7 +335,10 @@ describe("truncateConversationIfNeeded", () => {
318335 apiHandler : mockApiHandler ,
319336 } )
320337
321- expect ( result3 ) . toEqual ( result4 )
338+ expect ( result3 . messages ) . toEqual ( result4 . messages )
339+ expect ( result3 . summary ) . toEqual ( result4 . summary )
340+ expect ( result3 . cost ) . toEqual ( result4 . cost )
341+ expect ( result3 . prevContextTokens ) . toEqual ( result4 . prevContextTokens )
322342 } )
323343
324344 it ( "should consider incoming content when deciding to truncate" , async ( ) => {
@@ -344,7 +364,12 @@ describe("truncateConversationIfNeeded", () => {
344364 maxTokens,
345365 apiHandler : mockApiHandler ,
346366 } )
347- expect ( resultWithSmall ) . toEqual ( messagesWithSmallContent ) // No truncation
367+ expect ( resultWithSmall ) . toEqual ( {
368+ messages : messagesWithSmallContent ,
369+ summary : "" ,
370+ cost : 0 ,
371+ prevContextTokens : baseTokensForSmall + smallContentTokens ,
372+ } ) // No truncation
348373
349374 // Test case 2: Large content that will push us over the threshold
350375 const largeContent = [
@@ -368,7 +393,10 @@ describe("truncateConversationIfNeeded", () => {
368393 maxTokens,
369394 apiHandler : mockApiHandler ,
370395 } )
371- expect ( resultWithLarge ) . not . toEqual ( messagesWithLargeContent ) // Should truncate
396+ expect ( resultWithLarge . messages ) . not . toEqual ( messagesWithLargeContent ) // Should truncate
397+ expect ( resultWithLarge . summary ) . toBe ( "" )
398+ expect ( resultWithLarge . cost ) . toBe ( 0 )
399+ expect ( resultWithLarge . prevContextTokens ) . toBe ( baseTokensForLarge + largeContentTokens )
372400
373401 // Test case 3: Very large content that will definitely exceed threshold
374402 const veryLargeContent = [ { type : "text" as const , text : "X" . repeat ( 1000 ) } ]
@@ -387,7 +415,10 @@ describe("truncateConversationIfNeeded", () => {
387415 maxTokens,
388416 apiHandler : mockApiHandler ,
389417 } )
390- expect ( resultWithVeryLarge ) . not . toEqual ( messagesWithVeryLargeContent ) // Should truncate
418+ expect ( resultWithVeryLarge . messages ) . not . toEqual ( messagesWithVeryLargeContent ) // Should truncate
419+ expect ( resultWithVeryLarge . summary ) . toBe ( "" )
420+ expect ( resultWithVeryLarge . cost ) . toBe ( 0 )
421+ expect ( resultWithVeryLarge . prevContextTokens ) . toBe ( baseTokensForVeryLarge + veryLargeContentTokens )
391422 } )
392423
393424 it ( "should truncate if tokens are within TOKEN_BUFFER_PERCENTAGE of the threshold" , async ( ) => {
@@ -409,7 +440,140 @@ describe("truncateConversationIfNeeded", () => {
409440 maxTokens : modelInfo . maxTokens ,
410441 apiHandler : mockApiHandler ,
411442 } )
412- expect ( result ) . toEqual ( expectedResult )
443+ expect ( result ) . toEqual ( {
444+ messages : expectedResult ,
445+ summary : "" ,
446+ cost : 0 ,
447+ prevContextTokens : totalTokens ,
448+ } )
449+ } )
450+
451+ it ( "should use summarizeConversation when autoCondenseContext is true and tokens exceed threshold" , async ( ) => {
452+ // Mock the summarizeConversation function
453+ const mockSummary = "This is a summary of the conversation"
454+ const mockCost = 0.05
455+ const mockSummarizeResponse : condenseModule . SummarizeResponse = {
456+ messages : [
457+ { role : "user" , content : "First message" } ,
458+ { role : "assistant" , content : mockSummary , isSummary : true } ,
459+ { role : "user" , content : "Last message" } ,
460+ ] ,
461+ summary : mockSummary ,
462+ cost : mockCost ,
463+ newContextTokens : 100 ,
464+ }
465+
466+ const summarizeSpy = jest
467+ . spyOn ( condenseModule , "summarizeConversation" )
468+ . mockResolvedValue ( mockSummarizeResponse )
469+
470+ const modelInfo = createModelInfo ( 100000 , 30000 )
471+ const totalTokens = 70001 // Above threshold
472+ const messagesWithSmallContent = [ ...messages . slice ( 0 , - 1 ) , { ...messages [ messages . length - 1 ] , content : "" } ]
473+
474+ const result = await truncateConversationIfNeeded ( {
475+ messages : messagesWithSmallContent ,
476+ totalTokens,
477+ contextWindow : modelInfo . contextWindow ,
478+ maxTokens : modelInfo . maxTokens ,
479+ apiHandler : mockApiHandler ,
480+ autoCondenseContext : true ,
481+ systemPrompt : "System prompt" ,
482+ } )
483+
484+ // Verify summarizeConversation was called with the right parameters
485+ expect ( summarizeSpy ) . toHaveBeenCalledWith ( messagesWithSmallContent , mockApiHandler , "System prompt" )
486+
487+ // Verify the result contains the summary information
488+ expect ( result ) . toMatchObject ( {
489+ messages : mockSummarizeResponse . messages ,
490+ summary : mockSummary ,
491+ cost : mockCost ,
492+ prevContextTokens : totalTokens ,
493+ } )
494+ // newContextTokens might be present, but we don't need to verify its exact value
495+
496+ // Clean up
497+ summarizeSpy . mockRestore ( )
498+ } )
499+
500+ it ( "should fall back to truncateConversation when autoCondenseContext is true but summarization fails" , async ( ) => {
501+ // Mock the summarizeConversation function to return empty summary
502+ const mockSummarizeResponse : condenseModule . SummarizeResponse = {
503+ messages : messages , // Original messages unchanged
504+ summary : "" , // Empty summary indicates failure
505+ cost : 0.01 ,
506+ }
507+
508+ const summarizeSpy = jest
509+ . spyOn ( condenseModule , "summarizeConversation" )
510+ . mockResolvedValue ( mockSummarizeResponse )
511+
512+ const modelInfo = createModelInfo ( 100000 , 30000 )
513+ const totalTokens = 70001 // Above threshold
514+ const messagesWithSmallContent = [ ...messages . slice ( 0 , - 1 ) , { ...messages [ messages . length - 1 ] , content : "" } ]
515+
516+ // When truncating, always uses 0.5 fraction
517+ // With 4 messages after the first, 0.5 fraction means remove 2 messages
518+ const expectedMessages = [ messagesWithSmallContent [ 0 ] , messagesWithSmallContent [ 3 ] , messagesWithSmallContent [ 4 ] ]
519+
520+ const result = await truncateConversationIfNeeded ( {
521+ messages : messagesWithSmallContent ,
522+ totalTokens,
523+ contextWindow : modelInfo . contextWindow ,
524+ maxTokens : modelInfo . maxTokens ,
525+ apiHandler : mockApiHandler ,
526+ autoCondenseContext : true ,
527+ } )
528+
529+ // Verify summarizeConversation was called
530+ expect ( summarizeSpy ) . toHaveBeenCalled ( )
531+
532+ // Verify it fell back to truncation
533+ expect ( result . messages ) . toEqual ( expectedMessages )
534+ expect ( result . summary ) . toBe ( "" )
535+ expect ( result . prevContextTokens ) . toBe ( totalTokens )
536+ // The cost might be different than expected, so we don't check it
537+
538+ // Clean up
539+ summarizeSpy . mockRestore ( )
540+ } )
541+
542+ it ( "should not call summarizeConversation when autoCondenseContext is false" , async ( ) => {
543+ // Reset any previous mock calls
544+ jest . clearAllMocks ( )
545+ const summarizeSpy = jest . spyOn ( condenseModule , "summarizeConversation" )
546+
547+ const modelInfo = createModelInfo ( 100000 , 30000 )
548+ const totalTokens = 70001 // Above threshold
549+ const messagesWithSmallContent = [ ...messages . slice ( 0 , - 1 ) , { ...messages [ messages . length - 1 ] , content : "" } ]
550+
551+ // When truncating, always uses 0.5 fraction
552+ // With 4 messages after the first, 0.5 fraction means remove 2 messages
553+ const expectedMessages = [ messagesWithSmallContent [ 0 ] , messagesWithSmallContent [ 3 ] , messagesWithSmallContent [ 4 ] ]
554+
555+ const result = await truncateConversationIfNeeded ( {
556+ messages : messagesWithSmallContent ,
557+ totalTokens,
558+ contextWindow : modelInfo . contextWindow ,
559+ maxTokens : modelInfo . maxTokens ,
560+ apiHandler : mockApiHandler ,
561+ autoCondenseContext : false ,
562+ } )
563+
564+ // Verify summarizeConversation was not called
565+ expect ( summarizeSpy ) . not . toHaveBeenCalled ( )
566+
567+ // Verify it used truncation
568+ expect ( result ) . toEqual ( {
569+ messages : expectedMessages ,
570+ summary : "" ,
571+ cost : 0 ,
572+ prevContextTokens : totalTokens ,
573+ } )
574+
575+ // Clean up
576+ summarizeSpy . mockRestore ( )
413577 } )
414578} )
415579
@@ -449,7 +613,12 @@ describe("getMaxTokens", () => {
449613 maxTokens : modelInfo . maxTokens ,
450614 apiHandler : mockApiHandler ,
451615 } )
452- expect ( result1 ) . toEqual ( messagesWithSmallContent )
616+ expect ( result1 ) . toEqual ( {
617+ messages : messagesWithSmallContent ,
618+ summary : "" ,
619+ cost : 0 ,
620+ prevContextTokens : 39999 ,
621+ } )
453622
454623 // Above max tokens - truncate
455624 const result2 = await truncateConversationIfNeeded ( {
@@ -459,8 +628,11 @@ describe("getMaxTokens", () => {
459628 maxTokens : modelInfo . maxTokens ,
460629 apiHandler : mockApiHandler ,
461630 } )
462- expect ( result2 ) . not . toEqual ( messagesWithSmallContent )
463- expect ( result2 . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
631+ expect ( result2 . messages ) . not . toEqual ( messagesWithSmallContent )
632+ expect ( result2 . messages . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
633+ expect ( result2 . summary ) . toBe ( "" )
634+ expect ( result2 . cost ) . toBe ( 0 )
635+ expect ( result2 . prevContextTokens ) . toBe ( 50001 )
464636 } )
465637
466638 it ( "should use 20% of context window as buffer when maxTokens is undefined" , async ( ) => {
@@ -479,7 +651,12 @@ describe("getMaxTokens", () => {
479651 maxTokens : modelInfo . maxTokens ,
480652 apiHandler : mockApiHandler ,
481653 } )
482- expect ( result1 ) . toEqual ( messagesWithSmallContent )
654+ expect ( result1 ) . toEqual ( {
655+ messages : messagesWithSmallContent ,
656+ summary : "" ,
657+ cost : 0 ,
658+ prevContextTokens : 69999 ,
659+ } )
483660
484661 // Above max tokens - truncate
485662 const result2 = await truncateConversationIfNeeded ( {
@@ -489,8 +666,11 @@ describe("getMaxTokens", () => {
489666 maxTokens : modelInfo . maxTokens ,
490667 apiHandler : mockApiHandler ,
491668 } )
492- expect ( result2 ) . not . toEqual ( messagesWithSmallContent )
493- expect ( result2 . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
669+ expect ( result2 . messages ) . not . toEqual ( messagesWithSmallContent )
670+ expect ( result2 . messages . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
671+ expect ( result2 . summary ) . toBe ( "" )
672+ expect ( result2 . cost ) . toBe ( 0 )
673+ expect ( result2 . prevContextTokens ) . toBe ( 80001 )
494674 } )
495675
496676 it ( "should handle small context windows appropriately" , async ( ) => {
@@ -508,7 +688,7 @@ describe("getMaxTokens", () => {
508688 maxTokens : modelInfo . maxTokens ,
509689 apiHandler : mockApiHandler ,
510690 } )
511- expect ( result1 ) . toEqual ( messagesWithSmallContent )
691+ expect ( result1 . messages ) . toEqual ( messagesWithSmallContent )
512692
513693 // Above max tokens - truncate
514694 const result2 = await truncateConversationIfNeeded ( {
@@ -519,7 +699,7 @@ describe("getMaxTokens", () => {
519699 apiHandler : mockApiHandler ,
520700 } )
521701 expect ( result2 ) . not . toEqual ( messagesWithSmallContent )
522- expect ( result2 . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
702+ expect ( result2 . messages . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
523703 } )
524704
525705 it ( "should handle large context windows appropriately" , async ( ) => {
@@ -538,7 +718,7 @@ describe("getMaxTokens", () => {
538718 maxTokens : modelInfo . maxTokens ,
539719 apiHandler : mockApiHandler ,
540720 } )
541- expect ( result1 ) . toEqual ( messagesWithSmallContent )
721+ expect ( result1 . messages ) . toEqual ( messagesWithSmallContent )
542722
543723 // Above max tokens - truncate
544724 const result2 = await truncateConversationIfNeeded ( {
@@ -549,6 +729,6 @@ describe("getMaxTokens", () => {
549729 apiHandler : mockApiHandler ,
550730 } )
551731 expect ( result2 ) . not . toEqual ( messagesWithSmallContent )
552- expect ( result2 . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
732+ expect ( result2 . messages . length ) . toBe ( 3 ) // Truncated with 0.5 fraction
553733 } )
554734} )
0 commit comments