@@ -60,7 +60,13 @@ vi.mock("vscode", () => {
6060} )
6161
6262import * as vscode from "vscode"
63- import { VsCodeLmHandler } from "../vscode-lm"
63+ import {
64+ VsCodeLmHandler ,
65+ extractLeakedToolCalls ,
66+ trailingPartialToolMarkerLength ,
67+ middleOutTruncate ,
68+ truncateToolResultsToFitWindow ,
69+ } from "../vscode-lm"
6470import type { ApiHandlerOptions } from "../../../shared/api"
6571import type { Anthropic } from "@anthropic-ai/sdk"
6672import { openAiModelInfoSaneDefaults , vscodeLlmDefaultModelId , vscodeLlmModels } from "@roo-code/types"
@@ -1075,3 +1081,192 @@ describe("VsCodeLmHandler", () => {
10751081 } )
10761082 } )
10771083} )
1084+
1085+ describe ( "leaked tool-call recovery" , ( ) => {
1086+ // Builders keep the XML fixtures readable and prevent this file's own markup from being
1087+ // mistaken for a real tool call.
1088+ const invoke = ( name : string , body : string ) => `<in${ "voke" } name="${ name } ">${ body } </in${ "voke" } >`
1089+ const param = ( name : string , value : string ) => `<param${ "eter" } name="${ name } ">${ value } </param${ "eter" } >`
1090+
1091+ describe ( "extractLeakedToolCalls" , ( ) => {
1092+ it ( "recovers a known-tool block and strips it from the leftover text" , ( ) => {
1093+ const text = `Working on it.\n${ invoke ( "update_todo_list" , param ( "todos" , "[x] one\n[ ] two" ) ) } `
1094+
1095+ const { calls, leftoverText } = extractLeakedToolCalls ( text , new Set ( [ "update_todo_list" ] ) )
1096+
1097+ expect ( calls ) . toEqual ( [ { name : "update_todo_list" , input : { todos : "[x] one\n[ ] two" } } ] )
1098+ expect ( leftoverText ) . toBe ( "Working on it.\n" )
1099+ } )
1100+
1101+ it ( "recovers an unwrapped leak preceded by a stray token" , ( ) => {
1102+ const text = `court\n${ invoke ( "update_todo_list" , param ( "todos" , "[x] done" ) ) } `
1103+
1104+ const { calls, leftoverText } = extractLeakedToolCalls ( text , new Set ( [ "update_todo_list" ] ) )
1105+
1106+ expect ( calls ) . toEqual ( [ { name : "update_todo_list" , input : { todos : "[x] done" } } ] )
1107+ expect ( leftoverText ) . toBe ( "court\n" )
1108+ } )
1109+
1110+ it ( "recovers multiple params and strips function-call wrapper tags" , ( ) => {
1111+ const body = param ( "mode" , "code" ) + param ( "message" , "go" )
1112+ const text = `<function_calls>${ invoke ( "new_task" , body ) } </function_calls>`
1113+
1114+ const { calls, leftoverText } = extractLeakedToolCalls ( text , new Set ( [ "new_task" ] ) )
1115+
1116+ expect ( calls ) . toEqual ( [ { name : "new_task" , input : { mode : "code" , message : "go" } } ] )
1117+ expect ( leftoverText ) . toBe ( "" )
1118+ } )
1119+
1120+ it ( "passes through invoke blocks for tools that were not offered" , ( ) => {
1121+ const text = invoke ( "some_other_tool" , param ( "x" , "1" ) )
1122+
1123+ const { calls, leftoverText } = extractLeakedToolCalls ( text , new Set ( [ "update_todo_list" ] ) )
1124+
1125+ expect ( calls ) . toEqual ( [ ] )
1126+ expect ( leftoverText ) . toBe ( text )
1127+ } )
1128+
1129+ it ( "returns no calls for ordinary text" , ( ) => {
1130+ const { calls, leftoverText } = extractLeakedToolCalls ( "just a normal reply" , new Set ( [ "update_todo_list" ] ) )
1131+
1132+ expect ( calls ) . toEqual ( [ ] )
1133+ expect ( leftoverText ) . toBe ( "just a normal reply" )
1134+ } )
1135+ } )
1136+
1137+ describe ( "trailingPartialToolMarkerLength" , ( ) => {
1138+ it ( "holds back a split marker prefix at the end of a chunk" , ( ) => {
1139+ expect ( trailingPartialToolMarkerLength ( "some text <in" ) ) . toBe ( 3 )
1140+ } )
1141+
1142+ it ( "returns 0 for plain text and complete tags" , ( ) => {
1143+ expect ( trailingPartialToolMarkerLength ( "hello world" ) ) . toBe ( 0 )
1144+ expect ( trailingPartialToolMarkerLength ( "a < b" ) ) . toBe ( 0 )
1145+ expect ( trailingPartialToolMarkerLength ( "text <function_calls>" ) ) . toBe ( 0 )
1146+ } )
1147+ } )
1148+ } )
1149+
1150+ describe ( "context-window tool_result truncation" , ( ) => {
1151+ describe ( "middleOutTruncate" , ( ) => {
1152+ it ( "returns text unchanged when within the limit" , ( ) => {
1153+ expect ( middleOutTruncate ( "hello world" , 100 ) ) . toBe ( "hello world" )
1154+ } )
1155+
1156+ it ( "keeps the head and tail and inserts a truncation marker" , ( ) => {
1157+ const text = "A" . repeat ( 500 ) + "B" . repeat ( 500 )
1158+ const result = middleOutTruncate ( text , 200 )
1159+
1160+ expect ( result . length ) . toBeLessThanOrEqual ( 200 )
1161+ expect ( result ) . toContain ( "characters truncated to fit the model context window" )
1162+ expect ( result . startsWith ( "A" ) ) . toBe ( true )
1163+ expect ( result . endsWith ( "B" ) ) . toBe ( true )
1164+ } )
1165+
1166+ it ( "returns an empty string for a non-positive limit" , ( ) => {
1167+ expect ( middleOutTruncate ( "anything" , 0 ) ) . toBe ( "" )
1168+ } )
1169+ } )
1170+
1171+ describe ( "truncateToolResultsToFitWindow" , ( ) => {
1172+ const toolUseMessage = ( id : string ) : Anthropic . Messages . MessageParam => ( {
1173+ role : "assistant" ,
1174+ content : [
1175+ { type : "text" , text : "Calling a tool." } ,
1176+ { type : "tool_use" , id, name : "some_tool" , input : { a : 1 } } ,
1177+ ] ,
1178+ } )
1179+
1180+ const toolResultMessage = ( id : string , content : string ) : Anthropic . Messages . MessageParam => ( {
1181+ role : "user" ,
1182+ content : [
1183+ { type : "tool_result" , tool_use_id : id , content } ,
1184+ { type : "text" , text : "<environment_details>env</environment_details>" } ,
1185+ ] ,
1186+ } )
1187+
1188+ const findBlock = ( message : Anthropic . Messages . MessageParam , type : string ) =>
1189+ ( message . content as unknown as Array < { type : string ; [ key : string ] : unknown } > ) . find (
1190+ ( block ) => block . type === type ,
1191+ ) !
1192+
1193+ it ( "is a no-op when the conversation already fits the budget" , ( ) => {
1194+ const messages : Anthropic . Messages . MessageParam [ ] = [
1195+ toolUseMessage ( "t1" ) ,
1196+ toolResultMessage ( "t1" , "small result" ) ,
1197+ ]
1198+ const before = JSON . parse ( JSON . stringify ( messages ) )
1199+
1200+ truncateToolResultsToFitWindow ( messages , 100_000 )
1201+
1202+ expect ( messages ) . toEqual ( before )
1203+ } )
1204+
1205+ it ( "shrinks an oversized tool_result so the conversation fits the budget" , ( ) => {
1206+ const messages : Anthropic . Messages . MessageParam [ ] = [
1207+ toolUseMessage ( "t1" ) ,
1208+ toolResultMessage ( "t1" , "X" . repeat ( 50_000 ) ) ,
1209+ ]
1210+
1211+ truncateToolResultsToFitWindow ( messages , 10_000 )
1212+
1213+ const toolResult = findBlock ( messages [ 1 ] , "tool_result" )
1214+ expect ( toolResult . tool_use_id ) . toBe ( "t1" ) // pairing preserved
1215+ expect ( String ( toolResult . content ) . length ) . toBeLessThanOrEqual ( 10_000 )
1216+ expect ( String ( toolResult . content ) ) . toContain ( "characters truncated" )
1217+ } )
1218+
1219+ it ( "truncates the largest tool_result first and leaves small ones intact" , ( ) => {
1220+ const small = "small but real result"
1221+ const messages : Anthropic . Messages . MessageParam [ ] = [
1222+ toolUseMessage ( "t1" ) ,
1223+ toolResultMessage ( "t1" , "B" . repeat ( 40_000 ) ) ,
1224+ toolUseMessage ( "t2" ) ,
1225+ toolResultMessage ( "t2" , small ) ,
1226+ ]
1227+
1228+ truncateToolResultsToFitWindow ( messages , 12_000 )
1229+
1230+ expect ( String ( findBlock ( messages [ 1 ] , "tool_result" ) . content ) ) . toContain ( "characters truncated" )
1231+ expect ( findBlock ( messages [ 3 ] , "tool_result" ) . content ) . toBe ( small ) // untouched
1232+ } )
1233+
1234+ it ( "never truncates tool_use blocks, assistant text, or environment details" , ( ) => {
1235+ const messages : Anthropic . Messages . MessageParam [ ] = [
1236+ toolUseMessage ( "t1" ) ,
1237+ toolResultMessage ( "t1" , "X" . repeat ( 50_000 ) ) ,
1238+ ]
1239+
1240+ truncateToolResultsToFitWindow ( messages , 8_000 )
1241+
1242+ expect ( findBlock ( messages [ 0 ] , "text" ) . text ) . toBe ( "Calling a tool." )
1243+ expect ( findBlock ( messages [ 0 ] , "tool_use" ) ) . toMatchObject ( { id : "t1" , name : "some_tool" } )
1244+ expect ( findBlock ( messages [ 1 ] , "text" ) . text ) . toBe ( "<environment_details>env</environment_details>" )
1245+ } )
1246+
1247+ it ( "handles array-form tool_result content and keeps it valid" , ( ) => {
1248+ const messages : Anthropic . Messages . MessageParam [ ] = [
1249+ toolUseMessage ( "t1" ) ,
1250+ {
1251+ role : "user" ,
1252+ content : [
1253+ {
1254+ type : "tool_result" ,
1255+ tool_use_id : "t1" ,
1256+ content : [ { type : "text" , text : "Y" . repeat ( 40_000 ) } ] ,
1257+ } ,
1258+ ] ,
1259+ } ,
1260+ ]
1261+
1262+ truncateToolResultsToFitWindow ( messages , 8_000 )
1263+
1264+ const toolResult = findBlock ( messages [ 1 ] , "tool_result" )
1265+ expect ( toolResult . tool_use_id ) . toBe ( "t1" )
1266+ expect ( Array . isArray ( toolResult . content ) ) . toBe ( true )
1267+ const parts = toolResult . content as Array < { type : string ; text ?: string } >
1268+ expect ( parts [ 0 ] . type ) . toBe ( "text" )
1269+ expect ( String ( parts [ 0 ] . text ) ) . toContain ( "characters truncated" )
1270+ } )
1271+ } )
1272+ } )
0 commit comments