diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 37fb851720..5fc4051766 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -16,6 +16,14 @@ vi.mock("vscode", () => { ) {} } + class MockLanguageModelToolResultPart { + type = "tool_result" + constructor( + public callId: string, + public content: unknown[], + ) {} + } + return { workspace: { getConfiguration: vi.fn(() => ({ @@ -53,6 +61,7 @@ vi.mock("vscode", () => { }, LanguageModelTextPart: MockLanguageModelTextPart, LanguageModelToolCallPart: MockLanguageModelToolCallPart, + LanguageModelToolResultPart: MockLanguageModelToolResultPart, lm: { selectChatModels: vi.fn(), }, @@ -60,7 +69,13 @@ vi.mock("vscode", () => { }) import * as vscode from "vscode" -import { VsCodeLmHandler } from "../vscode-lm" +import { + VsCodeLmHandler, + extractLeakedToolCalls, + trailingPartialToolMarkerLength, + middleOutTruncate, + truncateToolResultsToFitWindow, +} from "../vscode-lm" import type { ApiHandlerOptions } from "../../../shared/api" import type { Anthropic } from "@anthropic-ai/sdk" import { openAiModelInfoSaneDefaults, vscodeLlmDefaultModelId, vscodeLlmModels } from "@roo-code/types" @@ -270,6 +285,257 @@ describe("VsCodeLmHandler", () => { }) }) + it("still trims oversized tool_results when the system prompt consumes the whole budget", async () => { + // A system prompt larger than the derived char budget drives the raw budget negative; the + // clamp keeps trimming active for the case where the request is most oversized. + const systemPrompt = "S".repeat(handler.getCondenseContextWindow() * 3) + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "assistant", + content: [{ type: "tool_use", id: "t1", name: "some_tool", input: { a: 1 } }], + }, + { + role: "user", + content: [{ type: "tool_result", tool_use_id: "t1", content: "X".repeat(50_000) }], + }, + ] + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("ok") + return + })(), + text: (async function* () { + yield "ok" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages, { taskId: "test-task" }) + for await (const _chunk of stream) { + // drain + } + + const sent = JSON.stringify(mockLanguageModelChat.sendRequest.mock.calls[0][0]) + expect(sent).toContain("characters truncated") + expect(sent).not.toContain("X".repeat(10_000)) + }) + + describe("leaked tool-call recovery during streaming", () => { + const salvageTools = [ + { + type: "function" as const, + function: { + name: "calculator", + description: "A simple calculator", + parameters: { type: "object", properties: { operation: { type: "string" } } }, + }, + }, + ] + + const streamTextParts = (parts: string[]) => { + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + for (const part of parts) { + yield new vscode.LanguageModelTextPart(part) + } + return + })(), + text: (async function* () { + yield parts.join("") + return + })(), + }) + } + + const streamMixedParts = (parts: Array) => { + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + for (const part of parts) { + yield typeof part === "string" + ? new vscode.LanguageModelTextPart(part) + : new vscode.LanguageModelToolCallPart("native-1", part.name, part.input) + } + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + } + + const drain = async () => { + const stream = handler.createMessage("system", [{ role: "user" as const, content: "hi" }], { + taskId: "test-task", + tools: salvageTools, + }) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + return chunks + } + + const collect = async (parts: string[]) => { + streamTextParts(parts) + return drain() + } + + it("recovers a tool call the model streamed as raw invoke XML", async () => { + const chunks = await collect([ + "Thinking. ", + 'add', + ]) + + expect(chunks.filter((chunk) => chunk.type === "text")).toEqual([{ type: "text", text: "Thinking. " }]) + expect(chunks.filter((chunk) => chunk.type === "tool_call")).toEqual([ + { + type: "tool_call", + id: expect.stringContaining("vscodelm-salvaged-"), + name: "calculator", + arguments: JSON.stringify({ operation: "add" }), + }, + ]) + }) + + it("detects a marker split across stream chunks", async () => { + const chunks = await collect([ + "abc sub', + ]) + + expect(chunks.filter((chunk) => chunk.type === "text")).toEqual([{ type: "text", text: "abc " }]) + expect(chunks.filter((chunk) => chunk.type === "tool_call")).toMatchObject([ + { name: "calculator", arguments: JSON.stringify({ operation: "sub" }) }, + ]) + }) + + it("emits a carried tail as plain text when it never becomes a marker", async () => { + const chunks = await collect(["hello chunk.type === "text")).toEqual([ + { type: "text", text: "hello " }, + { type: "text", text: " chunk.type === "tool_call")).toBe(false) + }) + + it("buffers across chunks that arrive after the marker", async () => { + const chunks = await collect([ + 'prose ', + '', + "mul", + "", + ]) + + expect(chunks.filter((chunk) => chunk.type === "text")).toEqual([{ type: "text", text: "prose " }]) + expect(chunks.filter((chunk) => chunk.type === "tool_call")).toMatchObject([ + { name: "calculator", arguments: JSON.stringify({ operation: "mul" }) }, + ]) + }) + + it("keeps an invoke block for an unknown tool as literal text", async () => { + const block = '1' + const chunks = await collect([block]) + + expect(chunks.filter((chunk) => chunk.type === "text")).toEqual([{ type: "text", text: block }]) + expect(chunks.some((chunk) => chunk.type === "tool_call")).toBe(false) + }) + + it("emits prose before the recovered tool call", async () => { + const chunks = await collect([ + "Thinking. ", + 'add', + ]) + + expect(chunks.map((chunk) => chunk.type)).toEqual(["text", "tool_call", "usage"]) + }) + + it("flushes buffered text before a native tool call so no text follows a tool_use", async () => { + streamMixedParts([ + 'partial ', + { name: "calculator", input: { operation: "div" } }, + ]) + const chunks = await drain() + + const lastText = chunks.map((chunk) => chunk.type).lastIndexOf("text") + const firstToolCall = chunks.map((chunk) => chunk.type).indexOf("tool_call") + expect(firstToolCall).toBeGreaterThan(lastText) + }) + + it("does not latch buffering on prose that merely mentions the tag", async () => { + const chunks = await collect(["never emit markup as text. ", "Streaming continues."]) + + expect(chunks.filter((chunk) => chunk.type === "text")).toEqual([ + { type: "text", text: "never emit markup as text. " }, + { type: "text", text: "Streaming continues." }, + ]) + expect(chunks.some((chunk) => chunk.type === "tool_call")).toBe(false) + }) + + it("does not recover an invoke block quoted inside a fenced code block", async () => { + const block = 'add' + const chunks = await collect(["Do NOT do this:\n```\n" + block + "\n```\n"]) + + expect(chunks.some((chunk) => chunk.type === "tool_call")).toBe(false) + }) + + it("flushes an over-long never-closing invoke as plain text before the stream ends", async () => { + // Defect 4: without a cap the buffer is only drained once the stream finishes, so the + // user sees nothing until then. Releasing it at the end looks identical in content — + // only the timing distinguishes the fix, so track how much of the source has been + // produced at the moment each text chunk reaches the consumer. + const filler = "x".repeat(5000) + const parts = ['', filler, filler, filler, filler] + let partsProduced = 0 + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + for (const part of parts) { + partsProduced++ + yield new vscode.LanguageModelTextPart(part) + } + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage("system", [{ role: "user" as const, content: "hi" }], { + taskId: "test-task", + tools: salvageTools, + }) + + let sawTextBeforeStreamEnd = false + let streamedText = "" + for await (const chunk of stream) { + if (chunk.type === "text") { + streamedText += chunk.text + if (partsProduced < parts.length) { + sawTextBeforeStreamEnd = true + } + } + } + + expect(sawTextBeforeStreamEnd).toBe(true) + expect(streamedText).toContain('') + expect(streamedText).toContain(filler) + }) + + it("sanitizes lone surrogates in the system prompt", async () => { + streamTextParts(["ok"]) + const stream = handler.createMessage("sys\uD800tem", [{ role: "user" as const, content: "hi" }]) + for await (const _chunk of stream) { + // drain + } + + expect(vscode.LanguageModelChatMessage.Assistant).toHaveBeenCalledWith("sys\uFFFDtem") + }) + }) + it("should handle native tool calls when tools are provided", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [ @@ -1075,3 +1341,430 @@ describe("VsCodeLmHandler", () => { }) }) }) + +describe("leaked tool-call recovery", () => { + // Builders keep the XML fixtures readable and prevent this file's own markup from being + // mistaken for a real tool call. + const invoke = (name: string, body: string) => `${body}` + const param = (name: string, value: string) => `${value}` + const wrap = (body: string) => `${body}` + + describe("extractLeakedToolCalls", () => { + it("recovers a known-tool block and strips it from the leftover text", () => { + const text = `Working on it.\n${wrap(invoke("update_todo_list", param("todos", "[x] one\n[ ] two")))}` + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toEqual([{ name: "update_todo_list", input: { todos: "[x] one\n[ ] two" } }]) + expect(leftoverText).toBe("Working on it.\n") + }) + + it("recovers a wrapped leak preceded by a stray token", () => { + const text = `court\n${wrap(invoke("update_todo_list", param("todos", "[x] done")))}` + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toEqual([{ name: "update_todo_list", input: { todos: "[x] done" } }]) + expect(leftoverText).toBe("court\n") + }) + + it("does not recover a bare invoke block with no function_calls wrapper", () => { + const text = `court\n${invoke("update_todo_list", param("todos", "[x] done"))}` + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("does not recover an invoke that follows an already-closed wrapper", () => { + const text = `${wrap("")}\n${invoke("update_todo_list", param("todos", "[x] done"))}` + + const { calls } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + }) + + it("recovers multiple params and strips function-call wrapper tags", () => { + const body = param("mode", "code") + param("message", "go") + const text = `${invoke("new_task", body)}` + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["new_task"])) + + expect(calls).toEqual([{ name: "new_task", input: { mode: "code", message: "go" } }]) + expect(leftoverText).toBe("") + }) + + it("passes through invoke blocks for tools that were not offered", () => { + const text = invoke("some_other_tool", param("x", "1")) + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toEqual([]) + expect(leftoverText).toBe(text) + }) + + it("returns no calls for ordinary text", () => { + const { calls, leftoverText } = extractLeakedToolCalls("just a normal reply", new Set(["update_todo_list"])) + + expect(calls).toEqual([]) + expect(leftoverText).toBe("just a normal reply") + }) + }) + + describe("trailingPartialToolMarkerLength", () => { + it("holds back a split marker prefix at the end of a chunk", () => { + expect(trailingPartialToolMarkerLength("some text { + expect(trailingPartialToolMarkerLength("hello world")).toBe(0) + expect(trailingPartialToolMarkerLength("a < b")).toBe(0) + expect(trailingPartialToolMarkerLength("text ")).toBe(0) + }) + + it("holds back an invoke tag whose name attribute has not arrived", () => { + expect(trailingPartialToolMarkerLength("text { + expect(trailingPartialToolMarkerLength(" { + expect(trailingPartialToolMarkerLength("text <" + "a".repeat(200))).toBe(0) + }) + }) + + describe("quoted markup", () => { + it("does not recover an invoke block inside a fenced code block", () => { + const text = "```\n" + invoke("update_todo_list", param("todos", "[x] one")) + "\n```" + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toContain("invoke") + }) + + it("does not recover an invoke block inside an inline code span", () => { + const text = "avoid `" + invoke("update_todo_list", param("todos", "x")) + "`" + + const { calls } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + }) + + it("does not recover an invoke block quoted in unfenced, backtick-free prose", () => { + const text = "You must never emit " + invoke("update_todo_list", param("todos", "x")) + " directly." + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("does not recover a quoted invoke block that ends its line", () => { + // Defect 3: an empty rest-of-line previously made this look like a genuine leak. + const text = "You must never emit " + invoke("update_todo_list", param("todos", "x")) + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("does not recover an invoke block inside a tilde fence", () => { + const text = "~~~\n" + invoke("update_todo_list", param("todos", "[x] one")) + "\n~~~" + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("does not recover an invoke inside a four-backtick fence containing a three-backtick fence", () => { + // A narrower inner fence must not close the wider outer one, so the invoke stays quoted. + const text = + "````\n```\n" + invoke("update_todo_list", param("todos", "[x] one")) + "\n```\n````" + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("does not recover an invoke inside a tilde fence containing a backtick fence line", () => { + const text = "~~~\n```\n" + invoke("update_todo_list", param("todos", "[x] one")) + "\n```\n~~~" + + const { calls, leftoverText } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toHaveLength(0) + expect(leftoverText).toBe(text) + }) + + it("recovers an invoke block that follows a closed code fence", () => { + const text = "```\nexample output\n```\n" + wrap(invoke("update_todo_list", param("todos", "[x] one"))) + + const { calls } = extractLeakedToolCalls(text, new Set(["update_todo_list"])) + + expect(calls).toEqual([{ name: "update_todo_list", input: { todos: "[x] one" } }]) + }) + + it("does not treat doubled angle brackets as trailing prose after stripping", () => { + // Defect 1: a single strip pass turns `<>` into a tag-looking ``, so the + // trailing-text check must strip repeatedly until stable. + const text = wrap(invoke("update_todo_list", param("todos", "x")) + "<