From 71b0a58cc9f1c85cdec769a25fd14e1aa190bd72 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 07:50:34 +0800 Subject: [PATCH 1/9] feat(api): abort signal support for vscode-lm + fake-ai completion tests - completePrompt: bridge CompletePromptOptions.abortSignal and timeoutMs into a request-local vscode.CancellationTokenSource (the VS Code LanguageModelChat API only accepts a CancellationToken, not an AbortSignal); apply timeoutMs only when it is a positive value - completePrompt: report aborted requests (external signal, timeout, or host cancellation) as errors with name = "AbortError" on both the error and success paths; remove the abort listener and dispose the token source in finally - createMessage: bridge metadata.abortSignal into the internal request CancellationTokenSource (Bedrock pattern: pre-aborted guard + { once: true } listener stored in a named const), fail fast with an AbortError when the signal is already aborted, surface host CancellationError with name = "AbortError", and detach the listener / dispose the source in finally - vscode-lm.spec.ts: add pre-aborted and mid-flight abort, timeout, listener attach/detach, and backward-compatibility tests - fake-ai.spec.ts: option pass-through tests already merged on main via #901; verified green without changes --- src/api/providers/__tests__/vscode-lm.spec.ts | 322 ++++++++++++++++++ src/api/providers/vscode-lm.ts | 103 +++++- 2 files changed, 422 insertions(+), 3 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 423f119f14..61a3898a14 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -65,6 +65,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 +79,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 +483,96 @@ 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 stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + 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 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 stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + 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. + controller.abort() + releaseStream() + await expect(firstChunk).resolves.toEqual({ done: false, value: { type: "text", text: "Hello" } }) + + // 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 stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + 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 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 +1142,225 @@ 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 () => { + 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() + controller.abort() + + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + expect(tokenSource.dispose).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, then let the mock stream finish. + controller.abort() + releaseStream() + + await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + }) + + 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 still completes, so the (unaborted-signal) + // completion resolves normally. + releaseStream() + await expect(promise).resolves.toBe("Completed text") + } 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 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" + ) + }) + }) }) describe("cleanMessageContent / deepClean", () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index c657e6c0d6..a002cc8eb5 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -375,6 +375,17 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan ): ApiStream { // Ensure clean state before starting a new request this.ensureCleanState() + + // The VS Code LanguageModelChat API cannot carry an AbortSignal, so a + // pre-aborted external signal is reported immediately instead of being + // sent to the host. + const externalAbortSignal = metadata?.abortSignal + if (externalAbortSignal?.aborted) { + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } + const client: vscode.LanguageModelChat = await this.getClient() // Process messages @@ -391,6 +402,18 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // Initialize cancellation token for the request this.currentRequestCancellation = new vscode.CancellationTokenSource() + 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 }) + } // Calculate input tokens before starting the stream const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) @@ -485,7 +508,12 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan this.ensureCleanState() 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 +536,17 @@ 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) + } + if (this.currentRequestCancellation) { + this.currentRequestCancellation.dispose() + this.currentRequestCancellation = null + } } } @@ -589,12 +628,43 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + const client = await this.getClient() + + // 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. + 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. + 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 { - const client = await this.getClient() 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 +672,39 @@ 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. + if (isAborted()) { + 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() } } } From 8ef34b1f3a23bd7d050bd0d6bb254f2e5602ceb2 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 09:31:20 +0800 Subject: [PATCH 2/9] fix(api): address CodeRabbit review on vscode-lm abort handling - createMessage: sendRequest now uses the request-local cancellation source, the finally block disposes that local source and clears the shared field only when it still points at this request, and the error path no longer calls ensureCleanState (prevents an older finishing request from cancelling/disposing a newer request's token) - completePrompt: a pre-aborted signal now fails fast before getClient() and again before sendRequest(), so a cancelled request never initializes or invokes the host - completePrompt: a host vscode.CancellationError is normalized to an AbortError alongside isAborted() - spec: the timeout test now expects an AbortError (the cancelled token aborts the completion); the pre-abort test asserts sendRequest is never called; added a CancellationError -> AbortError rejection test; the mock CancellationTokenSource cancel() now flips isCancellationRequested to match the real API --- src/api/providers/__tests__/vscode-lm.spec.ts | 56 ++++++++++--------- src/api/providers/vscode-lm.ts | 34 ++++++++--- 2 files changed, 58 insertions(+), 32 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 61a3898a14..9fa3a87c66 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(), } }), @@ -1220,20 +1225,6 @@ describe("VsCodeLmHandler", () => { }) it("should reject with an AbortError when the signal is already aborted", 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() @@ -1242,9 +1233,9 @@ describe("VsCodeLmHandler", () => { const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") - const tokenSource = tokenSourceInstance() - expect(tokenSource.cancel).toHaveBeenCalled() - expect(tokenSource.dispose).toHaveBeenCalled() + // Fails fast before invoking the host: no cancellation token source is + // created and the sendRequest is never called. + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) }) it("should reject with an AbortError when the signal aborts mid-flight", async () => { @@ -1309,10 +1300,12 @@ describe("VsCodeLmHandler", () => { expect(tokenSource.cancel).toHaveBeenCalled() expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) - // The token-agnostic mock stream still completes, so the (unaborted-signal) - // completion resolves normally. + // The token-agnostic mock stream completes, but the cancelled token + // still makes the completion abort. releaseStream() - await expect(promise).resolves.toBe("Completed text") + await expect(promise).rejects.toSatisfy( + (error) => error instanceof Error && error.name === "AbortError", + ) } finally { vi.useRealTimers() } @@ -1361,6 +1354,19 @@ describe("VsCodeLmHandler", () => { ) }) }) + + 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 a002cc8eb5..18ed09a62e 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -431,7 +431,7 @@ 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 @@ -505,8 +505,6 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan outputTokens: totalOutputTokens, } } catch (error: unknown) { - this.ensureCleanState() - if (error instanceof vscode.CancellationError) { // The host rejected because the request was cancelled: either the // bridged external signal aborted or the user cancelled the request @@ -543,8 +541,12 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan if (onExternalAbort !== undefined) { externalAbortSignal?.removeEventListener("abort", onExternalAbort) } - if (this.currentRequestCancellation) { - this.currentRequestCancellation.dispose() + // 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 } } @@ -628,6 +630,14 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Fail fast if the caller already aborted before we even started: do not + // initialize or invoke the host request for a cancelled request. + if (options?.abortSignal?.aborted) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + const client = await this.getClient() // The VS Code LanguageModelChat API cannot carry an AbortSignal: sendRequest @@ -661,6 +671,14 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan tokenSource.token.isCancellationRequested === true || externalAbortSignal?.aborted === true try { + // Re-check before invoking the host: the signal may have aborted while the + // client was being initialized above. + if (externalAbortSignal?.aborted) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + const response = await client.sendRequest( [vscode.LanguageModelChatMessage.User(prompt)], {}, @@ -684,8 +702,10 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } catch (error) { // Report an aborted request (external signal, timeout, or host // cancellation) as a standard AbortError instead of a generic - // completion error. - if (isAborted()) { + // 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 From aea94647451b2e542051fb30104ce6bd545d0be3 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 10:35:04 +0800 Subject: [PATCH 3/9] fix(api): handle vscode-lm createMessage aborts during client init and streaming - createMessage re-checks the external abort signal after client initialization (and before sendRequest), cancelling the local token source and throwing an AbortError when the signal aborted while getClient() was pending - createMessage re-checks the external abort signal at the top of the stream consumption loop so a late abort stops the stream instead of yielding stale chunks (the bridged listener still covers the normal mid-flight case) - spec: the mid-flight abort test now expects the stream to stop with an AbortError; added a regression test where client initialization is gated on a release promise and the signal aborts in that window - the generator rejects with AbortError and sendRequest is never called --- src/api/providers/__tests__/vscode-lm.spec.ts | 42 ++++++++++++++++++- src/api/providers/vscode-lm.ts | 17 ++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 9fa3a87c66..9fd4ea3aa1 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -508,6 +508,43 @@ describe("VsCodeLmHandler", () => { 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 + // pre-abort check and listener attachment happened before this window. + 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 stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const firstChunk = stream.next() + + // Release the client, then abort before the generator can reach 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. + expect(handler["client"]).not.toBeNull() + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + + // The local token source was cancelled for the request that never started. + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + }) + 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" }] @@ -537,10 +574,11 @@ describe("VsCodeLmHandler", () => { 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. + // (token-agnostic) mock stream finish. The late abort must stop the stream + // instead of yielding stale chunks. controller.abort() releaseStream() - await expect(firstChunk).resolves.toEqual({ done: false, value: { type: "text", text: "Hello" } }) + 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() diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 18ed09a62e..748159e7a6 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -422,6 +422,16 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan let accumulatedText: string = "" try { + // Re-check after client initialization: the external signal may have + // aborted while getClient() (or the token calculation) was pending. Bail + // before invoking the host so no request starts for a cancelled request. + if (externalAbortSignal?.aborted) { + cancellationTokenSource.cancel() + const abortError = new Error("Zoo Code : Request aborted") + 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.`, @@ -436,6 +446,13 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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") { From 4e72095d187d78c9b38028ac9a7556e5a2141417 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 14:18:19 +0800 Subject: [PATCH 4/9] fix(api): start vscode-lm request cancellation before client initialization --- src/api/providers/__tests__/vscode-lm.spec.ts | 49 +++++++- src/api/providers/vscode-lm.ts | 107 +++++++++++------- 2 files changed, 110 insertions(+), 46 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 9fd4ea3aa1..515177de43 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -513,7 +513,8 @@ describe("VsCodeLmHandler", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] // Gate model selection so the generator waits inside getClient(): the - // pre-abort check and listener attachment happened before this window. + // 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 @@ -524,6 +525,7 @@ describe("VsCodeLmHandler", () => { handler["client"] = null const controller = new AbortController() + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") const stream = handler.createMessage( systemPrompt, messages, @@ -531,18 +533,24 @@ describe("VsCodeLmHandler", () => { ) const firstChunk = stream.next() - // Release the client, then abort before the generator can reach sendRequest. + // 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. + // 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. + // 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 () => { @@ -1349,6 +1357,39 @@ describe("VsCodeLmHandler", () => { } }) + 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]) diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 748159e7a6..0f31897fd1 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -376,31 +376,15 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // Ensure clean state before starting a new request this.ensureCleanState() - // The VS Code LanguageModelChat API cannot carry an AbortSignal, so a + // 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 - if (externalAbortSignal?.aborted) { - const abortError = new Error("Zoo Code : Request aborted") - abortError.name = "AbortError" - throw abortError - } - - 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), - ] - // 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() const cancellationTokenSource = this.currentRequestCancellation @@ -415,16 +399,24 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) } - // Calculate input tokens before starting the stream - const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) - // Accumulate the text and count at the end of the stream to reduce token counting overhead. let accumulatedText: string = "" try { - // Re-check after client initialization: the external signal may have - // aborted while getClient() (or the token calculation) was pending. Bail - // before invoking the host so no request starts for a cancelled request. + // 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") @@ -432,6 +424,21 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan 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.`, @@ -522,6 +529,16 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan outputTokens: totalOutputTokens, } } catch (error: unknown) { + // 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) { // The host rejected because the request was cancelled: either the // bridged external signal aborted or the user cancelled the request @@ -647,24 +664,18 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { - // Fail fast if the caller already aborted before we even started: do not - // initialize or invoke the host request for a cancelled request. - if (options?.abortSignal?.aborted) { - const abortError = new Error("VSCode LM completion aborted") - abortError.name = "AbortError" - throw abortError - } - - const client = await this.getClient() - // 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. + // 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. + // 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) @@ -688,9 +699,21 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan tokenSource.token.isCancellationRequested === true || externalAbortSignal?.aborted === true try { - // Re-check before invoking the host: the signal may have aborted while the - // client was being initialized above. - if (externalAbortSignal?.aborted) { + // 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 From 8741154532b0c21453397f65c25b4ee11a29658a Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 14:56:22 +0800 Subject: [PATCH 5/9] fix(api): cancel vscode-lm request on premature generator closure + fix pre-abort test comment --- src/api/providers/__tests__/vscode-lm.spec.ts | 39 ++++++++++++++++++- src/api/providers/vscode-lm.ts | 6 +++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 515177de43..911aadb728 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -624,6 +624,40 @@ describe("VsCodeLmHandler", () => { 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 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[] = [ @@ -1279,9 +1313,10 @@ describe("VsCodeLmHandler", () => { 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: no cancellation token source is - // created and the sendRequest is never called. + // Fails fast before invoking the host: the pre-aborted signal cancels the + // request token and sendRequest is never called. expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() }) it("should reject with an AbortError when the signal aborts mid-flight", async () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 0f31897fd1..259e69aa70 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -575,6 +575,12 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan 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 From b5f12ebb2fd2a0b8247a0517fe45cfd05ddf9e79 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Sat, 5 Sep 2026 13:56:48 +0800 Subject: [PATCH 6/9] test(vscode-lm): assert abort/timeout edge paths for the mutation-diff gate Adds focused assertions for the remaining mutation-diff mutants on the abort-signal and timeoutMs paths: external-abort pre-checks and the in-stream abort re-throw, cancellation listener cleanup, timeout timer scheduling and clearing, and host CancellationError wrapping. --- src/api/providers/__tests__/vscode-lm.spec.ts | 376 ++++++++++++++++-- 1 file changed, 352 insertions(+), 24 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 911aadb728..d6a5a2c2a0 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -495,11 +495,8 @@ describe("VsCodeLmHandler", () => { const controller = new AbortController() controller.abort() - const stream = handler.createMessage( - systemPrompt, - messages, - makeCreateMessageMetadata({ abortSignal: controller.signal }), - ) + 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", ) @@ -526,11 +523,8 @@ describe("VsCodeLmHandler", () => { const controller = new AbortController() const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") - const stream = handler.createMessage( - systemPrompt, - messages, - makeCreateMessageMetadata({ abortSignal: controller.signal }), - ) + 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 @@ -573,11 +567,8 @@ describe("VsCodeLmHandler", () => { }) const controller = new AbortController() - const stream = handler.createMessage( - systemPrompt, - messages, - makeCreateMessageMetadata({ abortSignal: controller.signal }), - ) + 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)) @@ -609,11 +600,8 @@ describe("VsCodeLmHandler", () => { const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") - const stream = handler.createMessage( - systemPrompt, - messages, - makeCreateMessageMetadata({ abortSignal: controller.signal }), - ) + const meta = makeCreateMessageMetadata({ abortSignal: controller.signal }) + const stream = handler.createMessage(systemPrompt, messages, meta) for await (const _chunk of stream) { // drain } @@ -658,6 +646,254 @@ describe("VsCodeLmHandler", () => { 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[] = [ @@ -1313,8 +1549,7 @@ describe("VsCodeLmHandler", () => { 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 - // request token and sendRequest is never called. + // Fails fast before invoking the host: the pre-aborted signal cancels the token. expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) expect(tokenSourceInstance().cancel).toHaveBeenCalled() }) @@ -1344,11 +1579,20 @@ describe("VsCodeLmHandler", () => { const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) await vi.waitFor(() => expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1)) - // Abort while the stream is in flight, then let the mock stream finish. + // 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() - await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + // 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 () => { @@ -1452,6 +1696,90 @@ describe("VsCodeLmHandler", () => { 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]) From 2ee61f00c695ac91611fb77f912e6bc8c706e11c Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Sat, 5 Sep 2026 21:56:42 +0800 Subject: [PATCH 7/9] test(vscode-lm): assert exact abort error messages and listener cleanup Close the 18 surviving mutants of the PR #1300 mutation-diff gate on vscode-lm.ts: - assert the exact canonical abort error message at every abort path in both createMessage and completePrompt - rewrite the finally listener-detach guards to if (externalAbortSignal && handler) so the guard only runs when the signal exists, and direct the correlated LogicalOperator mutant - direct the StringLiteral mutants on the guard/bridge abort errors: the catch re-throws its own canonical abort error, so those literals are unobservable - add a test where the external signal aborts while the host cancellation token was not cancelled, exercising the right-hand operand of isAborted() - strengthen the listener-spy tests to assert removeEventListener is called with the same listener reference, and clear the timeout timer unconditionally --- src/api/providers/__tests__/vscode-lm.spec.ts | 107 ++++++++++++++++-- src/api/providers/vscode-lm.ts | 27 +++-- 2 files changed, 115 insertions(+), 19 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index d6a5a2c2a0..0b1b22733d 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -498,7 +498,10 @@ describe("VsCodeLmHandler", () => { 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) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request aborted", ) // No host request is started for an already-aborted signal. @@ -532,7 +535,12 @@ describe("VsCodeLmHandler", () => { releaseClient() controller.abort() - await expect(firstChunk).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request aborted", + ) // The abort landed after client initialization, so no host request started // and no input tokens were counted (the post-init re-check bails first). @@ -577,7 +585,12 @@ describe("VsCodeLmHandler", () => { // instead of yielding stale chunks. controller.abort() releaseStream() - await expect(firstChunk).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request aborted", + ) // The bridge relayed the abort to the request cancellation token. const tokenSource = tokenSourceInstance() @@ -609,7 +622,7 @@ describe("VsCodeLmHandler", () => { expect(addEventListenerSpy).toHaveBeenCalledTimes(1) expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) - expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addEventListenerSpy.mock.calls[0][1]) }) it("should cancel the request token when the consumer stops consuming early", async () => { @@ -681,7 +694,10 @@ describe("VsCodeLmHandler", () => { 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) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request aborted", ) // The pre-abort short-circuit runs before getClient(): the host is never contacted. @@ -720,7 +736,12 @@ describe("VsCodeLmHandler", () => { // 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") + await expect(firstChunk).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "Zoo Code : Request aborted", + ) // Bridge + post-init re-check + finally each cancel the request token. expect(tokenSourceInstance().cancel).toHaveBeenCalledTimes(3) @@ -1537,7 +1558,7 @@ describe("VsCodeLmHandler", () => { expect(addEventListenerSpy).toHaveBeenCalledTimes(1) expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) - expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addEventListenerSpy.mock.calls[0][1]) }) it("should reject with an AbortError when the signal is already aborted", async () => { @@ -1547,7 +1568,12 @@ describe("VsCodeLmHandler", () => { controller.abort() const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) - await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + 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) @@ -1595,6 +1621,47 @@ describe("VsCodeLmHandler", () => { ) }) + 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]) @@ -1629,7 +1696,10 @@ describe("VsCodeLmHandler", () => { // still makes the completion abort. releaseStream() await expect(promise).rejects.toSatisfy( - (error) => error instanceof Error && error.name === "AbortError", + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "VSCode LM completion aborted", ) } finally { vi.useRealTimers() @@ -1661,7 +1731,10 @@ describe("VsCodeLmHandler", () => { releaseClient() await expect(promise).rejects.toSatisfy( - (error) => error instanceof Error && error.name === "AbortError", + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "VSCode LM completion aborted", ) expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) } finally { @@ -1773,7 +1846,12 @@ describe("VsCodeLmHandler", () => { controller.abort() const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) - await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + 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) @@ -1807,7 +1885,12 @@ describe("VsCodeLmHandler", () => { handler["client"] = mockLanguageModelChat const promise = handler.completePrompt("Test prompt") - await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + await expect(promise).rejects.toSatisfy( + (error) => + error instanceof Error && + error.name === "AbortError" && + error.message === "VSCode LM completion aborted", + ) }) }) diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 259e69aa70..3b52b6c1f3 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -407,7 +407,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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 } @@ -419,7 +421,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // tokens or invoking the host so no work happens 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 } @@ -456,7 +460,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // A late abort while consuming must stop the stream instead of // yielding stale chunks. if (externalAbortSignal?.aborted) { + // 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 } @@ -572,8 +578,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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) + // 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 @@ -708,7 +715,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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 } @@ -720,7 +729,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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 } @@ -739,7 +750,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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 } @@ -762,13 +775,13 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } throw error } finally { - if (timeoutId !== undefined) { - clearTimeout(timeoutId) - } + // clearTimeout(undefined) is a spec no-op, so clear unconditionally. + clearTimeout(timeoutId) // { once: true } only detaches the listener when the signal actually // aborts, so remove it explicitly on every path. - if (onAbort !== undefined) { - externalAbortSignal?.removeEventListener("abort", onAbort) + // 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() } From 8a7d58041fc1ec61e0a3ea7b537e951800ab830b Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Sat, 5 Sep 2026 23:15:13 +0800 Subject: [PATCH 8/9] fix(vscode-lm): check request-local cancellation in createMessage and cover negative timeout --- src/api/providers/__tests__/vscode-lm.spec.ts | 127 ++++++++++++++++++ src/api/providers/vscode-lm.ts | 19 ++- 2 files changed, 144 insertions(+), 2 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 0b1b22733d..ca10795b1c 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -915,6 +915,93 @@ describe("VsCodeLmHandler", () => { 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 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() + // Token counting takes long enough that the caller's abort lands while it + // is in flight; the bridge cancels the request token during the count. + // @ts-ignore – access private method to drive the counting window + vi.spyOn(handler, "calculateTotalInputTokens").mockImplementation(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 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[] = [ @@ -1809,6 +1896,46 @@ describe("VsCodeLmHandler", () => { } }) + 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* () { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 3b52b6c1f3..4cf290dcd5 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -443,6 +443,19 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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.`, @@ -458,8 +471,10 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan // 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) { + // 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 From 1c753666b87ef267fb62951d423689f0e0966374 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Sun, 6 Sep 2026 16:45:50 +0800 Subject: [PATCH 9/9] fix(vscode-lm): re-check request-local token after client init Address CodeRabbit findings on the abort-signal series: - Re-check the request-local token immediately after getClient() so a request superseded while client initialization is pending aborts before token counting or host invocation. - Drive the counting-abort regression through the public path by mocking countTokens on the host chat mock, removing the @ts-ignore private-method suppression. - Add a gated-initialization supersession regression covering the new post-init guard. --- src/api/providers/__tests__/vscode-lm.spec.ts | 132 ++++++++++++++++-- src/api/providers/vscode-lm.ts | 9 +- 2 files changed, 123 insertions(+), 18 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index ca10795b1c..9e26a63754 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -16,6 +16,39 @@ vi.mock("vscode", () => { ) {} } + class MockLanguageModelToolResultPart { + type = "tool_result" + constructor( + public callId: string, + public content: unknown, + ) {} + } + + // 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". + class MockLanguageModelChatMessage { + role: string + content: unknown + constructor(role: string, content: unknown) { + this.role = role + this.content = content + } + static Assistant(content: string | unknown[]): MockLanguageModelChatMessage { + return new MockLanguageModelChatMessage( + "assistant", + Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], + ) + } + static User(content: string | unknown[]): MockLanguageModelChatMessage { + return new MockLanguageModelChatMessage( + "user", + Array.isArray(content) ? content : [new MockLanguageModelTextPart(content)], + ) + } + } + return { workspace: { getConfiguration: vi.fn(() => ({ @@ -46,18 +79,10 @@ 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, lm: { selectChatModels: vi.fn(), }, @@ -81,7 +106,9 @@ 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), } /** @@ -979,10 +1006,10 @@ describe("VsCodeLmHandler", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] const controller = new AbortController() - // Token counting takes long enough that the caller's abort lands while it - // is in flight; the bridge cancels the request token during the count. - // @ts-ignore – access private method to drive the counting window - vi.spyOn(handler, "calculateTotalInputTokens").mockImplementation(async () => { + // 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 }) @@ -1002,6 +1029,81 @@ describe("VsCodeLmHandler", () => { 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 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[] = [ diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 4cf290dcd5..02b053fa91 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -416,10 +416,13 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan 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 + // 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. - if (externalAbortSignal?.aborted) { + // 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")