Skip to content

Commit b4e1727

Browse files
author
Bertan Ari
committed
fix(vscode-lm): sanitize surrogates, recover leaked tool calls, and window-safe tool_result truncation
Hardens the VS Code Language Model provider (notably GitHub Copilot serving Anthropic Claude) against three failure modes: - Surrogate sanitization: a lone UTF-16 surrogate cannot be encoded as UTF-8, so the backend rejects the entire request with a 400. sanitizeSurrogates() replaces unpaired surrogates with U+FFFD while preserving valid pairs (emoji, CJK ext.), applied to string messages, tool results, and text parts. - Leaked tool-call recovery: some backends stream a tool call as raw <invoke> XML instead of a structured LanguageModelToolCallPart, leaving the turn with no tool_use block and stalling the task in a "no tools used" retry loop. extractLeakedToolCalls() and trailingPartialToolMarkerLength() detect the markup mid-stream (including markers split across chunk boundaries) and replay it as a real tool call, conservatively: only for <invoke> names matching a tool actually offered that turn, and only when tools were offered. - Window-safe tool_result truncation: Copilot's backend trims over-window requests without preserving tool_use/tool_result pairing, orphaning a tool_result and causing a 400 (unexpected tool_use_id). truncateToolResultsToFitWindow() and middleOutTruncate() shrink oversized tool_result payloads on our side (largest first, middle-out, pairing preserved) before sending. Ported from simurg79/Roo-Code#12.
1 parent 276e425 commit b4e1727

4 files changed

Lines changed: 594 additions & 11 deletions

File tree

‎src/api/providers/__tests__/vscode-lm.spec.ts‎

Lines changed: 196 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,13 @@ vi.mock("vscode", () => {
6060
})
6161

6262
import * 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"
6470
import type { ApiHandlerOptions } from "../../../shared/api"
6571
import type { Anthropic } from "@anthropic-ai/sdk"
6672
import { 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

Comments
 (0)