diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 423f119f14..911aadb728 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -26,12 +26,17 @@ vi.mock("vscode", () => { })), }, CancellationTokenSource: vi.fn(function () { + // Faithful to the real API: cancel() marks the token as requested so + // consumers that read isCancellationRequested observe the cancellation. + const token = { + isCancellationRequested: false, + onCancellationRequested: vi.fn(), + } return { - token: { - isCancellationRequested: false, - onCancellationRequested: vi.fn(), - }, - cancel: vi.fn(), + token, + cancel: vi.fn(() => { + token.isCancellationRequested = true + }), dispose: vi.fn(), } }), @@ -65,6 +70,7 @@ import type { ApiHandlerOptions } from "../../../shared/api" import type { Anthropic } from "@anthropic-ai/sdk" import { openAiModelInfoSaneDefaults, vscodeLlmDefaultModelId, vscodeLlmModels } from "@roo-code/types" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" const mockLanguageModelChat = { @@ -78,6 +84,18 @@ const mockLanguageModelChat = { countTokens: vi.fn(), } +/** + * Returns the instance created by the n-th `new vscode.CancellationTokenSource()` + * call recorded by the module mock, for asserting on the cancellation-token lifecycle. + */ +function tokenSourceInstance(index = 0) { + const result = (vscode.CancellationTokenSource as Mock).mock.results[index] + if (result?.type !== "return") { + return undefined + } + return result.value +} + describe("VsCodeLmHandler", () => { let handler: VsCodeLmHandler const defaultOptions: ApiHandlerOptions = { @@ -470,6 +488,176 @@ 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 reject with an AbortError when the external signal aborts during client initialization", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + // Gate model selection so the generator waits inside getClient(): the + // cancellation token source and the abort-bridge listener are established + // before this window, and the finally cleanup covers the getClient() await. + let releaseClient: () => void = () => {} + const clientGate = new Promise((resolve) => { + releaseClient = resolve + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), + ) + handler["client"] = null + + const controller = new AbortController() + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const firstChunk = stream.next() + + // Release the client, then abort before the next microtask can reach + // countTokens/sendRequest. + releaseClient() + controller.abort() + + await expect(firstChunk).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + + // The abort landed after client initialization, so no host request started + // and no input tokens were counted (the post-init re-check bails first). + expect(handler["client"]).not.toBeNull() + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + expect(mockLanguageModelChat.countTokens).toHaveBeenCalledTimes(0) + + // The local token source was cancelled for the request that never started, + // and the abort-bridge listener was removed by the finally covering the + // initialization window. + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + }) + + it("should bridge a mid-flight external abort to the request cancellation token", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + + let releaseStream: () => void = () => {} + const streamGate = new Promise((resolve) => { + releaseStream = resolve + }) + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + await streamGate + yield new vscode.LanguageModelTextPart("Hello") + })(), + text: (async function* () { + await streamGate + yield "Hello" + })(), + }) + + const controller = new AbortController() + const 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. The late abort must stop the stream + // instead of yielding stale chunks. + controller.abort() + releaseStream() + await expect(firstChunk).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + + // The bridge relayed the abort to the request cancellation token. + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).toHaveBeenCalled() + }) + + it("should attach and detach the abort bridge listener around the request", async () => { + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "Hello" }] + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Hello") + })(), + text: (async function* () { + yield "Hello" + })(), + }) + + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const 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 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[] = [ @@ -1039,6 +1227,260 @@ describe("VsCodeLmHandler", () => { const promise = handler.completePrompt("Test prompt") await expect(promise).rejects.toThrow("VSCode LM completion error: Completion failed") }) + + it("should work without options (backward compatible)", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Completed text") + }) + + it("should bridge the abort signal to a fresh cancellation token and dispose it", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + const result = await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(result).toBe("Completed text") + const tokenSource = tokenSourceInstance() + expect(tokenSource.cancel).not.toHaveBeenCalled() + expect(tokenSource.dispose).toHaveBeenCalled() + }) + + it("should attach and detach the abort listener around the completion", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + + it("should reject with an AbortError when the signal is already aborted", async () => { + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + controller.abort() + + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + + // Fails fast before invoking the host: the pre-aborted signal cancels the + // 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 () => { + 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 completes, but the cancelled token + // still makes the completion abort. + releaseStream() + await expect(promise).rejects.toSatisfy( + (error) => error instanceof Error && error.name === "AbortError", + ) + } finally { + vi.useRealTimers() + } + }) + + it("should reject with an AbortError when the timeout fires during client initialization", async () => { + // Gate client initialization so the timeout timer (started before getClient()) + // can fire while getClient() is still pending. + let releaseClient: () => void = () => {} + const clientGate = new Promise((resolve) => { + releaseClient = resolve + }) + ;(vscode.lm.selectChatModels as Mock).mockImplementation(() => + clientGate.then(() => [{ ...mockLanguageModelChat }]), + ) + handler["client"] = null + + vi.useFakeTimers() + try { + const promise = handler.completePrompt("Test prompt", { timeoutMs: 5000 }) + + // Advance past the timeout while getClient() is still pending: the timer + // fires and cancels the request token. + await vi.advanceTimersByTimeAsync(5000) + expect(tokenSourceInstance().cancel).toHaveBeenCalled() + + // Release the client; the post-init re-check must abort before sendRequest. + releaseClient() + + await expect(promise).rejects.toSatisfy( + (error) => error instanceof Error && error.name === "AbortError", + ) + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(0) + } finally { + vi.useRealTimers() + } + }) + + it("should apply both the abort signal and timeoutMs", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockResolvedValueOnce({ + stream: (async function* () { + yield new vscode.LanguageModelTextPart("Completed text") + return + })(), + text: (async function* () { + yield "Completed text" + return + })(), + }) + + handler["client"] = mockLanguageModelChat + + const controller = new AbortController() + const result = await handler.completePrompt("Test prompt", { + abortSignal: controller.signal, + timeoutMs: 10000, + }) + + expect(result).toBe("Completed text") + expect(mockLanguageModelChat.sendRequest).toHaveBeenCalledTimes(1) + }) + + it("should wrap non-abort completion errors without an AbortError name", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new Error("LM error")) + handler["client"] = mockLanguageModelChat + + const promise = handler.completePrompt("Test prompt") + await expect(promise).rejects.toSatisfy((error) => { + return ( + error instanceof Error && + error.name === "Error" && + error.message === "VSCode LM completion error: LM error" + ) + }) + }) + + it("should reject with an AbortError when the host raises a CancellationError", async () => { + const mockModel = { ...mockLanguageModelChat } + ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) + + // The host rejects with a CancellationError while the token flag is not set + // (no abort signal, no timeout), so only the error type signals the abort. + mockLanguageModelChat.sendRequest.mockRejectedValueOnce(new vscode.CancellationError()) + handler["client"] = mockLanguageModelChat + + const promise = handler.completePrompt("Test prompt") + await expect(promise).rejects.toSatisfy((error) => error instanceof Error && error.name === "AbortError") + }) }) describe("cleanMessageContent / deepClean", () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index c657e6c0d6..259e69aa70 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -375,30 +375,70 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan ): ApiStream { // Ensure clean state before starting a new request this.ensureCleanState() - const client: vscode.LanguageModelChat = await this.getClient() - // Process messages - const cleanedMessages = messages.map((msg) => ({ - ...msg, - content: this.cleanMessageContent(msg.content), - })) - - // Convert Anthropic messages to VS Code LM messages - const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ - vscode.LanguageModelChatMessage.Assistant(systemPrompt), - ...convertToVsCodeLmMessages(cleanedMessages), - ] + // The VS Code LanguageModelChat API cannot carry an AbortSignal, so the + // request cancellation is established before client initialization and a + // pre-aborted external signal is reported immediately instead of being + // sent to the host. + const externalAbortSignal = metadata?.abortSignal - // Initialize cancellation token for the request + // Initialize cancellation token for the request before getClient() so the + // client-initialization await is covered by the abort bridge and the + // finally cleanup below. this.currentRequestCancellation = new vscode.CancellationTokenSource() - - // Calculate input tokens before starting the stream - const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) + const cancellationTokenSource = this.currentRequestCancellation + + // Bridge the caller's abort signal (e.g. a task abort) into the request's + // cancellation token: the VS Code LM API cannot carry an AbortSignal + // directly, so cancellation is signalled to the host through the token. + // The listener is kept in a named const and removed in the finally block + // because { once: true } only detaches it when the signal actually aborts. + let onExternalAbort: (() => void) | undefined + if (externalAbortSignal) { + onExternalAbort = () => cancellationTokenSource.cancel() + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } // Accumulate the text and count at the end of the stream to reduce token counting overhead. let accumulatedText: string = "" try { + // Fail fast if the caller already aborted before we even started: do not + // initialize or invoke the host request for a cancelled request. + if (externalAbortSignal?.aborted) { + cancellationTokenSource.cancel() + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } + + const client: vscode.LanguageModelChat = await this.getClient() + + // Re-check immediately after client initialization: the signal may have + // aborted while getClient() was pending. Bail before counting input + // tokens or invoking the host so no work happens for a cancelled request. + if (externalAbortSignal?.aborted) { + cancellationTokenSource.cancel() + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } + + // Process messages + const cleanedMessages = messages.map((msg) => ({ + ...msg, + content: this.cleanMessageContent(msg.content), + })) + + // Convert Anthropic messages to VS Code LM messages + const vsCodeLmMessages: vscode.LanguageModelChatMessage[] = [ + vscode.LanguageModelChatMessage.Assistant(systemPrompt), + ...convertToVsCodeLmMessages(cleanedMessages), + ] + + // Calculate input tokens before starting the stream + const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages) + // Create the response stream with required options const requestOptions: vscode.LanguageModelChatRequestOptions = { justification: `Zoo Code would like to use '${client.name}' from '${client.vendor}', Click 'Allow' to proceed.`, @@ -408,11 +448,18 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan const response: vscode.LanguageModelChatResponse = await client.sendRequest( vsCodeLmMessages, requestOptions, - this.currentRequestCancellation.token, + cancellationTokenSource.token, ) // Consume the stream and handle both text and tool call chunks for await (const chunk of response.stream) { + // A late abort while consuming must stop the stream instead of + // yielding stale chunks. + if (externalAbortSignal?.aborted) { + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } if (chunk instanceof vscode.LanguageModelTextPart) { // Validate text part value if (typeof chunk.value !== "string") { @@ -482,10 +529,23 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan outputTokens: totalOutputTokens, } } catch (error: unknown) { - this.ensureCleanState() + // When the external signal wins during client initialization (the signal + // was already aborted while getClient() was pending or rejected), surface a + // standard abort error instead of wrapping the getClient() failure as a + // generic stream error. + if (externalAbortSignal?.aborted) { + const abortError = new Error("Zoo Code : Request aborted") + abortError.name = "AbortError" + throw abortError + } if (error instanceof vscode.CancellationError) { - throw new Error("Zoo Code : Request cancelled by user") + // The host rejected because the request was cancelled: either the + // bridged external signal aborted or the user cancelled the request + // in VS Code. Both are aborts, so surface a standard abort error. + const abortError = new Error("Zoo Code : Request cancelled by user") + abortError.name = "AbortError" + throw abortError } if (error instanceof Error) { @@ -508,6 +568,27 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan console.error("Zoo Code : Unknown stream error:", errorMessage) throw new Error(`Zoo Code : Response stream error: ${errorMessage}`) } + } finally { + // Detach the abort bridge listener on every path (success, error, and + // early consumer break); { once: true } alone would leak it when the + // request completes without the signal ever aborting. + if (onExternalAbort !== undefined) { + externalAbortSignal?.removeEventListener("abort", onExternalAbort) + } + // Cancel before disposing: VS Code's CancellationTokenSource.dispose() + // frees resources without cancelling the token, so a premature consumer + // closure (break/return) would otherwise leave the host request running and + // consuming model quota. Cancel is idempotent, so this is a no-op on paths + // where the token was already cancelled (aborted or timed-out request). + cancellationTokenSource.cancel() + // Dispose the request-local source. Clear the shared field only if it still + // points at this request's source: a newer request may have replaced it, and + // disposing the shared field here would cancel and dispose the newer request's + // token. + cancellationTokenSource.dispose() + if (this.currentRequestCancellation === cancellationTokenSource) { + this.currentRequestCancellation = null + } } } @@ -589,12 +670,65 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // The VS Code LanguageModelChat API cannot carry an AbortSignal: sendRequest + // only accepts a CancellationToken. Bridge the external signal and timeout + // into a request-local CancellationTokenSource instead. Cancellation is + // established before client initialization so the getClient() await is covered + // by the timeout, the abort bridge, and the finally cleanup below. + const tokenSource = new vscode.CancellationTokenSource() + const externalAbortSignal = options?.abortSignal + + // Apply the timeout only when it is a positive value: cancelling at once for a + // zero/negative timeout would abort every such request immediately. Starting + // the timer before getClient() means a timeout that fires during a slow client + // lookup still cancels the request before any sendRequest call. + let timeoutId: ReturnType | undefined + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + timeoutId = setTimeout(() => tokenSource.cancel(), options.timeoutMs) + } + + // Bridge the external abort signal: a pre-aborted signal cancels the token + // immediately, otherwise a one-shot listener relays the abort to the host. + let onAbort: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + tokenSource.cancel() + } else { + onAbort = () => tokenSource.cancel() + externalAbortSignal.addEventListener("abort", onAbort, { once: true }) + } + } + + // The request counts as aborted when the host-side token was cancelled + // (timeout or bridged signal) or when the external signal itself aborted. + const isAborted = () => + tokenSource.token.isCancellationRequested === true || externalAbortSignal?.aborted === true + try { + // Fail fast if the caller already aborted before we even started: do not + // initialize or invoke the host request for a cancelled request. + if (isAborted()) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + const client = await this.getClient() + + // Re-check after client initialization: the signal may have aborted (or the + // timeout fired) while getClient() was pending. Bail before invoking the host + // so a timeout that fired during a slow client lookup can never lead to a + // sendRequest call. + if (isAborted()) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + const response = await client.sendRequest( [vscode.LanguageModelChatMessage.User(prompt)], {}, - new vscode.CancellationTokenSource().token, + tokenSource.token, ) let result = "" for await (const chunk of response.stream) { @@ -602,12 +736,41 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan result += chunk.value } } + + // Guard against a quiet completion after the request was aborted. + if (isAborted()) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + return result } catch (error) { + // Report an aborted request (external signal, timeout, or host + // cancellation) as a standard AbortError instead of a generic + // completion error. A host CancellationError is treated as a cancellation + // even if the token flag was not observed (e.g. the host cancelled the + // request through the token without the flag being set). + if (isAborted() || error instanceof vscode.CancellationError) { + const abortError = new Error("VSCode LM completion aborted") + abortError.name = "AbortError" + throw abortError + } + if (error instanceof Error) { throw new Error(`VSCode LM completion error: ${error.message}`) } throw error + } finally { + if (timeoutId !== undefined) { + clearTimeout(timeoutId) + } + // { once: true } only detaches the listener when the signal actually + // aborts, so remove it explicitly on every path. + if (onAbort !== undefined) { + externalAbortSignal?.removeEventListener("abort", onAbort) + } + tokenSource.dispose() } } }