diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index c15be285ca..671136e688 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -24,6 +24,35 @@ vi.mock("vscode", () => { ) {} } + // A real class (not a plain object): internalCountTokens() and + // extractTextCountFromMessage() branch on `instanceof` against the message + // type, and a plain-object constructor throws "Right-hand side of + // 'instanceof' is not callable". The static factories stay vi.fn() spies so + // tests can assert on the exact content handed to the host (surrogate + // sanitization, message shape) while still yielding real instances. + class MockLanguageModelChatMessage { + role: string + content: unknown + constructor(role: string, content: unknown) { + this.role = role + this.content = content + } + static Assistant: Mock = vi.fn( + (content: string | unknown[]) => + new MockLanguageModelChatMessage( + "assistant", + Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], + ), + ) + static User: Mock = vi.fn( + (content: string | unknown[]) => + new MockLanguageModelChatMessage( + "user", + Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], + ), + ) + } + return { workspace: { getConfiguration: vi.fn(() => ({ @@ -34,12 +63,17 @@ vi.mock("vscode", () => { })), }, CancellationTokenSource: vi.fn(function () { + // Faithful to the real API: cancel() marks the token as requested so + // consumers that read isCancellationRequested observe the cancellation. + const token = { + isCancellationRequested: false, + onCancellationRequested: vi.fn(), + } return { - token: { - isCancellationRequested: false, - onCancellationRequested: vi.fn(), - }, - cancel: vi.fn(), + token, + cancel: vi.fn(() => { + token.isCancellationRequested = true + }), dispose: vi.fn(), } }), @@ -49,16 +83,7 @@ vi.mock("vscode", () => { this.name = "CancellationError" } }, - LanguageModelChatMessage: { - Assistant: vi.fn((content) => ({ - role: "assistant", - content: Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], - })), - User: vi.fn((content) => ({ - role: "user", - content: Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], - })), - }, + LanguageModelChatMessage: MockLanguageModelChatMessage, LanguageModelTextPart: MockLanguageModelTextPart, LanguageModelToolCallPart: MockLanguageModelToolCallPart, LanguageModelToolResultPart: MockLanguageModelToolResultPart, @@ -83,6 +108,7 @@ import type { ApiStreamChunk } from "../../transform/stream" import type { Anthropic } from "@anthropic-ai/sdk" import { openAiModelInfoSaneDefaults, vscodeLlmDefaultModelId, vscodeLlmModels } from "@roo-code/types" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { normalizeToolSchema } from "../../../utils/json-schema" import { getMcpServerTools } from "../../../core/prompts/tools/native-tools/mcp_server" import type { McpHub } from "../../../services/mcp/McpHub" @@ -97,7 +123,21 @@ const mockLanguageModelChat = { version: "1.0", maxInputTokens: 4096, sendRequest: vi.fn(), - countTokens: vi.fn(), + // Default to a valid numeric count so the createMessage input-token counting + // (which now reaches the host mock) does not log a non-numeric warning. + countTokens: vi.fn(async () => 0), +} + +/** + * Returns the instance created by the n-th `new vscode.CancellationTokenSource()` + * call recorded by the module mock, for asserting on the cancellation-token lifecycle. + */ +function tokenSourceInstance(index = 0) { + const result = (vscode.CancellationTokenSource as Mock).mock.results[index] + if (result?.type !== "return") { + return undefined + } + return result.value } describe("VsCodeLmHandler", () => { @@ -703,9 +743,10 @@ describe("VsCodeLmHandler", () => { tools: salvageTools, }) - await expect(stream.next()).rejects.toThrow( - new Error("Zoo Code : Request cancelled by user"), - ) + // The series normalizes host cancellation to a canonical AbortError + // (name "AbortError", same branded message), so assert the message + // only - an Error-instance comparison would also pin the name. + await expect(stream.next()).rejects.toThrow("Zoo Code : Request cancelled by user") await expect(stream.next()).resolves.toEqual({ done: true, value: undefined }) }) @@ -2053,6 +2094,156 @@ describe("VsCodeLmHandler", () => { expect(sent).not.toContain("X".repeat(400_000)) }) + it("sends an oversized request untouched when the context window resolves to zero", async () => { + // A non-positive window disables the budget gate entirely: no trimming and no refusal, even + // for a conversation far larger than any window. The mutants that force the gate open must + // refuse here (the tool_result cannot shrink below its floor against a non-positive raw + // budget), so asserting the send is what distinguishes them. + vi.spyOn(handler, "getCondenseContextWindow").mockReturnValue(0) + 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(5_000) }], + }, + ] + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("ok") + return + })(), + text: (async function* () { + yield "ok" + return + })(), + }) + + const stream = handler.createMessage("system", messages, { taskId: "test-task" }) + const chunks = await collectStream(stream) + + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled() + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + }) + + it("reports the exact raw budget and pairing warning when refusing a request", async () => { + // The refusal message carries the RAW (unclamped) budget: asserting the full text pins the + // budget arithmetic (window * 0.8 * 3 - systemPrompt.length) and both message fragments, + // which the looser regex assertions in the other refusal tests do not cover. + const contextWindow = handler.getCondenseContextWindow() + const targetRawBudgetChars = 1_000 + const systemPrompt = "S".repeat(Math.floor(contextWindow * 0.8 * 3) - targetRawBudgetChars) + 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(1_999) }], + }, + ] + + const stream = handler.createMessage(systemPrompt, messages, { taskId: "test-task" }) + await expect( + (async () => { + for await (const _chunk of stream) { + // drain + } + })(), + ).rejects.toThrow( + "Zoo Code : The request is too large for this model's context window " + + "(estimated 2,006 characters against a budget of 1,000), and it cannot be reduced further without " + + "breaking tool-call pairing. Condense the conversation or start a new task.", + ) + expect(mockLanguageModelChat.sendRequest).not.toHaveBeenCalled() + }) + + it("reports the exact raw budget when the tool schema consumes part of it", async () => { + // The raw budget subtracts the serialized tool schema: with no tools the -toolSchema + // operator is equivalent to +toolSchema, so pin the figure with a non-empty schema. The + // mutant that flips it inflates the reported budget by twice the schema size. + const contextWindow = handler.getCondenseContextWindow() + const targetRawBudgetChars = 1_000 + const tools = [ + { + type: "function" as const, + function: { + name: "calculator", + description: "A simple calculator", + parameters: { + type: "object", + properties: { + operation: { type: "string" }, + numbers: { type: "array", items: { type: "number" } }, + }, + }, + }, + }, + ] + const toolSchemaChars = JSON.stringify(tools).length + const systemPrompt = "S".repeat( + Math.floor(contextWindow * 0.8 * 3) - targetRawBudgetChars - toolSchemaChars, + ) + 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(1_999) }], + }, + ] + + const stream = handler.createMessage(systemPrompt, messages, { taskId: "test-task", tools }) + await expect( + (async () => { + for await (const _chunk of stream) { + // drain + } + })(), + ).rejects.toThrow( + "(estimated 2,006 characters against a budget of 1,000), and it cannot be reduced further without " + + "breaking tool-call pairing.", + ) + expect(mockLanguageModelChat.sendRequest).not.toHaveBeenCalled() + }) + + it("sends a request whose trimmed size equals the raw budget exactly", async () => { + // Admission is strict: a conversation that lands exactly on the raw budget fits, so the + // request must be sent. The >= mutant refuses exactly this boundary, which is the + // off-by-one this test exists to catch. + const contextWindow = handler.getCondenseContextWindow() + const targetRawBudgetChars = 1_000 + const systemPrompt = "S".repeat(Math.floor(contextWindow * 0.8 * 3) - targetRawBudgetChars) + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: [{ type: "text", text: "Y".repeat(targetRawBudgetChars) }], + }, + ] + + 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" }) + const chunks = await collectStream(stream) + + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled() + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + }) it("should handle native tool calls when tools are provided", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [ @@ -2251,574 +2442,2005 @@ describe("VsCodeLmHandler", () => { ) }) - it("should throw a Zoo Code branded error on stream error with error-like object", async () => { + it("should reject with an AbortError when the external signal is already aborted", async () => { const systemPrompt = "You are a helpful assistant" - const messages: Anthropic.Messages.MessageParam[] = [ - { - role: "user" as const, - content: "Hello", - }, - ] + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const controller = new AbortController() + controller.abort() - mockLanguageModelChat.sendRequest.mockRejectedValueOnce({ code: "STREAM_ERROR", details: "broken" }) + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + await expect(stream.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) - await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow( - "Zoo Code : Response stream error:", + // No host request is started for an already-aborted signal. + expect(mockLanguageModelChat.sendRequest).not.toHaveBeenCalled() + }) + + it("should reject with an AbortError when the external signal aborts during client initialization", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + // Gate model selection so the generator waits inside getClient(): the + // cancellation token source and the abort-bridge listener are established + // before this window, and the finally cleanup covers the getClient() await. + let releaseClient: () => void = () => {} + const clientGate = new Promise((resolve) => { + releaseClient = resolve + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), ) + handler["client"] = null - expect(consoleErrorSpy).toHaveBeenCalledWith( - "Zoo Code : Stream error object:", - expect.stringContaining("STREAM_ERROR"), + const controller = new AbortController() + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + const firstChunk = stream.next() + + // Release the client, then abort before the next microtask can reach + // countTokens/sendRequest. + releaseClient() + controller.abort() + + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", ) - consoleErrorSpy.mockRestore() + // The abort landed after client initialization, so no host request started + // and no input tokens were counted (the post-init re-check bails first). + expect(handler["client"]).not.toBeNull() + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(mockLanguageModelChat.countTokens).toHaveBeenCalledTimes(0) + + // The local token source was cancelled for the request that never started, + // and the abort-bridge listener was removed by the finally covering the + // initialization window. + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) }) - it("should log Zoo Code branded warning for unknown chunk type in stream", async () => { + + it("should bridge a mid-flight external abort to the request cancellation token", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) - + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ stream: (async function* () { - // Yield an unknown chunk type (not TextPart, not ToolCallPart) - yield { type: "unknown", foo: "bar" } as unknown as vscode.LanguageModelTextPart - return + await streamGate + yield new vscode.LanguageModelTextPart("Hello") })(), text: (async function* () { - yield "" - return + await streamGate + yield "Hello" })(), }) - const stream = handler.createMessage(systemPrompt, messages) - for await (const _chunk of stream) { - // drain - } - - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Unknown chunk type received:", - expect.objectContaining({ type: "unknown" }), + const controller = new AbortController() + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + const firstChunk = stream.next() + await vi.waitFor(() => expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1)) + + // Abort the external signal while the request is in flight, then let the + // (token-agnostic) mock stream finish. The late abort must stop the stream + // instead of yielding stale chunks. + controller.abort() + releaseStream() + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", ) - consoleWarnSpy.mockRestore() + // The bridge relayed the abort to the request cancellation token. + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + + // The cancellation must reach the host through the very token handed + // to sendRequest: passing a fresh token would leave the host request + // live while the local source looks cancelled. + expect(mockLanguageModelChat.sendRequest.mock.calls[0]?.[2]).toBe(tokenSource.token) }) - it("should log Zoo Code branded warning for invalid text part value", async () => { + it("should attach and detach the abort bridge listener around the request", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) - - // Create a TextPart with a non-string value (number) - const badTextPart = new vscode.LanguageModelTextPart(42 as unknown as string) mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ stream: (async function* () { - yield badTextPart - return + yield new vscode.LanguageModelTextPart("Hello") })(), text: (async function* () { - yield "" - return + yield "Hello" })(), }) - const stream = handler.createMessage(systemPrompt, messages) + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) for await (const _chunk of stream) { // drain } - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Invalid text part value received:", - 42, - ) - - consoleWarnSpy.mockRestore() + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addEventListenerSpy.mock.calls[0][1]) }) - it("should log Zoo Code branded warning for invalid tool callId", async () => { + it("should cancel the request token when the consumer stops consuming early", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) - - // Create a ToolCallPart with a non-string callId - const badToolCall = new vscode.LanguageModelToolCallPart(123 as unknown as string, "valid-name", {}) + // The mock stream yields one chunk and then stays pending, so the host + // request is still in flight when the consumer gives up. + let _releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + _releaseStream = resolve + }) mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ stream: (async function* () { - yield badToolCall - return + yield new vscode.LanguageModelTextPart("Hello") + await streamGate })(), text: (async function* () { - yield "" - return + yield "Hello" + await streamGate })(), }) const stream = handler.createMessage(systemPrompt, messages) - for await (const _chunk of stream) { - // drain - } - - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Invalid tool callId received:", - 123, - ) + await stream.next() + // Stop consuming before the stream finishes: the premature closure must + // cancel the request token so the host stops the request (dispose alone + // frees resources without cancelling). + // (the AsyncGenerator type requires the return value argument) + await stream.return(undefined) - consoleWarnSpy.mockRestore() + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + expect(tokenSource.dispose).toHaveBeenCalled() }) - it("should log Zoo Code branded warning for invalid tool input", async () => { + it("should report the canonical abort error and cancel the token twice when the external signal is already aborted", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + const controller = new AbortController() + controller.abort() - // Create a ToolCallPart with a string input (not an object) - const badToolCall = new vscode.LanguageModelToolCallPart( - "call-1", - "valid-name", - "not-an-object" as unknown as object, + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + await expect(stream.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", ) - mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ - stream: (async function* () { - yield badToolCall - return - })(), - text: (async function* () { - yield "" - return - })(), - }) - const stream = handler.createMessage(systemPrompt, messages) - for await (const _chunk of stream) { - // drain - } + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + // The pre-check and the finally each cancel the request token; the + // already-aborted signal cannot fire the bridge listener a second time. + expect(tokenSourceInstance().cancel).toHaveBeenCalledTimes(2) + }) - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Invalid tool input received:", - "not-an-object", + it("should fail fast before client initialization when the external signal is already aborted", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + handler["client"] = null + const selectCallsBefore = (vscode.lm.selectChatModels as Mock).mock.calls.length + + const controller = new AbortController() + controller.abort() + + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + await expect(stream.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", ) - consoleWarnSpy.mockRestore() + // The pre-abort short-circuit runs before getClient(): the host is never contacted. + expect((vscode.lm.selectChatModels as Mock).mock.calls.length).toBe(selectCallsBefore) }) - it("should log Zoo Code branded error when tool call processing fails", async () => { + it("should cancel the request token synchronously when the external signal aborts during client initialization", async () => { const systemPrompt = "You are a helpful assistant" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] - const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + let releaseClient: () => void = () => {} + const clientGate = new Promise((resolve) => { + releaseClient = resolve + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), + ) + handler["client"] = null + // Consume the beforeEach's queued selectChatModels value so the request's + // own client lookup is the one that parks on the gate above. + await (vscode.lm.selectChatModels as Mock)() + const selectCallsBefore = (vscode.lm.selectChatModels as Mock).mock.calls.length + + const controller = new AbortController() + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + const firstChunk = stream.next() + await vi.waitFor(() => + expect((vscode.lm.selectChatModels as Mock).mock.calls.length).toBe(selectCallsBefore + 1), + ) - // Create a ToolCallPart with circular input that will throw on JSON.stringify - const circularInput: Record = { name: "circular" } - circularInput.self = circularInput + releaseClient() + controller.abort() + // The abort bridge must cancel the request token synchronously on + // controller.abort(): the generator is still suspended inside getClient(), + // so the post-init re-check and the finally have not run yet. + expect(tokenSourceInstance().token.isCancellationRequested).toBe(true) + + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) - const badToolCall = new vscode.LanguageModelToolCallPart("call-1", "valid-name", circularInput) - mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ - stream: (async function* () { - yield badToolCall - return + // Bridge + post-init re-check + finally each cancel the request token. + expect(tokenSourceInstance().cancel).toHaveBeenCalledTimes(3) + }) + + it("should report the canonical abort error when client initialization fails while the external signal aborts", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + let rejectClient: (reason: unknown) => void = () => {} + const clientGate: Promise = new Promise((_resolve, reject) => { + rejectClient = reject + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), + ) + handler["client"] = null + // Consume the beforeEach's queued selectChatModels value so the request's + // own client lookup is the one that parks on the gate above. + await (vscode.lm.selectChatModels as Mock)() + const selectCallsBefore = (vscode.lm.selectChatModels as Mock).mock.calls.length + + const controller = new AbortController() + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + const firstChunk = stream.next() + await vi.waitFor(() => + expect((vscode.lm.selectChatModels as Mock).mock.calls.length).toBe(selectCallsBefore + 1), + ) + + // The signal aborts while client initialization is failing: the canonical + // abort error must win over the original client-initialization error. + controller.abort() + rejectClient(new Error("network down")) + + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + }) + + it("should brand host cancellation errors with the AbortError name", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new vscode.CancellationError()) + + await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request cancelled by user", + ) + }) + + it("should preserve message content in the host request", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Hello!") + return })(), text: (async function* () { - yield "" + yield "Hello!" return })(), }) - const stream = handler.createMessage(systemPrompt, messages, { - taskId: "test-task", - tools: [ - { - type: "function" as const, - function: { name: "test", description: "", parameters: { type: "object", properties: {} } }, - }, - ], + const stream = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream) { + // drain + } + + // The host request must receive the system prompt followed by the + // converted conversation messages, with content preserved. + const requestMessages = (mockLanguageModelChat.sendRequest as Mock).mock + .calls[0][0] as vscode.LanguageModelChatMessage[] + expect(requestMessages).toHaveLength(2) + expect(requestMessages[0].role).toBe("assistant") + expect(requestMessages[1].role).toBe("user") + const textValues = (content: string | vscode.LanguageModelChatMessage["content"]) => + typeof content === "string" + ? [content] + : content.map((part) => (part as vscode.LanguageModelTextPart).value) + expect(textValues(requestMessages[0].content)).toEqual([systemPrompt]) + expect(textValues(requestMessages[1].content)).toEqual(["Hello"]) + }) + + it("should release the request cancellation slot when the request completes", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Hello!") + return + })(), + text: (async function* () { + yield "Hello!" + return + })(), }) + + const stream = handler.createMessage(systemPrompt, messages) for await (const _chunk of stream) { // drain } - expect(consoleErrorSpy).toHaveBeenCalledWith( - "Zoo Code : Failed to process tool call:", - expect.any(Error), + // The finally must clear the slot once the request is done so the next + // request (or dispose) does not see a stale token source. + expect(handler["currentRequestCancellation"]).toBeNull() + }) + + it("should keep the new request's cancellation source when a previous request finishes", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const gateA = new Promise(() => {}) + const gateB = new Promise(() => {}) + // Each response yields one chunk before parking so a later .return() settles; + // a generator parked inside an inner await would defer the return indefinitely. + mockLanguageModelChat.sendRequest + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("A1") + await gateA + yield new vscode.LanguageModelTextPart("A2") + })(), + text: (async function* () { + yield "A1" + await gateA + yield "A2" + })(), + }) + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("B1") + await gateB + yield new vscode.LanguageModelTextPart("B2") + })(), + text: (async function* () { + yield "B1" + await gateB + yield "B2" + })(), + }) + + const streamA = handler.createMessage(systemPrompt, messages) + const firstChunkA = await streamA.next() + expect(firstChunkA.value).toEqual({ type: "text", text: "A1" }) + + // A second overlapping request establishes its own cancellation source; + // its ensureCleanState() synchronously cancels the first request's token. + const streamB = handler.createMessage(systemPrompt, messages) + const firstChunkB = await streamB.next() + expect(firstChunkB.value).toEqual({ type: "text", text: "B1" }) + + const sourceB = tokenSourceInstance(1) + expect(handler["currentRequestCancellation"]).toBe(sourceB) + + // Stopping the first request early must not clear the second request's + // source: only a request that still owns the slot may release it. + await streamA.return(undefined) + expect(handler["currentRequestCancellation"]).toBe(sourceB) + + await streamB.return(undefined) + expect(handler["currentRequestCancellation"]).toBeNull() + }) + + it("should stop yielding stale chunks when a newer request supersedes it mid-stream", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + let releaseA: () => void = () => {} + const streamAGate = new Promise((resolve) => { + releaseA = resolve + }) + mockLanguageModelChat.sendRequest + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("A1") + await streamAGate + yield new vscode.LanguageModelTextPart("A2") + })(), + text: (async function* () { + yield "A1" + await streamAGate + yield "A2" + })(), + }) + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("B1") + })(), + text: (async function* () { + yield "B1" + })(), + }) + + const streamA = handler.createMessage(systemPrompt, messages) + const firstChunkA = await streamA.next() + expect(firstChunkA.value).toEqual({ type: "text", text: "A1" }) + + // A second overlapping request supersedes the first one: its + // ensureCleanState() cancels A's token while A is parked between host chunks. + const streamB = handler.createMessage(systemPrompt, messages) + const firstChunkB = await streamB.next() + expect(firstChunkB.value).toEqual({ type: "text", text: "B1" }) + expect(tokenSourceInstance(0).token.isCancellationRequested).toBe(true) + + releaseA() + + // The per-chunk re-check must abort A instead of yielding A2's stale chunk; + // the external signal was never involved, so only the token check can see it. + await expect(streamA.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + await streamB.return(undefined) + expect(handler["currentRequestCancellation"]).toBeNull() + consoleErrorSpy.mockRestore() + }) + + it("should reject with an AbortError when a newer request supersedes during input token counting", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + // Request A carries no external signal: only the request-local token can + // observe the supersession, so this pins the post-count token check. + const releases: Array<() => void> = [] + mockLanguageModelChat.countTokens.mockImplementation(async () => { + await new Promise((resolve) => { + releases.push(resolve) + }) + return 0 + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("B1") + })(), + text: (async function* () { + yield "B1" + })(), + }) + + const streamA = handler.createMessage(systemPrompt, messages) + const firstA = streamA.next() + await vi.waitFor(() => expect(mockLanguageModelChat.countTokens).toHaveBeenCalled()) + + // A newer request supersedes A: its ensureCleanState() cancels A's token + // while A is parked inside countTokens. + const streamB = handler.createMessage(systemPrompt, messages) + const firstB = streamB.next() + await vi.waitFor(() => expect(tokenSourceInstance(0).token.isCancellationRequested).toBe(true)) + + for (const release of releases) { + release() + } + + // A must bail at the post-count check: it never reaches the host call, + // so the only sendRequest is B's. + await expect(firstA).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) + + await streamB.return(undefined) + await firstB + }) + + it("should reject with an AbortError when the external signal aborts before a quiet stream end", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const controller = new AbortController() + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Partial text") + await streamGate + return // Quiet end: no chunk and no throw after the abort. + })(), + text: (async function* () { + yield "Partial text" + await streamGate + return + })(), + }) + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const firstChunk = await stream.next() + expect(firstChunk.value).toEqual({ type: "text", text: "Partial text" }) + + // The caller aborts while the host stream is parked; the host then ends quietly. + controller.abort() + releaseStream() + + // The post-loop guard must reject with the canonical abort instead of + // counting output tokens and yielding a usage chunk for the partial text. + await expect(stream.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + }) + + it("should reject with the canonical abort instead of flushing pending salvage after a quiet post-abort stream end", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const controller = new AbortController() + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + const leaked = + 'add' + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + // A complete leaked tool-call block: the salvage path buffers it + // whole, so nothing is yielded while it is held. + yield new vscode.LanguageModelTextPart(leaked) + await streamGate + return // Quiet end: no chunk and no throw after the abort. + })(), + text: (async function* () { + yield leaked + await streamGate + return + })(), + }) + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ + abortSignal: controller.signal, + tools: [ + { + type: "function", + function: { + name: "calculator", + description: "A simple calculator", + parameters: { type: "object", properties: { operation: { type: "string" } } }, + }, + }, + ], + }), + ) + const firstChunk = stream.next() + await vi.waitFor(() => expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1)) + + // The caller aborts while the salvage buffer is still pending; the host + // then ends the stream quietly. + controller.abort() + releaseStream() + + // The post-loop guard must reject before flushSalvage can emit the + // buffered tool call: if the guard and the flush swapped places, + // firstChunk would resolve with a tool_call chunk instead of rejecting. + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + }) + + it("should reject with an AbortError when a superseded request's stream ends quietly", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + let releaseA: () => void = () => {} + const streamAGate = new Promise((resolve) => { + releaseA = resolve + }) + mockLanguageModelChat.sendRequest + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("A1") + await streamAGate + return // Quiet end after supersession. + })(), + text: (async function* () { + yield "A1" + await streamAGate + return + })(), + }) + .mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("B1") + })(), + text: (async function* () { + yield "B1" + })(), + }) + + const streamA = handler.createMessage(systemPrompt, messages) + const firstChunkA = await streamA.next() + expect(firstChunkA.value).toEqual({ type: "text", text: "A1" }) + + // B supersedes A: ensureCleanState() cancels A's token while A is parked + // between host chunks; the host stream then ends without yielding. + const streamB = handler.createMessage(systemPrompt, messages) + const firstChunkB = await streamB.next() + expect(firstChunkB.value).toEqual({ type: "text", text: "B1" }) + expect(tokenSourceInstance(0).token.isCancellationRequested).toBe(true) + + releaseA() + + // The external signal was never involved, so only the token check in the + // post-loop guard can turn the quiet end into the canonical abort. + await expect(streamA.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + await streamB.return(undefined) + expect(handler["currentRequestCancellation"]).toBeNull() + consoleErrorSpy.mockRestore() + }) + + it("should abort before sendRequest when the external signal fires while counting input tokens", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const controller = new AbortController() + // The host's token-counting call is in flight when the caller aborts; the + // bridge cancels the request token during the count, so the post-counting + // check must stop the request before sendRequest is invoked. + mockLanguageModelChat.countTokens.mockImplementationOnce(async () => { + controller.abort() + return 10 + }) + + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) + + await expect(stream.next()).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + // The post-counting check must stop the request before the host is invoked. + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + }) + + it("should abort a superseded request before counting tokens when superseded during client initialization", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + // Distinct per-request clients so A's host calls can be attributed to A + // alone; B proceeds on its own client and must not pollute A's counters. + const clientA = { ...mockLanguageModelChat, sendRequest: vi.fn(), countTokens: vi.fn() } + const clientB = { ...mockLanguageModelChat, sendRequest: vi.fn(), countTokens: vi.fn(async () => 0) } + clientB.sendRequest.mockResolvedValue({ + stream: (async function* () {})(), + text: (async function* () {})(), + }) + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + let releaseA: () => void = () => {} + const gateA = new Promise((resolve) => { + releaseA = resolve + }) + let releaseB: () => void = () => {} + const gateB = new Promise((resolve) => { + releaseB = resolve + }) + // Gate each client-initialization lookup deterministically. mockReset() drops + // any mockImplementationOnce entries left behind by earlier tests (clearAllMocks + // does not clear the once-queue); a stale entry would otherwise consume A's + // lookup and hand B the wrong gate. + const selectChatModels = vscode.lm.selectChatModels as Mock + selectChatModels.mockReset() + let lookup = 0 + selectChatModels.mockImplementation(() => { + const k = lookup++ + if (k === 0) { + return gateA.then(() => [clientA]) + } + if (k === 1) { + return gateB.then(() => [clientB]) + } + return Promise.resolve([mockLanguageModelChat]) + }) + handler["client"] = null + + // A parks inside getClient() on gateA. + const streamA = handler.createMessage(systemPrompt, messages) + const nextA = streamA.next() + + // B's ensureCleanState() synchronously cancels A's token; B parks on gateB. + const streamB = handler.createMessage(systemPrompt, messages) + const nextB = streamB.next() + expect(tokenSourceInstance(0).token.isCancellationRequested).toBe(true) + + // Release A so it resumes into the post-init guard, which must see the + // cancelled token and abort A before it counts tokens or invokes the host. + releaseA() + await expect(nextA).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + // Supersession must stop A before any host work: A never counted tokens or + // invoked the host. + expect(clientA.countTokens).not.toHaveBeenCalled() + expect(clientA.sendRequest).not.toHaveBeenCalled() + + // Release B so it runs to completion on clientB (empty stream), draining + // the request slot so afterEach's dispose() finds a clean handler. + releaseB() + await nextB + await streamB.next() + expect(handler["currentRequestCancellation"]).toBeNull() + consoleErrorSpy.mockRestore() + }) + + it("should normalize a superseded request's client-initialization failure to AbortError", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + // B proceeds on its own client and must not pollute A's counters. + const clientB = { ...mockLanguageModelChat, sendRequest: vi.fn(), countTokens: vi.fn(async () => 0) } + clientB.sendRequest.mockResolvedValue({ + stream: (async function* () {})(), + text: (async function* () {})(), + }) + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + let rejectA: (reason: Error) => void = () => {} + const gateA = new Promise((_resolve, reject) => { + rejectA = reject + }) + let releaseB: () => void = () => {} + const gateB = new Promise((resolve) => { + releaseB = resolve + }) + // Gate each client-initialization lookup deterministically (see the + // supersession test above for why mockReset() is required here). + const selectChatModels = vscode.lm.selectChatModels as Mock + selectChatModels.mockReset() + let lookup = 0 + selectChatModels.mockImplementation(() => { + const k = lookup++ + if (k === 0) { + return gateA + } + if (k === 1) { + return gateB.then(() => [clientB]) + } + return Promise.resolve([mockLanguageModelChat]) + }) + handler["client"] = null + + // A parks inside getClient() on gateA. + const streamA = handler.createMessage(systemPrompt, messages) + const nextA = streamA.next() + + // B's ensureCleanState() synchronously cancels A's token. + const streamB = handler.createMessage(systemPrompt, messages) + const nextB = streamB.next() + expect(tokenSourceInstance(0).token.isCancellationRequested).toBe(true) + + // A's client initialization fails while A is already superseded: the + // catch path must normalize this to AbortError the way completePrompt + // does through isAborted() (which includes the token state) instead of + // leaking the wrapped client-creation error. + rejectA(new Error("boom")) + await expect(nextA).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + // Release B so it runs to completion on clientB (empty stream), draining + // the request slot so afterEach's dispose() finds a clean handler. + releaseB() + await nextB + await streamB.next() + expect(handler["currentRequestCancellation"]).toBeNull() + consoleErrorSpy.mockRestore() + }) + + it("should throw a Zoo Code branded error on stream error with error-like object", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user" as const, + content: "Hello", + }, + ] + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + mockLanguageModelChat.sendRequest.mockRejectedValueOnce({ code: "STREAM_ERROR", details: "broken" }) + + await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow( + "Zoo Code : Response stream error:", + ) + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Zoo Code : Stream error object:", + expect.stringContaining("STREAM_ERROR"), + ) + + consoleErrorSpy.mockRestore() + }) + it("should log Zoo Code branded warning for unknown chunk type in stream", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + // Yield an unknown chunk type (not TextPart, not ToolCallPart) + yield { type: "unknown", foo: "bar" } as unknown as vscode.LanguageModelTextPart + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream) { + // drain + } + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Unknown chunk type received:", + expect.objectContaining({ type: "unknown" }), + ) + + consoleWarnSpy.mockRestore() + }) + + it("should log Zoo Code branded warning for invalid text part value", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + // Create a TextPart with a non-string value (number) + const badTextPart = new vscode.LanguageModelTextPart(42 as unknown as string) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield badTextPart + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream) { + // drain + } + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Invalid text part value received:", + 42, + ) + + consoleWarnSpy.mockRestore() + }) + + it("should log Zoo Code branded warning for invalid tool callId", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + // Create a ToolCallPart with a non-string callId + const badToolCall = new vscode.LanguageModelToolCallPart(123 as unknown as string, "valid-name", {}) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield badToolCall + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream) { + // drain + } + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Invalid tool callId received:", + 123, + ) + + consoleWarnSpy.mockRestore() + }) + + it("should log Zoo Code branded warning for invalid tool input", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + // Create a ToolCallPart with a string input (not an object) + const badToolCall = new vscode.LanguageModelToolCallPart( + "call-1", + "valid-name", + "not-an-object" as unknown as object, + ) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield badToolCall + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream) { + // drain + } + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Invalid tool input received:", + "not-an-object", + ) + + consoleWarnSpy.mockRestore() + }) + + it("should log Zoo Code branded error when tool call processing fails", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + // Create a ToolCallPart with circular input that will throw on JSON.stringify + const circularInput: Record = { name: "circular" } + circularInput.self = circularInput + + const badToolCall = new vscode.LanguageModelToolCallPart("call-1", "valid-name", circularInput) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield badToolCall + return + })(), + text: (async function* () { + yield "" + return + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages, { + taskId: "test-task", + tools: [ + { + type: "function" as const, + function: { name: "test", description: "", parameters: { type: "object", properties: {} } }, + }, + ], + }) + for await (const _chunk of stream) { + // drain + } + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Zoo Code : Failed to process tool call:", + expect.any(Error), + ) + + consoleErrorSpy.mockRestore() + }) + }) + + describe("getClient", () => { + it("should log Zoo Code branded debug when creating client with selector", async () => { + const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + + // @ts-ignore – access private method for coverage + await handler["getClient"]() + + expect(consoleDebugSpy).toHaveBeenCalledWith( + "Zoo Code : Creating client with selector:", + expect.any(Object), + ) + + consoleDebugSpy.mockRestore() + }) + + it("should throw a Zoo Code branded error when getClient fails to create client", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + ;(vscode.lm.selectChatModels as Mock).mockRejectedValueOnce(new Error("network error")) + handler["client"] = null + + // @ts-ignore – access private method for coverage + await expect(handler["getClient"]()).rejects.toThrow( + "Zoo Code : Failed to create client:", + ) + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Zoo Code : Client creation failed:", + expect.stringContaining("network error"), + ) + + consoleErrorSpy.mockRestore() + }) + }) + + describe("initializeClient", () => { + it("should log when client is already initialized", async () => { + const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + + handler["client"] = mockLanguageModelChat + await handler.initializeClient() + + expect(consoleDebugSpy).toHaveBeenCalledWith("Zoo Code : Client already initialized") + + consoleDebugSpy.mockRestore() + }) + + it("should log success when client is initialized", async () => { + const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + + await handler.initializeClient() + + expect(consoleDebugSpy).toHaveBeenCalledWith( + "Zoo Code : Client initialized successfully", + ) + + consoleDebugSpy.mockRestore() + }) + + it("should throw a Zoo Code branded error when client initialization fails", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + ;(vscode.lm.selectChatModels as Mock).mockRejectedValue(new Error("select failed")) + handler["client"] = null + + // Catch the unhandled rejection that may occur from the constructor's async call + const initPromise = handler.initializeClient() + + await expect(initPromise).rejects.toThrow("Zoo Code : Failed to initialize client:") + + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Zoo Code : Client initialization failed:", + expect.stringContaining("select failed"), + ) + + consoleErrorSpy.mockRestore() + }) + }) + + describe("getModel", () => { + it("should return model info when client exists", async () => { + const mockModel = { ...mockLanguageModelChat } + // The handler starts async initialization in the constructor. + // Make the test deterministic by explicitly (re)initializing here. + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + await handler.initializeClient() + + const model = handler.getModel() + expect(model.id).toBe("test-model") + expect(model.info).toBeDefined() + expect(model.info.contextWindow).toBe(4096) + }) + + it("should return fallback model info when no client exists", () => { + const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + + // Clear the client first + handler["client"] = null + const model = handler.getModel() + expect(model.id).toBe("test-vendor/test-family") + expect(model.info).toBeDefined() + expect(consoleDebugSpy).toHaveBeenCalledWith( + "Zoo Code : No client available, using fallback model info", + ) + + consoleDebugSpy.mockRestore() + }) + + it("should return basic model info when client exists", async () => { + const mockModel = { ...mockLanguageModelChat } + // The handler starts async initialization in the constructor. + // Make the test deterministic by explicitly (re)initializing here. + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + await handler.initializeClient() + + const model = handler.getModel() + expect(model.info).toBeDefined() + expect(model.info.contextWindow).toBe(4096) + }) + + it("should return fallback model info when no client exists", () => { + // Clear the client first + handler["client"] = null + const model = handler.getModel() + expect(model.info).toBeDefined() + }) + + it("should use the full advertised maxInputTokens without an upper cap", async () => { + // A large advertised window is surfaced as-is, not clamped to a smaller default. + const mockModel = { ...mockLanguageModelChat, maxInputTokens: 936000 } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + await handler.initializeClient() + + const model = handler.getModel() + expect(model.info.contextWindow).toBe(936000) + }) + + it("should pass through a small maxInputTokens unchanged", async () => { + const mockModel = { ...mockLanguageModelChat, maxInputTokens: 4096 } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + await handler.initializeClient() + + const model = handler.getModel() + expect(model.info.contextWindow).toBe(4096) + }) + + it("should fall back to sane defaults when maxInputTokens is not a number", async () => { + const mockModel = { ...mockLanguageModelChat, maxInputTokens: undefined as unknown as number } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) + handler["client"] = null + await handler.initializeClient() + + const model = handler.getModel() + expect(model.info.contextWindow).toBe(openAiModelInfoSaneDefaults.contextWindow) + }) + }) + + describe("getCondenseContextWindow", () => { + it("uses the static-table maxInputTokens for a known VS Code LM family", () => { + const opusHandler = new VsCodeLmHandler({ + vsCodeLmModelSelector: { vendor: "copilot", family: "claude-opus-4.8" }, + }) + expect(opusHandler.getCondenseContextWindow()).toBe(vscodeLlmModels["claude-opus-4.8"].maxInputTokens) + opusHandler.dispose() + }) + + it("falls back to the default-row maxInputTokens for an unknown family (catalog drift)", () => { + // `test-family` isn't a curated row (e.g. a selector left over from a dropped model), so the + // gate resolves the default row instead of the inflated live window. + handler["client"] = mockLanguageModelChat as unknown as vscode.LanguageModelChat + expect(handler.getCondenseContextWindow()).toBe(vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens) + }) + + it("falls back to the default-row maxInputTokens when no family is resolvable (no client, no selector family)", () => { + // No client and no selector family means `family` is undefined, so the gate uses the default + // row's maxInputTokens rather than the live getModel().info.contextWindow. + const noFamilyHandler = new VsCodeLmHandler({ vsCodeLmModelSelector: { vendor: "copilot" } }) + noFamilyHandler["client"] = null + expect(noFamilyHandler.getCondenseContextWindow()).toBe( + vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens, + ) + noFamilyHandler.dispose() + }) + + it("falls back to the derived window when the static row exists but maxInputTokens is non-positive", () => { + // A curated row exists but its maxInputTokens is <= 0, so the `> 0` guard fails and the gate + // falls back to getModel().info.contextWindow. + const family = "claude-opus-4.8" + const original = vscodeLlmModels[family].maxInputTokens + try { + ;(vscodeLlmModels[family] as { maxInputTokens: number }).maxInputTokens = 0 + const guardHandler = new VsCodeLmHandler({ + vsCodeLmModelSelector: { vendor: "copilot", family }, + }) + // Leave the client unset so `family` resolves from the selector, forcing the zeroed + // static row to be read instead of a live client's family. + guardHandler["client"] = null + expect(guardHandler.getCondenseContextWindow()).toBe(guardHandler.getModel().info.contextWindow) + expect(guardHandler.getCondenseContextWindow()).toBe(openAiModelInfoSaneDefaults.contextWindow) + guardHandler.dispose() + } finally { + ;(vscodeLlmModels[family] as { maxInputTokens: number }).maxInputTokens = original + } + }) + }) + + describe("countTokens", () => { + beforeEach(() => { + handler["client"] = mockLanguageModelChat + }) + + it("should count tokens when called outside of an active request", async () => { + // Ensure no active request cancellation token exists + handler["currentRequestCancellation"] = null + + mockLanguageModelChat.countTokens.mockResolvedValueOnce(42) + + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello world" }] + const result = await handler.countTokens(content) + + expect(result).toBe(42) + expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("Hello world", expect.any(Object)) + }) + + it("should count tokens when called during an active request", async () => { + // Simulate an active request with a cancellation token + const mockCancellation = { + token: { isCancellationRequested: false, onCancellationRequested: vi.fn() }, + cancel: vi.fn(), + dispose: vi.fn(), + } + handler["currentRequestCancellation"] = mockCancellation as unknown as vscode.CancellationTokenSource + + mockLanguageModelChat.countTokens.mockResolvedValueOnce(50) + + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Test content" }] + const result = await handler.countTokens(content) + + expect(result).toBe(50) + expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("Test content", mockCancellation.token) + }) + + it("should return 0 when no client is available", async () => { + handler["client"] = null + handler["currentRequestCancellation"] = null + + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello" }] + const result = await handler.countTokens(content) + + expect(result).toBe(0) + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : No client available for token counting", ) - consoleErrorSpy.mockRestore() + consoleWarnSpy.mockRestore() }) - }) - describe("getClient", () => { - it("should log Zoo Code branded debug when creating client with selector", async () => { + it("should handle image blocks with placeholder", async () => { + handler["currentRequestCancellation"] = null + mockLanguageModelChat.countTokens.mockResolvedValueOnce(5) + + const content: Anthropic.Messages.ContentBlockParam[] = [ + { type: "image", source: { type: "base64", media_type: "image/png", data: "abc" } }, + ] + const result = await handler.countTokens(content) + + expect(result).toBe(5) + expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("[IMAGE]", expect.any(Object)) + }) + + it("should return 0 and log when empty text is provided to internalCountTokens", async () => { + handler["currentRequestCancellation"] = null const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) - const mockModel = { ...mockLanguageModelChat } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - // @ts-ignore – access private method for coverage - await handler["getClient"]() + // @ts-ignore – access private method for coverage of line 234 + const result = await handler["internalCountTokens"]("") + expect(result).toBe(0) expect(consoleDebugSpy).toHaveBeenCalledWith( - "Zoo Code : Creating client with selector:", - expect.any(Object), + "Zoo Code : Empty text provided for token counting", ) consoleDebugSpy.mockRestore() }) - it("should throw a Zoo Code branded error when getClient fails to create client", async () => { - const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) - ;(vscode.lm.selectChatModels as Mock).mockRejectedValueOnce(new Error("network error")) - handler["client"] = null + it("should return 0 and log when non-numeric token count is received", async () => { + handler["currentRequestCancellation"] = null + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) - // @ts-ignore – access private method for coverage - await expect(handler["getClient"]()).rejects.toThrow( - "Zoo Code : Failed to create client:", + mockLanguageModelChat.countTokens.mockResolvedValueOnce("not-a-number" as unknown as number) + + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "test" }] + const result = await handler.countTokens(content) + + expect(result).toBe(0) + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Non-numeric token count received:", + "not-a-number", ) - expect(consoleErrorSpy).toHaveBeenCalledWith( - "Zoo Code : Client creation failed:", - expect.stringContaining("network error"), + consoleWarnSpy.mockRestore() + }) + + it("should return 0 and log when negative token count is received", async () => { + handler["currentRequestCancellation"] = null + const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + + mockLanguageModelChat.countTokens.mockResolvedValueOnce(-5) + + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "test" }] + const result = await handler.countTokens(content) + + expect(result).toBe(0) + expect(consoleWarnSpy).toHaveBeenCalledWith( + "Zoo Code : Negative token count received:", + -5, ) - consoleErrorSpy.mockRestore() + consoleWarnSpy.mockRestore() }) }) - describe("initializeClient", () => { - it("should log when client is already initialized", async () => { - const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + describe("completePrompt", () => { + it("should complete single prompt", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - handler["client"] = mockLanguageModelChat - await handler.initializeClient() + const responseText = "Completed text" + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart(responseText) + return + })(), + text: (async function* () { + yield responseText + return + })(), + }) - expect(consoleDebugSpy).toHaveBeenCalledWith("Zoo Code : Client already initialized") + // Override the default client with our test client to ensure it uses + // the mock implementation rather than the default fallback + handler["client"] = mockLanguageModelChat - consoleDebugSpy.mockRestore() + const result = await handler.completePrompt("Test prompt") + expect(result).toBe(responseText) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled() }) - it("should log success when client is initialized", async () => { - const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + it("should handle errors during completion", async () => { const mockModel = { ...mockLanguageModelChat } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - await handler.initializeClient() + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("Completion failed")) - expect(consoleDebugSpy).toHaveBeenCalledWith( - "Zoo Code : Client initialized successfully", - ) + // Make sure we're using the mock client + handler["client"] = mockLanguageModelChat - consoleDebugSpy.mockRestore() + const promise = handler.completePrompt("Test prompt") + await expect(promise).rejects.toThrow("VSCode LM completion error: Completion failed") }) - it("should throw a Zoo Code branded error when client initialization fails", async () => { - const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) - ;(vscode.lm.selectChatModels as Mock).mockRejectedValue(new Error("select failed")) - handler["client"] = null - - // Catch the unhandled rejection that may occur from the constructor's async call - const initPromise = handler.initializeClient() + it("should work without options (backward compatible)", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - await expect(initPromise).rejects.toThrow("Zoo Code : Failed to initialize client:") + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) - expect(consoleErrorSpy).toHaveBeenCalledWith( - "Zoo Code : Client initialization failed:", - expect.stringContaining("select failed"), - ) + handler["client"] = mockLanguageModelChat - consoleErrorSpy.mockRestore() + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Completed text") }) - }) - describe("getModel", () => { - it("should return model info when client exists", async () => { + it("should bridge the abort signal to a fresh cancellation token and dispose it", async () => { const mockModel = { ...mockLanguageModelChat } - // The handler starts async initialization in the constructor. - // Make the test deterministic by explicitly (re)initializing here. - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - await handler.initializeClient() + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - const model = handler.getModel() - expect(model.id).toBe("test-model") - expect(model.info).toBeDefined() - expect(model.info.contextWindow).toBe(4096) - }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) - it("should return fallback model info when no client exists", () => { - const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + handler["client"] = mockLanguageModelChat - // Clear the client first - handler["client"] = null - const model = handler.getModel() - expect(model.id).toBe("test-vendor/test-family") - expect(model.info).toBeDefined() - expect(consoleDebugSpy).toHaveBeenCalledWith( - "Zoo Code : No client available, using fallback model info", - ) + const controller = new AbortController() + const result = await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) - consoleDebugSpy.mockRestore() + expect(result).toBe("Completed text") + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).not.toHaveBeenCalled() + expect(tokenSource.dispose).toHaveBeenCalled() }) - it("should return basic model info when client exists", async () => { + it("should attach and detach the abort listener around the completion", async () => { const mockModel = { ...mockLanguageModelChat } - // The handler starts async initialization in the constructor. - // Make the test deterministic by explicitly (re)initializing here. - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - await handler.initializeClient() + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - const model = handler.getModel() - expect(model.info).toBeDefined() - expect(model.info.contextWindow).toBe(4096) - }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) - it("should return fallback model info when no client exists", () => { - // Clear the client first - handler["client"] = null - const model = handler.getModel() - expect(model.info).toBeDefined() - }) + handler["client"] = mockLanguageModelChat - it("should use the full advertised maxInputTokens without an upper cap", async () => { - // A large advertised window is surfaced as-is, not clamped to a smaller default. - const mockModel = { ...mockLanguageModelChat, maxInputTokens: 936000 } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - await handler.initializeClient() + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") - const model = handler.getModel() - expect(model.info.contextWindow).toBe(936000) + await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addEventListenerSpy.mock.calls[0][1]) }) - it("should pass through a small maxInputTokens unchanged", async () => { - const mockModel = { ...mockLanguageModelChat, maxInputTokens: 4096 } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - await handler.initializeClient() + it("should reject with an AbortError when the signal is already aborted", async () => { + handler["client"] = mockLanguageModelChat - const model = handler.getModel() - expect(model.info.contextWindow).toBe(4096) - }) + const controller = new AbortController() + controller.abort() - it("should fall back to sane defaults when maxInputTokens is not a number", async () => { - const mockModel = { ...mockLanguageModelChat, maxInputTokens: undefined as unknown as number } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValue([mockModel]) - handler["client"] = null - await handler.initializeClient() + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) - const model = handler.getModel() - expect(model.info.contextWindow).toBe(openAiModelInfoSaneDefaults.contextWindow) + // Fails fast before invoking the host: the pre-aborted signal cancels the token. + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() }) - }) - describe("getCondenseContextWindow", () => { - it("uses the static-table maxInputTokens for a known VS Code LM family", () => { - const opusHandler = new VsCodeLmHandler({ - vsCodeLmModelSelector: { vendor: "copilot", family: "claude-opus-4.8" }, + it("should reject with an AbortError when the signal aborts mid-flight", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("partial") + })(), + text: (async function* () { + await streamGate + yield "partial" + })(), }) - expect(opusHandler.getCondenseContextWindow()).toBe(vscodeLlmModels["claude-opus-4.8"].maxInputTokens) - opusHandler.dispose() - }) - it("falls back to the default-row maxInputTokens for an unknown family (catalog drift)", () => { - // `test-family` isn't a curated row (e.g. a selector left over from a dropped model), so the - // gate resolves the default row instead of the inflated live window. - handler["client"] = mockLanguageModelChat as unknown as vscode.LanguageModelChat - expect(handler.getCondenseContextWindow()).toBe(vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens) + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + await vi.waitFor(() => expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1)) + + // Abort while the stream is in flight. The bridge must cancel the request + // token synchronously on controller.abort(): the generator is still + // suspended inside the stream, so the finally cleanup has not run yet. + controller.abort() + expect(tokenSourceInstance().token.isCancellationRequested).toBe(true) + releaseStream() + + // The late abort must surface as the canonical abort error. + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) }) - it("falls back to the default-row maxInputTokens when no family is resolvable (no client, no selector family)", () => { - // No client and no selector family means `family` is undefined, so the gate uses the default - // row's maxInputTokens rather than the live getModel().info.contextWindow. - const noFamilyHandler = new VsCodeLmHandler({ vsCodeLmModelSelector: { vendor: "copilot" } }) - noFamilyHandler["client"] = null - expect(noFamilyHandler.getCondenseContextWindow()).toBe( - vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens, + it("should reject with the canonical AbortError when the external signal aborts even if the host token was not cancelled", async () => { + // Queue the host response but expect the original code never to consume it: + // the post-init re-check must abort before any sendRequest call. + mockLanguageModelChat.sendRequest.mockImplementationOnce(async () => ({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + })) + + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + // Neutralise the bridge's token cancel before aborting the signal so the + // host token flag stays false: only the signal's own aborted state is + // true, the exact operand the right-hand side of isAborted() must see. + ;(tokenSourceInstance().cancel as Mock).mockImplementation(() => {}) + controller.abort() + + // The token-agnostic mock stream completes, but the post-init re-check + // must still see the aborted external signal and surface the canonical + // abort error before any host request. + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", ) - noFamilyHandler.dispose() + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + + // The original path never made a host request, so the queued response + // was left unconsumed; reset the mock so it cannot leak into later tests. + mockLanguageModelChat.sendRequest.mockReset() }) - it("falls back to the derived window when the static row exists but maxInputTokens is non-positive", () => { - // A curated row exists but its maxInputTokens is <= 0, so the `> 0` guard fails and the gate - // falls back to getModel().info.contextWindow. - const family = "claude-opus-4.8" - const original = vscodeLlmModels[family].maxInputTokens + it("should cancel the token when timeoutMs elapses", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + await streamGate + yield "Completed text" + })(), + }) + + handler["client"] = mockLanguageModelChat + + vi.useFakeTimers() try { - ;(vscodeLlmModels[family] as { maxInputTokens: number }).maxInputTokens = 0 - const guardHandler = new VsCodeLmHandler({ - vsCodeLmModelSelector: { vendor: "copilot", family }, - }) - // Leave the client unset so `family` resolves from the selector, forcing the zeroed - // static row to be read instead of a live client's family. - guardHandler["client"] = null - expect(guardHandler.getCondenseContextWindow()).toBe(guardHandler.getModel().info.contextWindow) - expect(guardHandler.getCondenseContextWindow()).toBe(openAiModelInfoSaneDefaults.contextWindow) - guardHandler.dispose() + const promise = handler.completePrompt("Test prompt", { timeoutMs: 5000 }) + await vi.advanceTimersByTimeAsync(5000) + + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) + + // The token-agnostic mock stream completes, but the cancelled token + // still makes the completion abort. + releaseStream() + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) } finally { - ;(vscodeLlmModels[family] as { maxInputTokens: number }).maxInputTokens = original + vi.useRealTimers() } }) - }) - describe("countTokens", () => { - beforeEach(() => { - handler["client"] = mockLanguageModelChat - }) + it("should reject with an AbortError when the timeout fires during client initialization", async () => { + // Gate client initialization so the timeout timer (started before getClient()) + // can fire while getClient() is still pending. + let releaseClient: () => void = () => {} + const clientGate = new Promise((resolve) => { + releaseClient = resolve + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), + ) + handler["client"] = null - it("should count tokens when called outside of an active request", async () => { - // Ensure no active request cancellation token exists - handler["currentRequestCancellation"] = null + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 5000 }) - mockLanguageModelChat.countTokens.mockResolvedValueOnce(42) + // Advance past the timeout while getClient() is still pending: the timer + // fires and cancels the request token. + await vi.advanceTimersByTimeAsync(5000) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() - const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello world" }] - const result = await handler.countTokens(content) + // Release the client; the post-init re-check must abort before sendRequest. + releaseClient() - expect(result).toBe(42) - expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("Hello world", expect.any(Object)) + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + } finally { + vi.useRealTimers() + } }) - it("should count tokens when called during an active request", async () => { - // Simulate an active request with a cancellation token - const mockCancellation = { - token: { isCancellationRequested: false, onCancellationRequested: vi.fn() }, - cancel: vi.fn(), - dispose: vi.fn(), - } - handler["currentRequestCancellation"] = mockCancellation as unknown as vscode.CancellationTokenSource + it("should apply both the abort signal and timeoutMs", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - mockLanguageModelChat.countTokens.mockResolvedValueOnce(50) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) - const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Test content" }] - const result = await handler.countTokens(content) + handler["client"] = mockLanguageModelChat - expect(result).toBe(50) - expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("Test content", mockCancellation.token) + const controller = new AbortController() + const result = await handler.completePrompt("Test prompt", { + abortSignal: controller.signal, + timeoutMs: 10000, + }) + + expect(result).toBe("Completed text") + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) }) - it("should return 0 when no client is available", async () => { - handler["client"] = null - handler["currentRequestCancellation"] = null + it("should not treat a zero timeoutMs as an immediate timeout", async () => { + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + await streamGate + yield "Completed text" + })(), + }) - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + handler["client"] = mockLanguageModelChat - const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello" }] - const result = await handler.countTokens(content) + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 0 }) + for (let i = 0; i < 20; i++) { + await Promise.resolve() + } + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) - expect(result).toBe(0) - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : No client available for token counting", - ) + // timeoutMs: 0 must be treated as "no timeout": the condition + // timeoutMs > 0 must not schedule a zero-delay timer. + expect(vi.getTimerCount()).toBe(0) + await vi.advanceTimersByTimeAsync(10) - consoleWarnSpy.mockRestore() + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } }) - it("should handle image blocks with placeholder", async () => { - handler["currentRequestCancellation"] = null - mockLanguageModelChat.countTokens.mockResolvedValueOnce(5) + it("should not treat a negative timeoutMs as an immediate timeout", async () => { + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + await streamGate + yield "Completed text" + })(), + }) - const content: Anthropic.Messages.ContentBlockParam[] = [ - { type: "image", source: { type: "base64", media_type: "image/png", data: "abc" } }, - ] - const result = await handler.countTokens(content) + handler["client"] = mockLanguageModelChat - expect(result).toBe(5) - expect(mockLanguageModelChat.countTokens).toHaveBeenCalledWith("[IMAGE]", expect.any(Object)) - }) + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: -1 }) + for (let i = 0; i < 20; i++) { + await Promise.resolve() + } + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) - it("should return 0 and log when empty text is provided to internalCountTokens", async () => { - handler["currentRequestCancellation"] = null - const consoleDebugSpy = vi.spyOn(console, "debug").mockImplementation(() => {}) + // timeoutMs: -1 must be treated as "no timeout" like 0: the condition + // timeoutMs > 0 must not schedule a timer for a negative value. + expect(vi.getTimerCount()).toBe(0) + await vi.advanceTimersByTimeAsync(10) - // @ts-ignore – access private method for coverage of line 234 - const result = await handler["internalCountTokens"]("") + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } + }) - expect(result).toBe(0) - expect(consoleDebugSpy).toHaveBeenCalledWith( - "Zoo Code : Empty text provided for token counting", - ) + it("should schedule a fractional timeoutMs cancellation rounded up", async () => { + // Node truncates non-integer setTimeout delays (1500.5 fires at + // 1500 ms); the guard must round up so the delay matches + // mergeAbortSignalAndTimeout's 1501. Pin the delay handed to + // setTimeout rather than relying on fake-timer semantics. + const setTimeoutSpy = vi.spyOn(global, "setTimeout") + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + yield "Completed text" + })(), + }) - consoleDebugSpy.mockRestore() + handler["client"] = mockLanguageModelChat + + try { + const result = await handler.completePrompt("Test prompt", { timeoutMs: 1500.5 }) + expect(result).toBe("Completed text") + expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 1501) + } finally { + setTimeoutSpy.mockRestore() + } }) - it("should return 0 and log when non-numeric token count is received", async () => { - handler["currentRequestCancellation"] = null - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + it("should not treat an Infinity timeoutMs as an immediate timeout", async () => { + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + await streamGate + yield "Completed text" + })(), + }) - mockLanguageModelChat.countTokens.mockResolvedValueOnce("not-a-number" as unknown as number) + handler["client"] = mockLanguageModelChat - const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "test" }] - const result = await handler.countTokens(content) + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: Infinity }) + for (let i = 0; i < 20; i++) { + await Promise.resolve() + } + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) - expect(result).toBe(0) - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Non-numeric token count received:", - "not-a-number", - ) + // Infinity must mean "no timeout": Node would coerce a non-finite delay to 1ms, so the guard must reject non-finite values before setTimeout. + expect(vi.getTimerCount()).toBe(0) + await vi.advanceTimersByTimeAsync(1000) + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) - consoleWarnSpy.mockRestore() + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } }) - it("should return 0 and log when negative token count is received", async () => { - handler["currentRequestCancellation"] = null - const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + it("should cap an overflowing timeoutMs instead of firing immediately", async () => { + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Completed text") + })(), + text: (async function* () { + await streamGate + yield "Completed text" + })(), + }) - mockLanguageModelChat.countTokens.mockResolvedValueOnce(-5) + handler["client"] = mockLanguageModelChat - const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "test" }] - const result = await handler.countTokens(content) + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 2_147_483_648 }) + for (let i = 0; i < 20; i++) { + await Promise.resolve() + } + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) - expect(result).toBe(0) - expect(consoleWarnSpy).toHaveBeenCalledWith( - "Zoo Code : Negative token count received:", - -5, - ) + // 2^31 ms overflows Node's timer clamp to 1ms; the guard caps the delay at 2^31-1 so the request is not cancelled within the first second. + expect(vi.getTimerCount()).toBe(1) + await vi.advanceTimersByTimeAsync(1000) + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) - consoleWarnSpy.mockRestore() + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } }) - }) - - describe("completePrompt", () => { - it("should complete single prompt", async () => { - const mockModel = { ...mockLanguageModelChat } - ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - const responseText = "Completed text" + it("should clear the timeout timer when the completion finishes before it elapses", async () => { mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ stream: (async function* () { - yield new vscode.LanguageModelTextPart(responseText) + yield new vscode.LanguageModelTextPart("Completed text") return })(), text: (async function* () { - yield responseText + yield "Completed text" return })(), }) - // Override the default client with our test client to ensure it uses - // the mock implementation rather than the default fallback handler["client"] = mockLanguageModelChat - const result = await handler.completePrompt("Test prompt") - expect(result).toBe(responseText) - expect(mockLanguageModelChat.sendRequest).toHaveBeenCalled() + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 10_000 }) + const result = await promise + expect(result).toBe("Completed text") + + // The finally must have cleared the timer: advancing the clock past + // the timeout must not cancel the (already finished) request. + await vi.advanceTimersByTimeAsync(10_000) + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } }) - it("should handle errors during completion", async () => { + it("should reject before client initialization when the signal is already aborted", async () => { + handler["client"] = null + const selectCallsBefore = (vscode.lm.selectChatModels as Mock).mock.calls.length + + const controller = new AbortController() + controller.abort() + + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) + + // The pre-abort short-circuit runs before getClient(): the host is never contacted. + expect((vscode.lm.selectChatModels as Mock).mock.calls.length).toBe(selectCallsBefore) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + }) + + it("should wrap non-abort completion errors without an AbortError name", async () => { const mockModel = { ...mockLanguageModelChat } ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("Completion failed")) + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("LM error")) + handler["client"] = mockLanguageModelChat - // Make sure we're using the mock client + const promise = handler.completePrompt("Test prompt") + await expect(promise).rejects.toSatisfy((error) => { + return ( + error instanceof Error && + error.name === "Error" && + error.message === "VSCode LM completion error: LM error" + ) + }) + }) + + it("should reject with an AbortError when the host raises a CancellationError", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + // The host rejects with a CancellationError while the token flag is not set + // (no abort signal, no timeout), so only the error type signals the abort. + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new vscode.CancellationError()) handler["client"] = mockLanguageModelChat const promise = handler.completePrompt("Test prompt") - await expect(promise).rejects.toThrow("VSCode LM completion error: Completion failed") + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "The VS Code LM request was aborted", + ) }) }) diff --git a/src/api/providers/utils/__tests__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index f779d48e1f..1e2021f4f9 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -275,6 +275,56 @@ describe("abort-signal utilities", () => { await vi.waitFor(() => expect(result?.aborted).toBe(true)) }) + + it("treats a non-finite timeoutMs as no timeout", () => { + const controller = new AbortController() + + // Infinity must not reach AbortSignal.timeout(), which throws + // ERR_OUT_OF_RANGE for non-integer delays. + expect(() => mergeAbortSignalAndTimeout(undefined, Infinity)).not.toThrow() + expect(mergeAbortSignalAndTimeout(undefined, Infinity)).toBeUndefined() + expect(mergeAbortSignalAndTimeout(controller.signal, Infinity)).toBe(controller.signal) + }) + + it("caps an overflowing timeoutMs instead of aborting immediately", async () => { + // Node clamps AbortSignal.timeout delays above 2^31-1 ms to 1 ms; + // the helper caps the delay so an overflowing timeout does not abort + // the request within the first millisecond. + const result = mergeAbortSignalAndTimeout(undefined, 2_147_483_648) + + expect(result).toBeInstanceOf(AbortSignal) + expect(result?.aborted).toBe(false) + + await new Promise((resolve) => setTimeout(resolve, 20)) + expect(result?.aborted).toBe(false) + }) + + it("rounds a fractional timeoutMs up before scheduling", async () => { + // Node validates AbortSignal.timeout delays as uint32: 1.5 or 0.5 + // would throw ERR_OUT_OF_RANGE. Pin the rounded delay that actually + // reaches AbortSignal.timeout (1.5 -> 2, 0.5 -> 1, never 0) so a + // round-down mutation reddens this test instead of still aborting + // eventually. + const timeoutSpy = vi.spyOn(AbortSignal, "timeout") + let result: AbortSignal | undefined + let tiny: AbortSignal | undefined + try { + expect(() => { + result = mergeAbortSignalAndTimeout(undefined, 1.5) + }).not.toThrow() + expect(timeoutSpy).toHaveBeenCalledWith(2) + expect(() => { + tiny = mergeAbortSignalAndTimeout(undefined, 0.5) + }).not.toThrow() + expect(timeoutSpy).toHaveBeenCalledWith(1) + } finally { + timeoutSpy.mockRestore() + } + expect(result).toBeInstanceOf(AbortSignal) + expect(result?.aborted).toBe(false) + expect(tiny).toBeInstanceOf(AbortSignal) + await vi.waitFor(() => expect(result?.aborted).toBe(true)) + }) }) describe("mergeAbortSignals", () => { diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index cbb7a13aaf..8917a24ba1 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -1,18 +1,25 @@ /** * Merge an optional external abort signal with an optional timeout. * - * Timeout values <= 0 are treated as disabled. The timeout is created via the + * Timeout values <= 0 and non-finite values are treated as disabled; finite + * delays are rounded up and capped at 2^31-1 ms (Node clamps larger delays + * to 1 ms and rejects non-integer ones). The timeout is created via the * native AbortSignal.timeout() API, which self-manages its timer — callers do * not need to (and cannot) clear it manually. */ export function mergeAbortSignalAndTimeout(externalSignal?: AbortSignal, timeoutMs?: number): AbortSignal | undefined { - const hasTimeout = typeof timeoutMs === "number" && timeoutMs > 0 + // Node's AbortSignal.timeout() throws ERR_OUT_OF_RANGE for non-integer + // delays (Infinity and fractions like 1500.5 fail its uint32 validation) + // and clamps delays above 2^31-1 ms to 1 ms, which would abort almost + // immediately; treat non-finite values as "no timeout", round fractional + // delays up to whole milliseconds, and cap the rest. + const hasTimeout = typeof timeoutMs === "number" && Number.isFinite(timeoutMs) && timeoutMs > 0 if (!hasTimeout) { return externalSignal } - const timeoutSignal = AbortSignal.timeout(timeoutMs) + const timeoutSignal = AbortSignal.timeout(Math.min(Math.max(1, Math.ceil(timeoutMs)), 2_147_483_647)) if (!externalSignal) { return timeoutSignal diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 7326d3ec8a..cf6f4cccfb 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -27,6 +27,7 @@ import { import { CONTEXT_WINDOW_EXCEEDED_STATUS } from "../../core/context/context-management/context-error-handling" import { BaseProvider } from "./base-provider" +import { createAbortError } from "./utils/abort-signal" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** @@ -759,6 +760,8 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan protected options: ApiHandlerOptions private client: vscode.LanguageModelChat | null private disposable: vscode.Disposable | null + // Provider label for the shared abort contract (createAbortError). + private readonly providerName = "VS Code LM" private currentRequestCancellation: vscode.CancellationTokenSource | null constructor(options: ApiHandlerOptions) { @@ -1070,64 +1073,29 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan ): ApiStream { // Ensure clean state before starting a new request this.ensureCleanState() - const client: vscode.LanguageModelChat = await this.getClient() - // Process messages - const cleanedMessages = messages.map((msg) => ({ - ...msg, - content: this.cleanMessageContent(msg.content), - })) - - // Keep context-window trimming on OUR side. Copilot's backend trims an over-window request - // without preserving tool_use/tool_result pairing, which orphans a tool_result and triggers a - // 400 ("unexpected tool_use_id"). See truncateToolResultsToFitWindow. - const contextWindowTokens = this.getCondenseContextWindow() - if (Number.isFinite(contextWindowTokens) && contextWindowTokens > 0) { - const toolSchemaChars = metadata?.tools ? JSON.stringify(metadata.tools).length : 0 - const rawBudgetChars = - contextWindowTokens * VSCODE_LM_INPUT_BUDGET_FRACTION * VSCODE_LM_BUDGET_CHARS_PER_TOKEN - - systemPrompt.length - - toolSchemaChars - // A system prompt or tool schema large enough to consume the whole budget would leave a - // non-positive budget, which disables trimming exactly when the request is most oversized. - const messagesBudgetChars = Math.max(MIN_TOOL_RESULT_CHARS, rawBudgetChars) - const { remainingChars } = truncateToolResultsToFitWindow(cleanedMessages, messagesBudgetChars) - - // Shrinking tool_results cannot always reach the budget: each keeps MIN_TOOL_RESULT_CHARS, - // and the excess may be non-tool content (a huge paste, tool_use inputs, or the system - // prompt) that we must not touch. Dropping messages here would orphan a tool_result from - // its tool_use — the exact 400 this guard exists to prevent — so fail loudly instead of - // sending a request we already know is over the window. - // Admission is judged against the RAW budget, not the clamped one: the clamp exists only - // to keep trimming productive, so accepting up to it would send a request the window - // genuinely cannot hold whenever the raw budget falls below MIN_TOOL_RESULT_CHARS. - if (remainingChars > rawBudgetChars) { - // `status` is what makes checkContextWindowExceededError recognise this as a context - // -window failure; without it the task takes its generic retry path and re-sends the - // same over-window history instead of condensing. - throw Object.assign( - new Error( - "Zoo Code : The request is too large for this model's context window " + - `(estimated ${remainingChars.toLocaleString("en-US")} characters against a budget of ` + - `${Math.max(0, Math.floor(rawBudgetChars)).toLocaleString("en-US")}), and it cannot be reduced further without ` + - "breaking tool-call pairing. Condense the conversation or start a new task.", - ), - { status: CONTEXT_WINDOW_EXCEEDED_STATUS }, - ) - } - } - - // Convert Anthropic messages to VS Code LM messages - const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ - vscode.LanguageModelChatMessage.Assistant(sanitizeSurrogates(systemPrompt)), - ...convertToVsCodeLmMessages(cleanedMessages), - ] + // The VS Code LanguageModelChat API cannot carry an AbortSignal, so the + // request cancellation is established before client initialization and a + // pre-aborted external signal is reported immediately instead of being + // sent to the host. + const externalAbortSignal = metadata?.abortSignal - // Initialize cancellation token for the request + // Initialize cancellation token for the request before getClient() so the + // client-initialization await is covered by the abort bridge and the + // finally cleanup below. this.currentRequestCancellation = new vscode.CancellationTokenSource() - - // Calculate input tokens before starting the stream - const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) + const cancellationTokenSource = this.currentRequestCancellation + + // Bridge the caller's abort signal (e.g. a task abort) into the request's + // cancellation token: the VS Code LM API cannot carry an AbortSignal + // directly, so cancellation is signalled to the host through the token. + // The listener is kept in a named const and removed in the finally block + // because { once: true } only detaches it when the signal actually aborts. + let onExternalAbort: (() => void) | undefined + if (externalAbortSignal) { + onExternalAbort = () => cancellationTokenSource.cancel() + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } // Accumulate the text and count at the end of the stream to reduce token counting overhead. let accumulatedText: string = "" @@ -1213,6 +1181,89 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } try { + // Fail fast if the caller already aborted before we even started: do not + // initialize or invoke the host request for a cancelled request. + if (externalAbortSignal?.aborted) { + cancellationTokenSource.cancel() + throw createAbortError(this.providerName) + } + + const client: vscode.LanguageModelChat = await this.getClient() + + // Re-check immediately after client initialization: the request-local token + // may have been cancelled while getClient() was pending — either because the + // external signal aborted (bridged into the token) or because a newer request + // superseded this one and cancelled its token. Bail before counting input + // tokens or invoking the host so no work happens for a cancelled request. + // The token is the superset here: the bridge makes every external abort cancel it. + if (cancellationTokenSource.token.isCancellationRequested) { + cancellationTokenSource.cancel() + throw createAbortError(this.providerName) + } + + // Process messages + const cleanedMessages = messages.map((msg) => ({ + ...msg, + content: this.cleanMessageContent(msg.content), + })) + + // Keep context-window trimming on OUR side. Copilot's backend trims an over-window request + // without preserving tool_use/tool_result pairing, which orphans a tool_result and triggers a + // 400 ("unexpected tool_use_id"). See truncateToolResultsToFitWindow. + const contextWindowTokens = this.getCondenseContextWindow() + if (Number.isFinite(contextWindowTokens) && contextWindowTokens > 0) { + const toolSchemaChars = metadata?.tools ? JSON.stringify(metadata.tools).length : 0 + const rawBudgetChars = + contextWindowTokens * VSCODE_LM_INPUT_BUDGET_FRACTION * VSCODE_LM_BUDGET_CHARS_PER_TOKEN - + systemPrompt.length - + toolSchemaChars + // A system prompt or tool schema large enough to consume the whole budget would leave a + // non-positive budget, which disables trimming exactly when the request is most oversized. + const messagesBudgetChars = Math.max(MIN_TOOL_RESULT_CHARS, rawBudgetChars) + const { remainingChars } = truncateToolResultsToFitWindow(cleanedMessages, messagesBudgetChars) + + // Shrinking tool_results cannot always reach the budget: each keeps MIN_TOOL_RESULT_CHARS, + // and the excess may be non-tool content (a huge paste, tool_use inputs, or the system + // prompt) that we must not touch. Dropping messages here would orphan a tool_result from + // its tool_use — the exact 400 this guard exists to prevent — so fail loudly instead of + // sending a request we already know is over the window. + // Admission is judged against the RAW budget, not the clamped one: the clamp exists only + // to keep trimming productive, so accepting up to it would send a request the window + // genuinely cannot hold whenever the raw budget falls below MIN_TOOL_RESULT_CHARS. + if (remainingChars > rawBudgetChars) { + // `status` is what makes checkContextWindowExceededError recognise this as a context + // -window failure; without it the task takes its generic retry path and re-sends the + // same over-window history instead of condensing. + throw Object.assign( + new Error( + "Zoo Code : The request is too large for this model's context window " + + `(estimated ${remainingChars.toLocaleString("en-US")} characters against a budget of ` + + `${Math.max(0, Math.floor(rawBudgetChars)).toLocaleString("en-US")}), and it cannot be reduced further without ` + + "breaking tool-call pairing. Condense the conversation or start a new task.", + ), + { status: CONTEXT_WINDOW_EXCEEDED_STATUS }, + ) + } + } + + // Convert Anthropic messages to VS Code LM messages + const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ + vscode.LanguageModelChatMessage.Assistant(sanitizeSurrogates(systemPrompt)), + ...convertToVsCodeLmMessages(cleanedMessages), + ] + + // Calculate input tokens before starting the stream + const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) + + // Re-check the request-local token after counting: the external signal may + // have aborted while counting was in flight (bridged into the token) or a + // newer request may have superseded this one and cancelled its token. Bail + // before invoking the host so no work happens for a cancelled request. The + // token is the superset here: the bridge makes every external abort cancel it. + if (cancellationTokenSource.token.isCancellationRequested) { + throw createAbortError(this.providerName) + } + // Create the response stream with required options const requestOptions: vscode.LanguageModelChatRequestOptions = { justification: `Zoo Code would like to use '${client.name}' from '${client.vendor}', Click 'Allow' to proceed.`, @@ -1222,11 +1273,18 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan const response: vscode.LanguageModelChatResponse = await client.sendRequest( vsCodeLmMessages, requestOptions, - this.currentRequestCancellation.token, + cancellationTokenSource.token, ) // Consume the stream and handle both text and tool call chunks for await (const chunk of response.stream) { + // A late abort while consuming must stop the stream instead of + // yielding stale chunks. The request-local token also covers local + // supersession (a newer request cancels this one), which the external + // signal alone cannot see. + if (externalAbortSignal?.aborted || cancellationTokenSource.token.isCancellationRequested) { + throw createAbortError(this.providerName) + } if (chunk instanceof vscode.LanguageModelTextPart) { // Validate text part value if (typeof chunk.value !== "string") { @@ -1326,6 +1384,17 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } } + // Guard against a quiet stream end after the request was aborted or superseded: + // the in-loop check only fires when the host delivers another chunk, so a stream + // that ends normally after cancellation would otherwise complete with partial + // text and a usage chunk instead of rejecting with the canonical abort. + // The guard runs before the salvage flush: an aborted request must not emit + // any further content, and the per-request salvage state is moot once the + // request is rejected. + if (externalAbortSignal?.aborted || cancellationTokenSource.token.isCancellationRequested) { + throw createAbortError(this.providerName) + } + // Flush any leaked tool-call recovery state accumulated during streaming. yield* flushSalvage() @@ -1339,10 +1408,23 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan outputTokens: totalOutputTokens, } } catch (error: unknown) { - this.ensureCleanState() + // When this request is already cancelled while getClient() was pending or + // rejected — the external signal aborted (bridged into the token) or a + // newer request superseded this one and cancelled its token — surface a + // standard abort error instead of leaking the raw client-initialization + // failure. completePrompt normalizes the same case through isAborted(), + // which includes the token state; both methods follow that contract. + if (externalAbortSignal?.aborted || cancellationTokenSource.token.isCancellationRequested) { + throw createAbortError(this.providerName) + } if (error instanceof vscode.CancellationError) { - throw new Error("Zoo Code : Request cancelled by user") + // The host rejected because the request was cancelled: either the + // bridged external signal aborted or the user cancelled the request + // in VS Code. Both are aborts, so surface a standard abort error. + const abortError = new Error("Zoo Code : Request cancelled by user") + abortError.name = "AbortError" + throw abortError } if (error instanceof Error) { @@ -1365,6 +1447,27 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan console.error("Zoo Code : Unknown stream error:", errorMessage) throw new Error(`Zoo Code : Response stream error: ${errorMessage}`) } + } finally { + // Detach the abort bridge listener on every path (success, error, and + // early consumer break); { once: true } alone would leak it when the + // request completes without the signal ever aborting. + if (externalAbortSignal && onExternalAbort) { + externalAbortSignal.removeEventListener("abort", onExternalAbort) + } + // Cancel before disposing: VS Code's CancellationTokenSource.dispose() + // frees resources without cancelling the token, so a premature consumer + // closure (break/return) would otherwise leave the host request running and + // consuming model quota. Cancel is idempotent, so this is a no-op on paths + // where the token was already cancelled (aborted or timed-out request). + cancellationTokenSource.cancel() + // Dispose the request-local source. Clear the shared field only if it still + // points at this request's source: a newer request may have replaced it, and + // disposing the shared field here would cancel and dispose the newer request's + // token. + cancellationTokenSource.dispose() + if (this.currentRequestCancellation === cancellationTokenSource) { + this.currentRequestCancellation = null + } } } @@ -1445,13 +1548,68 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan return this.getModel().info.contextWindow } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // The VS Code LanguageModelChat API cannot carry an AbortSignal: sendRequest + // only accepts a CancellationToken. Bridge the external signal and timeout + // into a request-local CancellationTokenSource instead. Cancellation is + // established before client initialization so the getClient() await is covered + // by the timeout, the abort bridge, and the finally cleanup below. + const tokenSource = new vscode.CancellationTokenSource() + const externalAbortSignal = options?.abortSignal + + // Apply the timeout only when it is a positive value: cancelling at once for a + // zero/negative timeout would abort every such request immediately. Starting + // the timer before getClient() means a timeout that fires during a slow client + // lookup still cancels the request before any sendRequest call. + let timeoutId: ReturnType | undefined + const timeoutMs = options?.timeoutMs + if (timeoutMs !== undefined && Number.isFinite(timeoutMs) && timeoutMs > 0) { + // Node clamps delays above 2^31-1 ms to 1 ms (and Infinity coerces the + // same way), which would cancel almost immediately, and truncates + // fractional delays; round up and cap so the semantics match + // mergeAbortSignalAndTimeout, and treat non-finite values as + // "no timeout". + timeoutId = setTimeout(() => tokenSource.cancel(), Math.min(Math.ceil(timeoutMs), 2_147_483_647)) + } + + // Bridge the external abort signal: a pre-aborted signal cancels the token + // immediately, otherwise a one-shot listener relays the abort to the host. + let onAbort: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + tokenSource.cancel() + } else { + onAbort = () => tokenSource.cancel() + externalAbortSignal.addEventListener("abort", onAbort, { once: true }) + } + } + + // The request counts as aborted when the host-side token was cancelled + // (timeout or bridged signal) or when the external signal itself aborted. + const isAborted = () => + tokenSource.token.isCancellationRequested === true || externalAbortSignal?.aborted === true + try { + // Fail fast if the caller already aborted before we even started: do not + // initialize or invoke the host request for a cancelled request. + if (isAborted()) { + throw createAbortError(this.providerName) + } + const client = await this.getClient() + + // Re-check after client initialization: the signal may have aborted (or the + // timeout fired) while getClient() was pending. Bail before invoking the host + // so a timeout that fired during a slow client lookup can never lead to a + // sendRequest call. + if (isAborted()) { + throw createAbortError(this.providerName) + } + const response = await client.sendRequest( [vscode.LanguageModelChatMessage.User(prompt)], {}, - new vscode.CancellationTokenSource().token, + tokenSource.token, ) let result = "" for await (const chunk of response.stream) { @@ -1459,12 +1617,36 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan result += chunk.value } } + + // Guard against a quiet completion after the request was aborted. + if (isAborted()) { + throw createAbortError(this.providerName) + } + return result } catch (error) { + // Report an aborted request (external signal, timeout, or host + // cancellation) as a standard AbortError instead of a generic + // completion error. A host CancellationError is treated as a cancellation + // even if the token flag was not observed (e.g. the host cancelled the + // request through the token without the flag being set). + if (isAborted() || error instanceof vscode.CancellationError) { + throw createAbortError(this.providerName) + } + if (error instanceof Error) { throw new Error(`VSCode LM completion error: ${error.message}`) } throw error + } finally { + // clearTimeout(undefined) is a spec no-op, so clear unconditionally. + clearTimeout(timeoutId) + // { once: true } only detaches the listener when the signal actually + // aborts, so remove it explicitly on every path. + if (externalAbortSignal && onAbort) { + externalAbortSignal.removeEventListener("abort", onAbort) + } + tokenSource.dispose() } } }