diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 423f119f14..d6a5a2c2a0 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -26,12 +26,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(), } }), @@ -65,6 +70,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 { clearAllMocks } from "../../../test-utils/reset" const mockLanguageModelChat = { @@ -78,6 +84,18 @@ const mockLanguageModelChat = { countTokens: vi.fn(), } +/** + * 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", () => { let handler: VsCodeLmHandler const defaultOptions: ApiHandlerOptions = { @@ -470,6 +488,412 @@ describe("VsCodeLmHandler", () => { ) }) + 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 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", + ) + + // 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 + + 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") + + // 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 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" }] + + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Hello") + })(), + text: (async function* () { + await streamGate + yield "Hello" + })(), + }) + + 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") + + // The bridge relayed the abort to the request cancellation token. + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + }) + + 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" }] + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Hello") + })(), + text: (async function* () { + yield "Hello" + })(), + }) + + 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(addEventListenerSpy).toHaveBeenCalledTimes(1) + expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + + 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" }] + + // 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 new vscode.LanguageModelTextPart("Hello") + await streamGate + })(), + text: (async function* () { + yield "Hello" + await streamGate + })(), + }) + + const stream = handler.createMessage(systemPrompt, messages) + 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) + + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + expect(tokenSource.dispose).toHaveBeenCalled() + }) + + 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 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", + ) + + 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) + }) + + 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", + ) + + // 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 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" }] + + 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), + ) + + 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") + + // 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 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[] = [ @@ -1039,6 +1463,352 @@ 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", expect.any(Function)) + }) + + 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") + + // 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 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", + ) + } 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", + ) + 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 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") + + // 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") + }) }) describe("cleanMessageContent / deepClean", () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index c657e6c0d6..259e69aa70 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -375,30 +375,70 @@ 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), - })) - - // Convert Anthropic messages to VS Code LM messages - const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ - vscode.LanguageModelChatMessage.Assistant(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() + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } + + const client: vscode.LanguageModelChat = await this.getClient() + + // Re-check immediately after client initialization: the signal may have + // aborted while getClient() was pending. Bail before counting input + // tokens or invoking the host so no work happens for a cancelled request. + if (externalAbortSignal?.aborted) { + cancellationTokenSource.cancel() + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } + + // Process messages + const cleanedMessages = messages.map((msg) => ({ + ...msg, + content: this.cleanMessageContent(msg.content), + })) + + // Convert Anthropic messages to VS Code LM messages + const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ + vscode.LanguageModelChatMessage.Assistant(systemPrompt), + ...convertToVsCodeLmMessages(cleanedMessages), + ] + + // Calculate input tokens before starting the stream + const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) + // 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.`, @@ -408,11 +448,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. + if (externalAbortSignal?.aborted) { + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } if (chunk instanceof vscode.LanguageModelTextPart) { // Validate text part value if (typeof chunk.value !== "string") { @@ -482,10 +529,23 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan outputTokens: totalOutputTokens, } } catch (error: unknown) { - this.ensureCleanState() + // When the external signal wins during client initialization (the signal + // was already aborted while getClient() was pending or rejected), surface a + // standard abort error instead of wrapping the getClient() failure as a + // generic stream error. + if (externalAbortSignal?.aborted) { + 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) { @@ -508,6 +568,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 (onExternalAbort !== undefined) { + 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 + } } } @@ -589,12 +670,65 @@ 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()) { + const abortError = new Error("VSCode LM completion aborted") + 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()) { + const abortError = new Error("VSCode LM completion aborted") + 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) { @@ -602,12 +736,41 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan result += chunk.value } } + + // Guard against a quiet completion after the request was aborted. + if (isAborted()) { + const abortError = new Error("VSCode LM completion aborted") + 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 { + if (timeoutId !== undefined) { + clearTimeout(timeoutId) + } + // { once: true } only detaches the listener when the signal actually + // aborts, so remove it explicitly on every path. + if (onAbort !== undefined) { + externalAbortSignal?.removeEventListener("abort", onAbort) + } + tokenSource.dispose() } } }