diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 86a9a30c38..01e8c120da 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, @@ -82,6 +107,7 @@ import type { ApiHandlerOptions } from "../../../shared/api" 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" @@ -96,7 +122,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", () => { @@ -683,6 +723,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[] = [ @@ -881,178 +1071,938 @@ 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 === "Zoo Code : Request aborted", ) - consoleWarnSpy.mockRestore() + // The bridge relayed the abort to the request cancellation token. + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() }) - 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 === "Zoo Code : Request aborted", + ) - const badToolCall = new vscode.LanguageModelToolCallPart("call-1", "valid-name", circularInput) - mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ - stream: (async function* () { - yield badToolCall - return - })(), - text: (async function* () { + // 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 === "Zoo Code : Request 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 "Hello!" + return + })(), + }) + + 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 + } + + // 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 === "Zoo Code : Request aborted", + ) + + await streamB.return(undefined) + expect(handler["currentRequestCancellation"]).toBeNull() + consoleErrorSpy.mockRestore() + }) + + 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 === "Zoo Code : Request 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 })(), @@ -1450,6 +2400,454 @@ describe("VsCodeLmHandler", () => { const promise = handler.completePrompt("Test prompt") await expect(promise).rejects.toThrow("VSCode LM completion error: Completion failed") }) + + it("should work without options (backward compatible)", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Completed text") + }) + + it("should bridge the abort signal to a fresh cancellation token and dispose it", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + 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 result = await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(result).toBe("Completed text") + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).not.toHaveBeenCalled() + expect(tokenSource.dispose).toHaveBeenCalled() + }) + + it("should attach and detach the abort listener around the completion", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + 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 addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + 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 reject with an AbortError when the signal is already aborted", async () => { + handler["client"] = mockLanguageModelChat + + 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 === "VSCode LM completion aborted", + ) + + // Fails fast before invoking the host: the pre-aborted signal cancels the token. + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + }) + + 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" + })(), + }) + + 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 === "VSCode LM completion aborted", + ) + }) + + 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 === "VSCode LM completion aborted", + ) + 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("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 { + 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 === "VSCode LM completion aborted", + ) + } finally { + vi.useRealTimers() + } + }) + + 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 + + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 5000 }) + + // 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() + + // Release the client; the post-init re-check must abort before sendRequest. + releaseClient() + + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "VSCode LM completion aborted", + ) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + } finally { + vi.useRealTimers() + } + }) + + it("should apply both the abort signal and timeoutMs", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + 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 result = await handler.completePrompt("Test prompt", { + abortSignal: controller.signal, + timeoutMs: 10000, + }) + + expect(result).toBe("Completed text") + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) + }) + + 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" + })(), + }) + + handler["client"] = mockLanguageModelChat + + 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) + + // 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) + + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } + }) + + 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" + })(), + }) + + handler["client"] = mockLanguageModelChat + + 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) + + // 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) + + releaseStream() + const result = await promise + expect(result).toBe("Completed text") + expect(tokenSourceInstance().token.isCancellationRequested).toBe(false) + } finally { + vi.useRealTimers() + } + }) + + it("should clear the timeout timer when the completion finishes before it elapses", async () => { + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + 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 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 === "VSCode LM completion 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("LM error")) + handler["client"] = mockLanguageModelChat + + 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.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "VSCode LM completion aborted", + ) + }) }) describe("cleanMessageContent / deepClean", () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 2771227680..aeb783729a 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -1041,69 +1041,129 @@ 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 = "" 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() + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("Zoo Code : Request aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + + 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() + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("Zoo Code : Request aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + + // 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) { + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("Zoo Code : Request aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + // 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.`, @@ -1113,11 +1173,22 @@ 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) { + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("Zoo Code : Request aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } if (chunk instanceof vscode.LanguageModelTextPart) { // Validate text part value if (typeof chunk.value !== "string") { @@ -1178,6 +1249,18 @@ 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. + if (externalAbortSignal?.aborted || cancellationTokenSource.token.isCancellationRequested) { + // Stryker disable next-line StringLiteral: caught by the abort branch in the catch below, which re-throws its own canonical abort error with the identical message, so this string is never observed + const abortError = new Error("Zoo Code : Request aborted") + // Stryker disable next-line StringLiteral: caught by the abort branch in the catch below, which re-throws its own canonical abort error with the identical name, so this string is never observed + abortError.name = "AbortError" + throw abortError + } + // Count tokens in the accumulated text after stream completion const totalOutputTokens: number = await this.internalCountTokens(accumulatedText) @@ -1188,10 +1271,25 @@ 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) { + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } 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) { @@ -1214,6 +1312,28 @@ 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. + // Stryker disable next-line LogicalOperator: set only when the signal is truthy, so && and || are equivalent + 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 + } } } @@ -1295,12 +1415,69 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } 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 + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + timeoutId = setTimeout(() => tokenSource.cancel(), options.timeoutMs) + } + + // 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()) { + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("VSCode LM completion aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + 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()) { + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("VSCode LM completion aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + const response = await client.sendRequest( [vscode.LanguageModelChatMessage.User(prompt)], {}, - new vscode.CancellationTokenSource().token, + tokenSource.token, ) let result = "" for await (const chunk of response.stream) { @@ -1308,12 +1485,43 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan result += chunk.value } } + + // Guard against a quiet completion after the request was aborted. + if (isAborted()) { + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + const abortError = new Error("VSCode LM completion aborted") + // Stryker disable next-line StringLiteral: caught below; the catch re-throws its own canonical abort error here + abortError.name = "AbortError" + throw abortError + } + 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) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + 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. + // Stryker disable next-line LogicalOperator: set only when the signal is truthy, so && and || are equivalent + if (externalAbortSignal && onAbort) { + externalAbortSignal.removeEventListener("abort", onAbort) + } + tokenSource.dispose() } } }