diff --git a/src/api/providers/__tests__/complete-prompt-options.spec.ts b/src/api/providers/__tests__/complete-prompt-options.spec.ts new file mode 100644 index 0000000000..f9925cd119 --- /dev/null +++ b/src/api/providers/__tests__/complete-prompt-options.spec.ts @@ -0,0 +1,29 @@ +import { describe, it, expect } from "vitest" + +import type { CompletePromptOptions } from "../../index" + +describe("CompletePromptOptions", () => { + it("should allow abortSignal property", () => { + const controller = new AbortController() + const options: CompletePromptOptions = { abortSignal: controller.signal } + expect(options.abortSignal).toBe(controller.signal) + }) + + it("should allow timeoutMs property", () => { + const options: CompletePromptOptions = { timeoutMs: 5000 } + expect(options.timeoutMs).toBe(5000) + }) + + it("should allow both abortSignal and timeoutMs together", () => { + const controller = new AbortController() + const options: CompletePromptOptions = { abortSignal: controller.signal, timeoutMs: 10000 } + expect(options.abortSignal).toBe(controller.signal) + expect(options.timeoutMs).toBe(10000) + }) + + it("should allow empty options object", () => { + const options: CompletePromptOptions = {} + expect(options.abortSignal).toBeUndefined() + expect(options.timeoutMs).toBeUndefined() + }) +}) diff --git a/src/api/providers/__tests__/opencode-go.spec.ts b/src/api/providers/__tests__/opencode-go.spec.ts index 98209ae325..a669dfb6a9 100644 --- a/src/api/providers/__tests__/opencode-go.spec.ts +++ b/src/api/providers/__tests__/opencode-go.spec.ts @@ -9,8 +9,12 @@ vitest.mock("vscode", () => ({ }, })) -import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import { + Anthropic, + APIConnectionTimeoutError as AnthropicTimeoutError, + APIUserAbortError as AnthropicAbortError, +} from "@anthropic-ai/sdk" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { opencodeGoDefaultModelId, opencodeGoModels, isOpencodeGoAnthropicFormatModel } from "@roo-code/types" @@ -19,6 +23,7 @@ import { getModels } from "../fetchers/modelCache" import { ApiHandlerOptions } from "../../../shared/api" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" +import { makeCreateMessageMetadata } from "../../../test-utils/api" vitest.mock("openai") vitest.mock("delay", () => ({ @@ -48,15 +53,22 @@ const mockAnthropicCreate = vitest.fn() } }) -vitest.mock("@anthropic-ai/sdk", () => ({ - Anthropic: vitest.fn(function () { - return { - messages: { - create: mockAnthropicCreate, - }, - } - }), -})) +// The real SDK error classes are re-exported alongside the mocked client so +// tests can emulate the SDK's abort/timeout rejections and the provider's +// instanceof checks resolve against the same class identity. +vitest.mock("@anthropic-ai/sdk", async () => { + const actual = await vi.importActual("@anthropic-ai/sdk") + return { + ...actual, + Anthropic: vitest.fn(function () { + return { + messages: { + create: mockAnthropicCreate, + }, + } + }), + } +}) describe("OpencodeGoHandler", () => { const mockOptions: ApiHandlerOptions = { @@ -188,6 +200,7 @@ describe("OpencodeGoHandler", () => { max_completion_tokens: 40_960, temperature: expect.any(Number), }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -202,6 +215,7 @@ describe("OpencodeGoHandler", () => { model: "glm-5.1", reasoning_effort: "medium", }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -352,10 +366,124 @@ describe("OpencodeGoHandler", () => { await collectStream(handler.createMessage("sys", messages)) - expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ max_completion_tokens: 999 })) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ max_completion_tokens: 999 }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), + ) + }) + + it("rethrows non-abort errors from the OpenAI stream unchanged", async () => { + // A mid-stream failure that is not an abort (e.g. a connection + // reset) must propagate unchanged — the catch only normalizes + // aborts to a DOM-standard AbortError. + const streamError = new Error("connection reset") + mockCreate.mockImplementation(async () => + (async function* () { + yield { choices: [{ delta: { content: "partial" }, index: 0 }], index: 0 } + throw streamError + })(), + ) + + const handler = new OpencodeGoHandler(mockOptions) + + const error = await collectStream(handler.createMessage("sys", [{ role: "user", content: "hi" }])).then( + () => undefined, + (e: unknown) => e, + ) + + expect(error).toBe(streamError) }) }) + describe("createMessage abort signal bridging", () => { + it("rejects with an AbortError when the external signal is already aborted", async () => { + mockCreate.mockImplementation(async () => { + throw new DOMException("The operation was aborted.", "AbortError") + }) + + const handler = new OpencodeGoHandler(mockOptions) + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "sys", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("aborts the in-flight request when the external signal fires mid-stream", async () => { + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + return (async function* () { + yield { + choices: [{ delta: { content: "partial" }, index: 0 }], + index: 0, + } + await new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("The operation was aborted.", "AbortError")) + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + })() + }) + + const handler = new OpencodeGoHandler(mockOptions) + const controller = new AbortController() + + const consumed = collectStream( + handler.createMessage( + "sys", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ), + ) + + // Let the request start and the first chunk be yielded before aborting. + await new Promise((resolve) => setTimeout(resolve, 25)) + controller.abort() + + await expect(consumed).rejects.toMatchObject({ name: "AbortError" }) + expect(capturedSignal?.aborted).toBe(true) + }) + + it("detaches the bridged abort listener when the request completes normally", async () => { + // The listener is added with { once: true }, so it only detaches on + // abort. A task-scoped signal spanning many requests must not + // accumulate a listener per request: assert explicit removal after a + // normal (non-aborted) completion. + mockCreate.mockImplementation(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" }, index: 0 }], + index: 0, + }, + { + choices: [{ delta: {}, index: 0 }], + index: 0, + usage: { prompt_tokens: 2, completion_tokens: 3 }, + }, + ]), + ) + + const handler = new OpencodeGoHandler(mockOptions) + const controller = new AbortController() + const removeListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const stream = handler.createMessage( + "sys", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const chunks = await collectStream(stream) + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + expect(removeListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(controller.signal.aborted).toBe(false) + }) + }) describe("completePrompt", () => { it("returns the message content for a non-streaming completion", async () => { mockCreate.mockResolvedValue({ choices: [{ message: { content: "the answer" } }] }) @@ -369,6 +497,7 @@ describe("OpencodeGoHandler", () => { max_completion_tokens: 40_960, reasoning_effort: "medium", }), + {}, ) }) @@ -394,7 +523,7 @@ describe("OpencodeGoHandler", () => { mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) const handler = new OpencodeGoHandler({ ...mockOptions, includeMaxTokens: true, modelMaxTokens: 4321 }) await handler.completePrompt("ping") - expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ max_completion_tokens: 4321 })) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ max_completion_tokens: 4321 }), {}) }) }) @@ -456,6 +585,7 @@ describe("OpencodeGoHandler", () => { stream: true, system: expect.arrayContaining([expect.objectContaining({ type: "text", text: "sys" })]), }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) // The OpenAI chat completions endpoint must NOT be used for this model. expect(mockCreate).not.toHaveBeenCalled() @@ -537,6 +667,7 @@ describe("OpencodeGoHandler", () => { // so the model default is used. max_tokens: 65_536, }), + undefined, ) expect(mockCreate).not.toHaveBeenCalled() }) @@ -552,7 +683,7 @@ describe("OpencodeGoHandler", () => { modelMaxTokens: 2048, }) await handler.completePrompt("ping") - expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ max_tokens: 2048 })) + expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ max_tokens: 2048 }), undefined) }) it("completePrompt rethrows non-Error values unchanged from the Anthropic path", async () => { @@ -567,6 +698,237 @@ describe("OpencodeGoHandler", () => { expect(await handler.completePrompt("ping")).toBe("") }) + it("completePrompt passes abort signal through to Anthropic client", async () => { + mockAnthropicCreate.mockResolvedValue({ content: [{ type: "text", text: "response" }] }) + const controller = new AbortController() + const handler = new OpencodeGoHandler(anthropicOptions) + await handler.completePrompt("ping", { abortSignal: controller.signal }) + expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + signal: controller.signal, + }) + }) + + it("completePrompt passes both signal and timeoutMs through to Anthropic client", async () => { + mockAnthropicCreate.mockResolvedValue({ content: [{ type: "text", text: "response" }] }) + const controller = new AbortController() + const handler = new OpencodeGoHandler(anthropicOptions) + await handler.completePrompt("ping", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + signal: controller.signal, + timeout: 10000, + }) + }) + + it("completePrompt passes only timeoutMs when no signal is provided", async () => { + mockAnthropicCreate.mockResolvedValue({ content: [{ type: "text", text: "response" }] }) + const handler = new OpencodeGoHandler(anthropicOptions) + await handler.completePrompt("ping", { timeoutMs: 5000 }) + expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeout: 5000, + }) + }) + + it("completePrompt omits the timeout option when timeoutMs is 0 (Anthropic path)", async () => { + // The SDK treats timeout: 0 as an immediate abort, so the "disabled" + // value must never be forwarded — assert the absence of the option. + mockAnthropicCreate.mockResolvedValue({ content: [{ type: "text", text: "response" }] }) + const handler = new OpencodeGoHandler(anthropicOptions) + await handler.completePrompt("ping", { timeoutMs: 0 }) + const call = mockAnthropicCreate.mock.calls[mockAnthropicCreate.mock.calls.length - 1] + const requestOptions = call[1] as { timeout?: number } | undefined + // When no option is forwarded the provider omits the SDK options + // argument entirely, so absence means: undefined arg OR an arg + // without a timeout key. + expect(Object.keys(requestOptions ?? {})).not.toContain("timeout") + }) + + it("completePrompt preserves abort identity when the caller aborts (Anthropic path)", async () => { + // Emulate the Anthropic SDK: an aborted request signal rejects with + // APIUserAbortError ("Request was aborted." — the trailing period would + // fail task-level abort detection, so the provider must normalize it). + mockAnthropicCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + throw new AnthropicAbortError() + } + throw new Error("boom") + }) + const handler = new OpencodeGoHandler(anthropicOptions) + const controller = new AbortController() + controller.abort() + + const error = await handler.completePrompt("ping", { abortSignal: controller.signal }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("completePrompt surfaces request timeouts as an AbortError (Anthropic path)", async () => { + // Emulate the Anthropic SDK: when the request timeout fires, the SDK + // surfaces APIConnectionTimeoutError ("Request timed out.") once retries + // are exhausted — verified against @anthropic-ai/sdk against a hung server. + mockAnthropicCreate.mockImplementation(async (_params: unknown, options: { timeout?: number }) => { + await new Promise((resolve) => setTimeout(resolve, options?.timeout ?? 50)) + throw new AnthropicTimeoutError() + }) + const handler = new OpencodeGoHandler(anthropicOptions) + + const error = await handler.completePrompt("ping", { timeoutMs: 50 }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("completePrompt works without options (backward compatible, Anthropic path)", async () => { + mockAnthropicCreate.mockResolvedValue({ content: [{ type: "text", text: "response" }] }) + const handler = new OpencodeGoHandler(anthropicOptions) + const result = await handler.completePrompt("ping") + expect(result).toBe("response") + expect(mockAnthropicCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, + ) + }) + + describe("completePrompt (OpenAI path)", () => { + const openaiOptions: ApiHandlerOptions = { + opencodeGoApiKey: "test-key", + apiModelId: "glm-5.1", // OpenAI-format model + } + + beforeEach(() => { + vitest.clearAllMocks() + }) + + it("completePrompt returns text for OpenAI path", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + + const handler = new OpencodeGoHandler(openaiOptions) + expect(await handler.completePrompt("ping")).toBe("response") + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String), stream: false }), + {}, // empty object when no options + ) + }) + + it("completePrompt passes abort signal through to OpenAI client", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const controller = new AbortController() + const handler = new OpencodeGoHandler(openaiOptions) + + await handler.completePrompt("ping", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String), stream: false }), + { signal: controller.signal }, + ) + }) + + it("completePrompt passes both signal and timeoutMs through to OpenAI client", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const controller = new AbortController() + const handler = new OpencodeGoHandler(openaiOptions) + + await handler.completePrompt("ping", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String), stream: false }), + { signal: controller.signal, timeout: 10000 }, + ) + }) + + it("completePrompt passes only timeoutMs when no signal is provided", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const handler = new OpencodeGoHandler(openaiOptions) + + await handler.completePrompt("ping", { timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String), stream: false }), + { timeout: 5000 }, + ) + }) + + it("completePrompt omits the timeout option when timeoutMs is 0 (OpenAI path)", async () => { + // The OpenAI SDK treats timeout: 0 as an immediate abort, so the + // "disabled" value must never be forwarded — assert the absence of + // the option. + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const handler = new OpencodeGoHandler(openaiOptions) + await handler.completePrompt("ping", { timeoutMs: 0 }) + const call = mockCreate.mock.calls[mockCreate.mock.calls.length - 1] + const requestOptions = call[1] as { timeout?: number } | undefined + expect(requestOptions).not.toHaveProperty("timeout") + }) + + it("completePrompt preserves abort identity when the caller aborts (OpenAI path)", async () => { + // Emulate the OpenAI SDK: an aborted request signal rejects with + // APIUserAbortError ("Request was aborted." — the trailing period would + // fail task-level abort detection, so the provider must normalize it). + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + throw new APIUserAbortError() + } + throw new Error("boom") + }) + const handler = new OpencodeGoHandler(openaiOptions) + const controller = new AbortController() + controller.abort() + + const error = await handler.completePrompt("ping", { abortSignal: controller.signal }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("completePrompt surfaces request timeouts as an AbortError (OpenAI path)", async () => { + // Emulate the OpenAI SDK: when the request timeout fires, the SDK + // surfaces APIConnectionTimeoutError ("Request timed out.") once retries + // are exhausted — verified against openai v5.23.2 against a hung server. + mockCreate.mockImplementation(async (_params: unknown, options: { timeout?: number }) => { + await new Promise((resolve) => setTimeout(resolve, options?.timeout ?? 50)) + throw new APIConnectionTimeoutError() + }) + const handler = new OpencodeGoHandler(openaiOptions) + + const error = await handler.completePrompt("ping", { timeoutMs: 50 }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("completePrompt works without options (backward compatible, OpenAI path)", async () => { + mockCreate.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const handler = new OpencodeGoHandler(openaiOptions) + + const result = await handler.completePrompt("ping") + expect(result).toBe("response") + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String), stream: false }), + {}, // empty object when no options + ) + }) + }) it("omits tools and tool_choice from the Anthropic request when no tools are provided", async () => { const handler = new OpencodeGoHandler(anthropicOptions) const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hi" }] @@ -712,7 +1074,10 @@ describe("OpencodeGoHandler", () => { await collectStream(handler.createMessage("sys", messages)) - expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ max_tokens: 8192 })) + expect(mockAnthropicCreate).toHaveBeenCalledWith( + expect.objectContaining({ max_tokens: 8192 }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), + ) }) it("falls back to the model max_tokens when includeMaxTokens is on but modelMaxTokens is unset", async () => { @@ -722,7 +1087,10 @@ describe("OpencodeGoHandler", () => { await collectStream(handler.createMessage("sys", messages)) // qwen3.7-max maxTokens (65_536) clamped to 20% of 1M context => 65_536. - expect(mockAnthropicCreate).toHaveBeenCalledWith(expect.objectContaining({ max_tokens: 65_536 })) + expect(mockAnthropicCreate).toHaveBeenCalledWith( + expect.objectContaining({ max_tokens: 65_536 }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), + ) }) it("accumulates output tokens across message_delta events into the final cost", async () => { @@ -789,6 +1157,21 @@ describe("OpencodeGoHandler", () => { await collectStream(handler.createMessage("sys", messages)) }).rejects.toThrow("Opencode Go completion error: rate limited") }) + + it("preserves abort identity for aborted Anthropic requests from createMessage", async () => { + // A cancelled /v1/messages request (the SDK rejects with + // APIUserAbortError) must surface as a DOM-standard AbortError, not + // the wrapped "completion error" reserved for other failures. + mockAnthropicCreate.mockRejectedValue(new AnthropicAbortError()) + const handler = new OpencodeGoHandler(anthropicOptions) + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hi" }] + await expect(async () => { + await collectStream(handler.createMessage("sys", messages)) + }).rejects.toMatchObject({ + name: "AbortError", + message: "The Opencode Go request was aborted", + }) + }) }) describe("isOpencodeGoAnthropicFormatModel", () => { diff --git a/src/api/providers/__tests__/unbound.spec.ts b/src/api/providers/__tests__/unbound.spec.ts index 9b45713386..2269fc77ff 100644 --- a/src/api/providers/__tests__/unbound.spec.ts +++ b/src/api/providers/__tests__/unbound.spec.ts @@ -1,18 +1,27 @@ import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { UnboundHandler } from "../unbound" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" +import { makeCreateMessageMetadata } from "../../../test-utils/api" -vi.mock("openai", () => { - const createMock = vi.fn() +// Single hoisted mock shared by the `openai` factory and every test so tests +// can configure the SDK `create` call without untyped access casts. +const sharedMockCreate = vi.hoisted(() => vi.fn()) + +// The real SDK error classes are re-exported alongside the mocked client so +// tests can emulate the SDK's abort/timeout rejections (APIUserAbortError, +// APIConnectionTimeoutError) and the provider's instanceof checks resolve. +vi.mock("openai", async () => { + const actual = await vi.importActual("openai") return { + ...actual, default: vi.fn(function () { return { chat: { completions: { - create: createMock, + create: sharedMockCreate, }, }, } @@ -175,7 +184,80 @@ describe("UnboundHandler", () => { mode: "architect", }, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), + ) + }) + + it("wraps non-abort pre-stream failures via handleOpenAIError", async () => { + // A non-abort rejection from create() (e.g. an upstream 500) must be + // routed through handleOpenAIError, not the AbortError normalization + // path: assert the wrapped identity and the preserved message. + sharedMockCreate.mockRejectedValue(new Error("upstream 500")) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const stream = handler.createMessage("system", [{ role: "user", content: "hi" }], { + taskId: "t", + tools: [], + }) + + const error = await collectStream(stream).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toBeInstanceOf(Error) + expect((error as Error).message).toBe("Unbound completion error: upstream 500") + expect((error as Error).name).not.toBe("AbortError") + }) + + it("emits tool_call_partial chunks for native tool calls in the stream", async () => { + // Native tool calls arrive on delta.tool_calls and must be re-emitted + // as raw tool_call_partial chunks for NativeToolCallParser to assemble. + sharedMockCreate.mockResolvedValue( + asyncStreamFrom([ + { + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: '{"city": "NYC"}' }, + }, + ], + }, + }, + ], + }, + { choices: [{ delta: { content: "done" } }], usage: { prompt_tokens: 1, completion_tokens: 1 } }, + ]), + ) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const chunks = await collectStream( + handler.createMessage("system", [{ role: "user", content: "hi" }], { + taskId: "t", + tools: [], + }), ) + + expect(chunks).toContainEqual({ + type: "tool_call_partial", + index: 0, + id: "call_1", + name: "get_weather", + arguments: '{"city": "NYC"}', + }) + expect(chunks).toContainEqual({ type: "text", text: "done" }) }) it("completePrompt returns the response text", async () => { @@ -195,6 +277,225 @@ describe("UnboundHandler", () => { expect.objectContaining({ messages: [{ role: "system", content: "Write a haiku" }], }), + {}, + ) + }) + + it("completePrompt should pass abort signal through to client", async () => { + const controller = new AbortController() + sharedMockCreate.mockResolvedValue({ + choices: [{ message: { content: "completed text" } }], + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + await handler.completePrompt("Write a haiku", { abortSignal: controller.signal }) + expect(sharedMockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal }), + ) + }) + + it("completePrompt should pass timeout through to client", async () => { + sharedMockCreate.mockResolvedValue({ + choices: [{ message: { content: "completed text" } }], + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + await handler.completePrompt("Write a haiku", { timeoutMs: 5000 }) + expect(sharedMockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), ) }) + + it("completePrompt should omit the timeout option when timeoutMs is 0", async () => { + // The OpenAI SDK treats timeout: 0 as an immediate abort, so the + // "disabled" value must never be forwarded — assert the absence of + // the option (a forwarded timeout: 0 would fail this assertion). + sharedMockCreate.mockResolvedValue({ + choices: [{ message: { content: "completed text" } }], + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + await handler.completePrompt("Write a haiku", { timeoutMs: 0 }) + const call = sharedMockCreate.mock.calls[sharedMockCreate.mock.calls.length - 1] + const requestOptions = call[1] as { timeout?: number } | undefined + expect(requestOptions).not.toHaveProperty("timeout") + }) + + it("completePrompt should preserve abort identity when the caller aborts", async () => { + // Emulate the OpenAI SDK: an aborted request signal rejects with + // APIUserAbortError ("Request was aborted." — the trailing period would + // fail task-level abort detection, so the provider must normalize it). + sharedMockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + throw new APIUserAbortError() + } + throw new Error("boom") + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + const controller = new AbortController() + controller.abort() + + const error = await handler.completePrompt("Write a haiku", { abortSignal: controller.signal }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("completePrompt should surface request timeouts as an AbortError", async () => { + // Emulate the OpenAI SDK: when the request timeout fires, the SDK + // surfaces APIConnectionTimeoutError ("Request timed out.") once retries + // are exhausted — verified against openai v5.23.2 against a hung server. + sharedMockCreate.mockImplementation(async (_params: unknown, options: { timeout?: number }) => { + await new Promise((resolve) => setTimeout(resolve, options?.timeout ?? 50)) + throw new APIConnectionTimeoutError() + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const error = await handler.completePrompt("Write a haiku", { timeoutMs: 50 }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + it("completePrompt should work without options (backward compatible)", async () => { + sharedMockCreate.mockResolvedValue({ + choices: [{ message: { content: "completed text" } }], + }) + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const result = await handler.completePrompt("Write a haiku") + expect(result).toBe("completed text") + }) + + describe("createMessage abort signal bridging", () => { + it("rejects the request with an AbortError when the external signal is already aborted", async () => { + let requestError: unknown + sharedMockCreate.mockImplementation(async () => { + // The real SDK rejects with an AbortError when its request signal is aborted. + requestError = new DOMException("The operation was aborted.", "AbortError") + throw requestError + }) + + const controller = new AbortController() + controller.abort() + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const stream = handler.createMessage( + "system", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + // The bridge surfaces a DOM-standard AbortError (series standard) + // instead of the wrapped completion error. + await expect(collectStream(stream)).rejects.toMatchObject({ + name: "AbortError", + message: "The Unbound request was aborted", + }) + expect(requestError).toMatchObject({ name: "AbortError" }) + }) + + it("aborts the in-flight request when the external signal fires mid-stream", async () => { + let capturedSignal: AbortSignal | undefined + sharedMockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + await new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("The operation was aborted.", "AbortError")) + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + })() + }) + + const controller = new AbortController() + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const consumed = collectStream( + handler.createMessage( + "system", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ), + ) + + // Let the request start and the first chunk be yielded before aborting. + await new Promise((resolve) => setTimeout(resolve, 25)) + controller.abort() + + await expect(consumed).rejects.toMatchObject({ name: "AbortError" }) + expect(capturedSignal?.aborted).toBe(true) + }) + + it("detaches the bridged abort listener when the request completes normally", async () => { + // The listener is added with { once: true }, so it only detaches on + // abort. A task-scoped signal spanning many requests must not + // accumulate a listener per request: assert explicit removal after a + // normal (non-aborted) completion. + sharedMockCreate.mockImplementation(async () => + asyncStreamFrom([ + { choices: [{ delta: { content: "ok" } }] }, + { choices: [{ delta: {} }], usage: { prompt_tokens: 1, completion_tokens: 1 } }, + ]), + ) + + const controller = new AbortController() + const removeListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const handler = new UnboundHandler({ + unboundApiKey: "test-key", + unboundModelId: "openai/gpt-4o", + }) + + const chunks = await collectStream( + handler.createMessage( + "system", + [{ role: "user", content: "hi" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ), + ) + + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + expect(removeListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(controller.signal.aborted).toBe(false) + }) + }) }) diff --git a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts index ffad3fa0d1..e111c9f524 100644 --- a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts +++ b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts @@ -10,10 +10,10 @@ vitest.mock("vscode", () => ({ })) import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { VercelAiGatewayHandler } from "../vercel-ai-gateway" -import { makeApiHandlerOptions } from "../../../test-utils/api" +import { makeApiHandlerOptions, makeCreateMessageMetadata } from "../../../test-utils/api" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" import { vercelAiGatewayDefaultModelId, VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types" @@ -287,6 +287,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ temperature: customTemp, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -302,6 +303,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -321,6 +323,7 @@ describe("VercelAiGatewayHandler", () => { temperature: undefined, max_completion_tokens: 128000, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -393,6 +396,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ max_completion_tokens: 64000, // max tokens for sonnet 4 }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -468,6 +472,7 @@ describe("VercelAiGatewayHandler", () => { }), ]), }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -485,6 +490,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ tool_choice: "auto", }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -502,6 +508,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ parallel_tool_calls: true, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -519,6 +526,7 @@ describe("VercelAiGatewayHandler", () => { tools: expect.any(Array), parallel_tool_calls: true, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -615,6 +623,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ stream_options: { include_usage: true }, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) }) @@ -653,6 +662,7 @@ describe("VercelAiGatewayHandler", () => { temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE, max_completion_tokens: 64000, }), + undefined, ) }) @@ -671,6 +681,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ temperature: customTemp, }), + undefined, ) }) @@ -703,10 +714,239 @@ describe("VercelAiGatewayHandler", () => { const result = await handler.completePrompt("Test") expect(result).toBe("") }) + + it("should pass abort signal through to client", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const controller = new AbortController() + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: "response" }, + finish_reason: "stop", + index: 0, + }, + ], + })) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal }), + ) + }) + + it("should pass timeout through to client", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: "response" }, + finish_reason: "stop", + index: 0, + }, + ], + })) + + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), + ) + }) + + it("should omit the timeout option when timeoutMs is 0", async () => { + // The OpenAI SDK treats timeout: 0 as an immediate abort, so the + // "disabled" value must never be forwarded — assert the absence of + // the option (a forwarded timeout: 0 would fail this assertion). + const handler = new VercelAiGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: "response" }, + finish_reason: "stop", + index: 0, + }, + ], + })) + + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + const call = mockCreate.mock.calls[mockCreate.mock.calls.length - 1] + const requestOptions = call[1] as { timeout?: number } | undefined + // When no option is forwarded the provider omits the SDK options + // argument entirely, so absence means: undefined arg OR an arg + // without a timeout key. + expect(Object.keys(requestOptions ?? {})).not.toContain("timeout") + }) + + it("should preserve abort identity when the caller aborts", async () => { + // Emulate the OpenAI SDK: an aborted request signal rejects with + // APIUserAbortError ("Request was aborted." — the trailing period would + // fail task-level abort detection, so the provider must normalize it). + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + throw new APIUserAbortError() + } + throw new Error("boom") + }) + + const handler = new VercelAiGatewayHandler(mockOptions) + const controller = new AbortController() + controller.abort() + + const error = await handler.completePrompt("test prompt", { abortSignal: controller.signal }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("should surface request timeouts as an AbortError", async () => { + // Emulate the OpenAI SDK: when the request timeout fires, the SDK + // surfaces APIConnectionTimeoutError ("Request timed out.") once retries + // are exhausted — verified against openai v5.23.2 against a hung server. + mockCreate.mockImplementation(async (_params: unknown, options: { timeout?: number }) => { + await new Promise((resolve) => setTimeout(resolve, options?.timeout ?? 50)) + throw new APIConnectionTimeoutError() + }) + + const handler = new VercelAiGatewayHandler(mockOptions) + + const error = await handler.completePrompt("test prompt", { timeoutMs: 50 }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + it("should work without options (backward compatible)", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: "response" }, + finish_reason: "stop", + index: 0, + }, + ], + })) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + }) + + describe("createMessage abort signal bridging", () => { + it("rejects with an AbortError when the external signal is already aborted", async () => { + mockCreate.mockImplementation(async () => { + throw new DOMException("The operation was aborted.", "AbortError") + }) + + const handler = new VercelAiGatewayHandler(mockOptions) + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "test prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("aborts the in-flight request when the external signal fires mid-stream", async () => { + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + return (async function* () { + yield { + choices: [{ delta: { content: "partial" }, index: 0 }], + index: 0, + } + await new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("The operation was aborted.", "AbortError")) + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + })() + }) + + const handler = new VercelAiGatewayHandler(mockOptions) + const controller = new AbortController() + + const consumed = collectStream( + handler.createMessage( + "test prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ), + ) + + // Let the request start and the first chunk be yielded before aborting. + await new Promise((resolve) => setTimeout(resolve, 25)) + controller.abort() + + await expect(consumed).rejects.toMatchObject({ name: "AbortError" }) + expect(capturedSignal?.aborted).toBe(true) + }) + + it("detaches the bridged abort listener when the request completes normally", async () => { + // The listener is added with { once: true }, so it only detaches on + // abort. A task-scoped signal spanning many requests must not + // accumulate a listener per request: assert explicit removal after a + // normal (non-aborted) completion. + mockCreate.mockImplementation(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" }, index: 0 }], + index: 0, + }, + { + choices: [{ delta: {}, index: 0 }], + index: 0, + usage: { prompt_tokens: 2, completion_tokens: 3 }, + }, + ]), + ) + + const handler = new VercelAiGatewayHandler(mockOptions) + const controller = new AbortController() + const removeListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const stream = handler.createMessage( + "test prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const chunks = await collectStream(stream) + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + expect(removeListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(controller.signal.aborted).toBe(false) + }) }) describe("temperature support", () => { it("applies temperature for supported models", async () => { + // Pin the response: a later describe's mock implementation may have + // left the shared mock in a state this test does not expect. + mockCreate.mockResolvedValueOnce({ + choices: [ + { + message: { role: "assistant", content: "Test completion response" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 8, + completion_tokens: 4, + total_tokens: 12, + }, + }) + const handler = new VercelAiGatewayHandler( makeApiHandlerOptions({ ...mockOptions, @@ -721,6 +961,7 @@ describe("VercelAiGatewayHandler", () => { expect.objectContaining({ temperature: 0.9, }), + undefined, ) }) }) diff --git a/src/api/providers/__tests__/zoo-gateway.spec.ts b/src/api/providers/__tests__/zoo-gateway.spec.ts index 66131d7cb1..7a1eafdb43 100644 --- a/src/api/providers/__tests__/zoo-gateway.spec.ts +++ b/src/api/providers/__tests__/zoo-gateway.spec.ts @@ -26,7 +26,7 @@ vitest.mock("../../../i18n", () => ({ t: (key: string) => key, })) -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { zooGatewayDefaultModelId, ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types" @@ -36,6 +36,7 @@ import { Package } from "../../../shared/package" import { clearZooCodeToken } from "../../../services/zoo-code-auth" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" +import { makeCreateMessageMetadata } from "../../../test-utils/api" vitest.mock("openai") vitest.mock("delay", () => ({ @@ -289,6 +290,7 @@ describe("ZooGatewayHandler", () => { "X-Zoo-Task-ID": "task-123", "X-Zoo-Mode": "code", }, + signal: expect.any(AbortSignal), }), ) }) @@ -434,6 +436,7 @@ describe("ZooGatewayHandler", () => { temperature: ZOO_GATEWAY_DEFAULT_TEMPERATURE, max_completion_tokens: 64000, }), + {}, ) }) @@ -456,8 +459,192 @@ describe("ZooGatewayHandler", () => { await expect(handler.completePrompt("Test")).resolves.toBe("") }) + + it("should pass abort signal through to client", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const controller = new AbortController() + mockCreate.mockImplementation(async () => ({ + choices: [{ message: { role: "assistant", content: "response" } }], + })) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal }), + ) + }) + + it("should pass timeout through to client", async () => { + const handler = new ZooGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [{ message: { role: "assistant", content: "response" } }], + })) + + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), + ) + }) + + it("should omit the timeout option when timeoutMs is 0", async () => { + // The OpenAI SDK treats timeout: 0 as an immediate abort, so the + // "disabled" value must never be forwarded — assert the absence of + // the option (a forwarded timeout: 0 would fail this assertion). + const handler = new ZooGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [{ message: { role: "assistant", content: "response" } }], + })) + + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + const call = mockCreate.mock.calls[mockCreate.mock.calls.length - 1] + const requestOptions = call[1] as { timeout?: number } | undefined + expect(requestOptions).not.toHaveProperty("timeout") + }) + + it("should preserve abort identity when the caller aborts", async () => { + // Emulate the OpenAI SDK: an aborted request signal rejects with + // APIUserAbortError ("Request was aborted." — the trailing period would + // fail task-level abort detection, so the provider must normalize it). + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + throw new APIUserAbortError() + } + throw new Error("boom") + }) + + const handler = new ZooGatewayHandler(mockOptions) + const controller = new AbortController() + controller.abort() + + const error = await handler.completePrompt("test prompt", { abortSignal: controller.signal }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + + it("should surface request timeouts as an AbortError", async () => { + // Emulate the OpenAI SDK: when the request timeout fires, the SDK + // surfaces APIConnectionTimeoutError ("Request timed out.") once retries + // are exhausted — verified against openai v5.23.2 against a hung server. + mockCreate.mockImplementation(async (_params: unknown, options: { timeout?: number }) => { + await new Promise((resolve) => setTimeout(resolve, options?.timeout ?? 50)) + throw new APIConnectionTimeoutError() + }) + + const handler = new ZooGatewayHandler(mockOptions) + + const error = await handler.completePrompt("test prompt", { timeoutMs: 50 }).then( + () => undefined, + (e: unknown) => e, + ) + expect(error).toMatchObject({ name: "AbortError" }) + expect((error as Error).message.endsWith("aborted")).toBe(true) + expect((error as Error).message).not.toContain("completion error") + }) + it("should work without options (backward compatible)", async () => { + const handler = new ZooGatewayHandler(mockOptions) + mockCreate.mockImplementation(async () => ({ + choices: [{ message: { role: "assistant", content: "response" } }], + })) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) }) + describe("createMessage abort signal bridging", () => { + it("rejects with an AbortError when the external signal is already aborted", async () => { + mockCreate.mockImplementation(async () => { + throw new DOMException("The operation was aborted.", "AbortError") + }) + + const handler = new ZooGatewayHandler(mockOptions) + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("aborts the in-flight request when the external signal fires mid-stream", async () => { + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementation(async (_params: unknown, options: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + return (async function* () { + yield { + choices: [{ delta: { content: "partial" }, index: 0 }], + index: 0, + } + await new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("The operation was aborted.", "AbortError")) + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + })() + }) + + const handler = new ZooGatewayHandler(mockOptions) + const controller = new AbortController() + + const consumed = collectStream( + handler.createMessage( + "prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ), + ) + + // Let the request start and the first chunk be yielded before aborting. + await new Promise((resolve) => setTimeout(resolve, 25)) + controller.abort() + + await expect(consumed).rejects.toMatchObject({ name: "AbortError" }) + expect(capturedSignal?.aborted).toBe(true) + }) + + it("detaches the bridged abort listener when the request completes normally", async () => { + // The listener is added with { once: true }, so it only detaches on + // abort. A task-scoped signal spanning many requests must not + // accumulate a listener per request: assert explicit removal after a + // normal (non-aborted) completion. + mockCreate.mockImplementation(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" }, index: 0 }], + index: 0, + }, + { + choices: [{ delta: {}, index: 0 }], + index: 0, + usage: { prompt_tokens: 2, completion_tokens: 3 }, + }, + ]), + ) + + const handler = new ZooGatewayHandler(mockOptions) + const controller = new AbortController() + const removeListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + + const stream = handler.createMessage( + "prompt", + [{ role: "user", content: "hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const chunks = await collectStream(stream) + expect(chunks).toContainEqual({ type: "text", text: "ok" }) + expect(removeListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(controller.signal.aborted).toBe(false) + }) + }) describe("classifyGatewayApiError", () => { it("returns sign_in on 401", () => { expect(classifyGatewayApiError(makeApiError(401))).toEqual({ kind: "sign_in" }) diff --git a/src/api/providers/opencode-go.ts b/src/api/providers/opencode-go.ts index 9456ac8fdb..cd3556d72e 100644 --- a/src/api/providers/opencode-go.ts +++ b/src/api/providers/opencode-go.ts @@ -1,6 +1,10 @@ -import { Anthropic } from "@anthropic-ai/sdk" +import { + Anthropic, + APIConnectionTimeoutError as AnthropicTimeoutError, + APIUserAbortError as AnthropicAbortError, +} from "@anthropic-ai/sdk" import { CacheControlEphemeral } from "@anthropic-ai/sdk/resources" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { type ModelInfo, @@ -28,6 +32,7 @@ import { convertOpenAIToolsToAnthropic, convertOpenAIToolChoiceToAnthropic, } from "../../core/prompts/tools/native-tools/converters" +import { createAbortError } from "./utils/abort-signal" /** * API handler for the Opencode "Go" subscription plan. @@ -165,8 +170,40 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio ): ApiStream { const { id: modelId, info, format, temperature, reasoningEffort, maxTokens } = await this.resolveModel() + // Per-request controller so an external abort signal (e.g. task + // cancellation) can interrupt the in-flight streaming request. + // Bridge it to our controller using the Bedrock pattern: + // - pre-aborted guard: check if already aborted before adding listener + // - { once: true }: remove listener after first abort to avoid leaks + // The listener is stored so it can be detached when the request ends: + // { once: true } only removes it on abort, so a task-scoped signal + // would otherwise accumulate one listener per request. + const controller = new AbortController() + const externalAbortSignal = metadata?.abortSignal + const abortListener = () => controller.abort() + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + controller.abort() + } else { + externalAbortSignal.addEventListener("abort", abortListener, { once: true }) + } + } + if (format === "anthropic") { - yield* this.streamAnthropicMessage(modelId, info, temperature, maxTokens, systemPrompt, messages, metadata) + try { + yield* this.streamAnthropicMessage( + modelId, + info, + temperature, + maxTokens, + systemPrompt, + messages, + controller.signal, + metadata, + ) + } finally { + externalAbortSignal?.removeEventListener("abort", abortListener) + } return } @@ -198,42 +235,54 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio }), } - const completion = await this.client.chat.completions.create(body) + try { + const completion = await this.client.chat.completions.create(body, { signal: controller.signal }) - for await (const chunk of completion) { - const delta = chunk.choices[0]?.delta + for await (const chunk of completion) { + const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { type: "text", text: delta.content } - } + if (delta?.content) { + yield { type: "text", text: delta.content } + } - // Several Go-plan models (GLM, DeepSeek) stream reasoning via this field. - const reasoningText = extractReasoningFromDelta(delta) - if (reasoningText) { - yield { type: "reasoning", text: reasoningText } - } + // Several Go-plan models (GLM, DeepSeek) stream reasoning via this field. + const reasoningText = extractReasoningFromDelta(delta) + if (reasoningText) { + yield { type: "reasoning", text: reasoningText } + } - // Emit raw tool call chunks - NativeToolCallParser handles state management. - if (delta?.tool_calls) { - for (const toolCall of delta.tool_calls) { - yield { - type: "tool_call_partial", - index: toolCall.index, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + // Emit raw tool call chunks - NativeToolCallParser handles state management. + if (delta?.tool_calls) { + for (const toolCall of delta.tool_calls) { + yield { + type: "tool_call_partial", + index: toolCall.index, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } } } - } - if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, - cacheReadTokens: chunk.usage.prompt_tokens_details?.cached_tokens || undefined, + if (chunk.usage) { + yield { + type: "usage", + inputTokens: chunk.usage.prompt_tokens || 0, + outputTokens: chunk.usage.completion_tokens || 0, + cacheReadTokens: chunk.usage.prompt_tokens_details?.cached_tokens || undefined, + } } } + } catch (error) { + // Preserve abort identity (series standard): surface a cancelled + // request as a DOM-standard AbortError rather than leaking the + // raw SDK abort error. + if (controller.signal.aborted) { + throw createAbortError("Opencode Go") + } + throw error + } finally { + externalAbortSignal?.removeEventListener("abort", abortListener) } } @@ -256,6 +305,7 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio maxTokens: number | undefined, systemPrompt: string, messages: Anthropic.Messages.MessageParam[], + abortSignal: AbortSignal, metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { const cacheControl: CacheControlEphemeral = { type: "ephemeral" } @@ -306,8 +356,18 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio // errors propagate unchanged, matching the OpenAI streaming path. let stream try { - stream = await this.anthropicClient.messages.create(requestParams) + stream = await this.anthropicClient.messages.create(requestParams, { signal: abortSignal }) } catch (error) { + // Preserve abort identity (series standard): a cancelled request + // must surface as a DOM-standard AbortError, not a wrapped + // completion error. + if ( + abortSignal.aborted || + error instanceof AnthropicAbortError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Opencode Go") + } if (error instanceof Error) { throw new Error(`Opencode Go completion error: ${error.message}`) } @@ -491,24 +551,53 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio if (format === "anthropic") { try { - const message = await this.anthropicClient.messages.create({ - model: modelId, - // Honour the same includeMaxTokens/modelMaxTokens override - // logic as the streaming path so non-streaming completions - // respect the user's max-output slider instead of always - // falling back to the model default. - max_tokens: - this.options.includeMaxTokens === true - ? this.options.modelMaxTokens || maxTokens || 16_384 - : (maxTokens ?? 16_384), - temperature: this.supportsTemperature(modelId) ? (temperature ?? 1.0) : undefined, - messages: [{ role: "user", content: prompt }], - stream: false, - }) + // Build request options with abortSignal and/or timeout handling. + // timeoutMs <= 0 means "no explicit timeout": omit the SDK timeout + // option entirely — the SDKs treat timeout: 0 as an immediate + // abort, which would cancel the request right away. + const requestOptions: Anthropic.RequestOptions = {} + if (options?.abortSignal) { + requestOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + requestOptions.timeout = options.timeoutMs + } + + const message = await this.anthropicClient.messages.create( + { + model: modelId, + // Honour the same includeMaxTokens/modelMaxTokens override + // logic as the streaming path so non-streaming completions + // respect the user's max-output slider instead of always + // falling back to the model default. + max_tokens: + this.options.includeMaxTokens === true + ? this.options.modelMaxTokens || maxTokens || 16_384 + : (maxTokens ?? 16_384), + temperature: this.supportsTemperature(modelId) ? (temperature ?? 1.0) : undefined, + messages: [{ role: "user", content: prompt }], + stream: false, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = message.content.find(({ type }) => type === "text") return content?.type === "text" ? content.text : "" } catch (error) { + // Preserve abort identity (series standard): caller-initiated + // cancellations and request timeouts must surface as a + // DOM-standard AbortError, not a wrapped completion error. The + // Anthropic SDK reports both with messages ending in a period + // ("Request was aborted.", "Request timed out."), which would not + // match task-level abort detection (message ending in "aborted"). + if ( + options?.abortSignal?.aborted || + error instanceof AnthropicAbortError || + error instanceof AnthropicTimeoutError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Opencode Go") + } if (error instanceof Error) { throw new Error(`Opencode Go completion error: ${error.message}`) } @@ -535,9 +624,35 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio reasoningEffort as OpenAI.Chat.ChatCompletionCreateParams["reasoning_effort"] } - const response = await this.client.chat.completions.create(requestOptions) + // Build request options with abortSignal and/or timeout for OpenAI path. + // timeoutMs <= 0 means "no explicit timeout": omit the SDK timeout + // option entirely — the OpenAI SDK treats timeout: 0 as an immediate + // abort, which would cancel the request right away. + const createOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + createOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + createOptions.timeout = options.timeoutMs + } + + const response = await this.client.chat.completions.create(requestOptions, createOptions) return response.choices[0]?.message.content || "" } catch (error) { + // Preserve abort identity (series standard): caller-initiated + // cancellations and request timeouts must surface as a + // DOM-standard AbortError, not a wrapped completion error. The + // OpenAI SDK reports both with messages ending in a period + // ("Request was aborted.", "Request timed out."), which would not + // match task-level abort detection (message ending in "aborted"). + if ( + options?.abortSignal?.aborted || + error instanceof APIUserAbortError || + error instanceof APIConnectionTimeoutError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Opencode Go") + } if (error instanceof Error) { throw new Error(`Opencode Go completion error: ${error.message}`) } diff --git a/src/api/providers/unbound.ts b/src/api/providers/unbound.ts index 0848e0804b..ebe8eb3fb9 100644 --- a/src/api/providers/unbound.ts +++ b/src/api/providers/unbound.ts @@ -1,5 +1,5 @@ import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { type ModelInfo, @@ -23,6 +23,7 @@ import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" +import { createAbortError } from "./utils/abort-signal" import { extractReasoningFromDelta } from "./utils/extract-reasoning" // Unbound usage includes extra fields for Anthropic cache tokens. @@ -158,46 +159,79 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand tool_choice: metadata?.tool_choice, } - let stream - try { - stream = await this.client.chat.completions.create(completionParams) - } catch (error) { - throw handleOpenAIError(error, this.providerName) + // Per-request controller so an external abort signal (e.g. task + // cancellation) can interrupt the in-flight streaming request. + // Bridge it to our controller using the Bedrock pattern: + // - pre-aborted guard: check if already aborted before adding listener + // - { once: true }: remove listener after first abort to avoid leaks + // The listener is stored so it can be detached when the request ends: + // { once: true } only removes it on abort, so a task-scoped signal + // would otherwise accumulate one listener per request. + const controller = new AbortController() + const externalAbortSignal = metadata?.abortSignal + const abortListener = () => controller.abort() + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + controller.abort() + } else { + externalAbortSignal.addEventListener("abort", abortListener, { once: true }) + } } - let lastUsage: any = undefined - - for await (const chunk of stream) { - const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { type: "text", text: delta.content } + try { + let stream + try { + stream = await this.client.chat.completions.create(completionParams, { signal: controller.signal }) + } catch (error) { + // Preserve abort identity (series standard): a cancelled request + // must surface as a DOM-standard AbortError, not a wrapped + // completion error. + if ( + controller.signal.aborted || + error instanceof APIUserAbortError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Unbound") + } + throw handleOpenAIError(error, this.providerName) } + let lastUsage: any = undefined - const reasoningText = extractReasoningFromDelta(delta) - if (reasoningText) { - yield { type: "reasoning", text: reasoningText } - } + for await (const chunk of stream) { + const delta = chunk.choices[0]?.delta + + if (delta?.content) { + yield { type: "text", text: delta.content } + } - // Handle native tool calls - if (delta && "tool_calls" in delta && Array.isArray(delta.tool_calls)) { - for (const toolCall of delta.tool_calls) { - yield { - type: "tool_call_partial", - index: toolCall.index, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + const reasoningText = extractReasoningFromDelta(delta) + if (reasoningText) { + yield { type: "reasoning", text: reasoningText } + } + + // Handle native tool calls + if (delta && "tool_calls" in delta && Array.isArray(delta.tool_calls)) { + for (const toolCall of delta.tool_calls) { + yield { + type: "tool_call_partial", + index: toolCall.index, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } } } - } - if (chunk.usage) { - lastUsage = chunk.usage + if (chunk.usage) { + lastUsage = chunk.usage + } } - } - if (lastUsage) { - yield this.processUsageMetrics(lastUsage, info) + if (lastUsage) { + yield this.processUsageMetrics(lastUsage, info) + } + } finally { + externalAbortSignal?.removeEventListener("abort", abortListener) } } @@ -212,11 +246,36 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand messages: openAiMessages, temperature: temperature, } + // Build request options with abortSignal and/or timeout. + // timeoutMs <= 0 means "no explicit timeout": omit the SDK timeout + // option entirely — the OpenAI SDK treats timeout: 0 as an immediate + // abort, which would cancel the request right away. + const createOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + createOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + createOptions.timeout = options.timeoutMs + } let response: OpenAI.Chat.ChatCompletion try { - response = await this.client.chat.completions.create(completionParams) + response = await this.client.chat.completions.create(completionParams, createOptions) } catch (error) { + // Preserve abort identity (series standard): caller-initiated + // cancellations and request timeouts must surface as a + // DOM-standard AbortError, not a wrapped completion error. The + // OpenAI SDK reports both with messages ending in a period + // ("Request was aborted.", "Request timed out."), which would not + // match task-level abort detection (message ending in "aborted"). + if ( + options?.abortSignal?.aborted || + error instanceof APIUserAbortError || + error instanceof APIConnectionTimeoutError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Unbound") + } throw handleOpenAIError(error, this.providerName) } return response.choices[0]?.message.content || "" diff --git a/src/api/providers/utils/__tests__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index ebc7edf3d3..aba72c181f 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -1,4 +1,10 @@ -import { mergeAbortSignalAndTimeout, mergeAbortSignals } from "../abort-signal" +import { + createAbortError, + isRequestAborted, + mergeAbortSignalAndTimeout, + mergeAbortSignals, + throwIfAborted, +} from "../abort-signal" describe("abort-signal utilities", () => { describe("mergeAbortSignalAndTimeout", () => { @@ -99,4 +105,84 @@ describe("abort-signal utilities", () => { expect(result.aborted).toBe(true) }) }) + + describe("throwIfAborted", () => { + it("does not throw when signal is undefined", () => { + expect(() => throwIfAborted()).not.toThrow() + }) + + it("does not throw when signal is not aborted", () => { + const controller = new AbortController() + + expect(() => throwIfAborted(controller.signal)).not.toThrow() + }) + + it("throws an AbortError when signal is already aborted", () => { + const controller = new AbortController() + controller.abort() + + let caught: unknown + try { + throwIfAborted(controller.signal) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + }) + }) + + describe("isRequestAborted", () => { + it("returns true when the caller signal is aborted", () => { + const controller = new AbortController() + controller.abort() + + expect(isRequestAborted(new Error("boom"), controller.signal)).toBe(true) + expect(isRequestAborted(undefined, controller.signal)).toBe(true) + }) + + it("returns true for a native AbortError or the OpenAI SDK APIUserAbortError", () => { + const native = new Error("This operation was aborted") + native.name = "AbortError" + expect(isRequestAborted(native)).toBe(true) + + const sdk = new Error("whatever") + sdk.name = "APIUserAbortError" + expect(isRequestAborted(sdk)).toBe(true) + }) + + it("matches the OpenAI SDK abort message exactly, not as a substring", () => { + expect(isRequestAborted(new Error("Request was aborted."))).toBe(true) + expect(isRequestAborted(new Error("Request was aborted"))).toBe(false) + expect(isRequestAborted(new Error("Request was aborted. Please retry"))).toBe(false) + }) + + it("returns false for unrelated errors, nullish errors, and live signals", () => { + expect(isRequestAborted(new Error("the abort failed"))).toBe(false) + expect(isRequestAborted(undefined)).toBe(false) + expect(isRequestAborted(null)).toBe(false) + + const controller = new AbortController() + expect(isRequestAborted(new Error("boom"), controller.signal)).toBe(false) + }) + }) + + describe("createAbortError", () => { + it("builds an error satisfying the Task.ts abort contract", () => { + const error = createAbortError("LM Studio") + + expect(error).toBeInstanceOf(Error) + expect(error.name).toBe("AbortError") + expect(error.message).toBe("The LM Studio request was aborted") + }) + + it("interpolates the provider name", () => { + expect(createAbortError("Qwen Code").message).toBe("The Qwen Code request was aborted") + }) + + it("returns a fresh error on each call", () => { + expect(createAbortError("X")).not.toBe(createAbortError("X")) + }) + }) }) diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index 73e0356f7b..26f57c3e9a 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -35,3 +35,61 @@ export function mergeAbortSignals(primarySignal: AbortSignal, secondarySignal?: return AbortSignal.any([primarySignal, secondarySignal]) } + +/** + * Throw an AbortError if the given signal is already aborted. + * + * Use as a fast-fail guard at the top of request-building code paths so + * callers receive a consistent `name === "AbortError"` when the operation + * was cancelled before it started, without building or issuing the request. + */ +export function throwIfAborted(signal?: AbortSignal): void { + if (!signal?.aborted) { + return + } + + const abortError = new Error("This operation was aborted") + abortError.name = "AbortError" + throw abortError +} + +/** + * Request options this series passes to the OpenAI SDK call. The SDK's + * `RequestOptions` declares `signal` as `AbortSignal | null | undefined`, + * which does not satisfy the builder's base constraint, so the builder is + * typed with only the options this series sets. The built config is still + * assignable to the SDK's `RequestOptions`. + */ +export type OpenAiRequestOptions = { + signal?: AbortSignal +} + +/** + * Whether a failure indicates an aborted request: the caller's signal fired, + * the SDK raised a native abort error, or the error carries the OpenAI SDK + * abort error message (exactly "Request was aborted."). The message check + * is an exact match on purpose: a substring match would misclassify + * unrelated errors that merely mention aborting. + */ +export function isRequestAborted(error: unknown, signal?: AbortSignal): boolean { + const candidate = error as { name?: string; message?: string } + return ( + Boolean(signal?.aborted) || + candidate?.name === "AbortError" || + candidate?.name === "APIUserAbortError" || + candidate?.message === "Request was aborted." + ) +} + +/** + * Fresh error satisfying the Task.ts abort contract: `name === + * "AbortError"` and a message ending in "aborted" (no trailing period). The + * OpenAI SDK's own abort error does not satisfy this contract (name "Error", + * message "Request was aborted."), so raw SDK abort errors must be + * normalized instead of rethrown. + */ +export function createAbortError(providerName: string): Error { + const abortError = new Error(`The ${providerName} request was aborted`) + abortError.name = "AbortError" + return abortError +} diff --git a/src/api/providers/vercel-ai-gateway.ts b/src/api/providers/vercel-ai-gateway.ts index bf434e5a00..4139ce9561 100644 --- a/src/api/providers/vercel-ai-gateway.ts +++ b/src/api/providers/vercel-ai-gateway.ts @@ -1,5 +1,5 @@ import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { vercelAiGatewayDefaultModelId, @@ -17,6 +17,7 @@ import { addCacheBreakpoints } from "../transform/caching/vercel-ai-gateway" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" +import { createAbortError } from "./utils/abort-signal" // Extend OpenAI's CompletionUsage to include Vercel AI Gateway specific fields interface VercelAiGatewayUsage extends OpenAI.CompletionUsage { @@ -69,52 +70,83 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp parallel_tool_calls: metadata?.parallelToolCalls ?? true, } - const completion = await this.client.chat.completions.create(body) - - for await (const chunk of completion) { - // Vercel AI Gateway reports mid-stream failures as an in-band error chunk - // rather than throwing, so surface it instead of returning an empty response. - if ("error" in chunk && chunk.error) { - const raw = chunk.error as { message?: unknown } - const message = - typeof raw.message === "string" && raw.message.length > 0 - ? raw.message - : "Vercel AI Gateway stream error" - throw new Error(message) + // Per-request controller so an external abort signal (e.g. task + // cancellation) can interrupt the in-flight streaming request. + // Bridge it to our controller using the Bedrock pattern: + // - pre-aborted guard: check if already aborted before adding listener + // - { once: true }: remove listener after first abort to avoid leaks + // The listener is stored so it can be detached when the request ends: + // { once: true } only removes it on abort, so a task-scoped signal + // would otherwise accumulate one listener per request. + const controller = new AbortController() + const externalAbortSignal = metadata?.abortSignal + const abortListener = () => controller.abort() + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + controller.abort() + } else { + externalAbortSignal.addEventListener("abort", abortListener, { once: true }) } + } - const delta = chunk.choices[0]?.delta - if (delta?.content) { - yield { - type: "text", - text: delta.content, + try { + const completion = await this.client.chat.completions.create(body, { signal: controller.signal }) + + for await (const chunk of completion) { + // Vercel AI Gateway reports mid-stream failures as an in-band error chunk + // rather than throwing, so surface it instead of returning an empty response. + if ("error" in chunk && chunk.error) { + const raw = chunk.error as { message?: unknown } + const message = + typeof raw.message === "string" && raw.message.length > 0 + ? raw.message + : "Vercel AI Gateway stream error" + throw new Error(message) } - } - // Emit raw tool call chunks - NativeToolCallParser handles state management - if (delta?.tool_calls) { - for (const toolCall of delta.tool_calls) { + const delta = chunk.choices[0]?.delta + if (delta?.content) { yield { - type: "tool_call_partial", - index: toolCall.index, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + type: "text", + text: delta.content, } } - } - if (chunk.usage) { - const usage = chunk.usage as VercelAiGatewayUsage - yield { - type: "usage", - inputTokens: usage.prompt_tokens || 0, - outputTokens: usage.completion_tokens || 0, - cacheWriteTokens: usage.cache_creation_input_tokens || undefined, - cacheReadTokens: usage.prompt_tokens_details?.cached_tokens || undefined, - totalCost: usage.cost ?? 0, + // Emit raw tool call chunks - NativeToolCallParser handles state management + if (delta?.tool_calls) { + for (const toolCall of delta.tool_calls) { + yield { + type: "tool_call_partial", + index: toolCall.index, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } + } } + + if (chunk.usage) { + const usage = chunk.usage as VercelAiGatewayUsage + yield { + type: "usage", + inputTokens: usage.prompt_tokens || 0, + outputTokens: usage.completion_tokens || 0, + cacheWriteTokens: usage.cache_creation_input_tokens || undefined, + cacheReadTokens: usage.prompt_tokens_details?.cached_tokens || undefined, + totalCost: usage.cost ?? 0, + } + } + } + } catch (error) { + // Preserve abort identity (series standard): surface a cancelled + // request as a DOM-standard AbortError rather than leaking the + // raw SDK abort error. + if (controller.signal.aborted) { + throw createAbortError("Vercel AI Gateway") } + throw error + } finally { + externalAbortSignal?.removeEventListener("abort", abortListener) } } @@ -133,10 +165,38 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp } requestOptions.max_completion_tokens = info.maxTokens + // Build request options with abortSignal and/or timeout. + // timeoutMs <= 0 means "no explicit timeout": omit the SDK timeout + // option entirely — the OpenAI SDK treats timeout: 0 as an immediate + // abort, which would cancel the request right away. + const createOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + createOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + createOptions.timeout = options.timeoutMs + } - const response = await this.client.chat.completions.create(requestOptions) + const response = await this.client.chat.completions.create( + requestOptions, + Object.keys(createOptions).length > 0 ? createOptions : undefined, + ) return response.choices[0]?.message.content || "" } catch (error) { + // Preserve abort identity (series standard): caller-initiated + // cancellations and request timeouts must surface as a + // DOM-standard AbortError, not a wrapped completion error. The + // OpenAI SDK reports both with messages ending in a period + // ("Request was aborted.", "Request timed out."), which would not + // match task-level abort detection (message ending in "aborted"). + if ( + options?.abortSignal?.aborted || + error instanceof APIUserAbortError || + error instanceof APIConnectionTimeoutError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Vercel AI Gateway") + } if (error instanceof Error) { throw new Error(`Vercel AI Gateway completion error: ${error.message}`) } diff --git a/src/api/providers/zoo-gateway.ts b/src/api/providers/zoo-gateway.ts index 4ff059df61..3f96b69de8 100644 --- a/src/api/providers/zoo-gateway.ts +++ b/src/api/providers/zoo-gateway.ts @@ -1,6 +1,6 @@ import * as vscode from "vscode" import { Anthropic } from "@anthropic-ai/sdk" -import OpenAI from "openai" +import OpenAI, { APIConnectionTimeoutError, APIUserAbortError } from "openai" import { zooGatewayDefaultModelId, @@ -22,6 +22,7 @@ import { addCacheBreakpoints } from "../transform/caching/vercel-ai-gateway" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { NOT_PROVIDED } from "./constants" import { RouterProvider } from "./router-provider" +import { createAbortError } from "./utils/abort-signal" function getApiErrorStatus(error: unknown): number | undefined { if (typeof error === "object" && error !== null && "status" in error) { @@ -219,9 +220,29 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio parallel_tool_calls: metadata?.parallelToolCalls ?? true, } + // Per-request controller so an external abort signal (e.g. task + // cancellation) can interrupt the in-flight streaming request. + // Bridge it to our controller using the Bedrock pattern: + // - pre-aborted guard: check if already aborted before adding listener + // - { once: true }: remove listener after first abort to avoid leaks + // The listener is stored so it can be detached when the request ends: + // { once: true } only removes it on abort, so a task-scoped signal + // would otherwise accumulate one listener per request. + const controller = new AbortController() + const externalAbortSignal = metadata?.abortSignal + const abortListener = () => controller.abort() + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + controller.abort() + } else { + externalAbortSignal.addEventListener("abort", abortListener, { once: true }) + } + } + try { const completion = await this.client.chat.completions.create(body, { headers: requestHeaders, + signal: controller.signal, }) for await (const chunk of completion) { @@ -266,6 +287,12 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio } } } catch (error) { + // Preserve abort identity (series standard): surface a cancelled + // request as a DOM-standard AbortError before the gateway error + // surfacing/telemetry path. + if (controller.signal.aborted) { + throw createAbortError("Zoo Gateway") + } try { await surfaceGatewayApiError(error) } catch (surfaceError) { @@ -275,6 +302,8 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio ) } throw error + } finally { + externalAbortSignal?.removeEventListener("abort", abortListener) } } @@ -295,10 +324,35 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio } requestOptions.max_completion_tokens = info.maxTokens + // Build request options with abortSignal and/or timeout. + // timeoutMs <= 0 means "no explicit timeout": omit the SDK timeout + // option entirely — the OpenAI SDK treats timeout: 0 as an immediate + // abort, which would cancel the request right away. + const createOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + createOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + createOptions.timeout = options.timeoutMs + } - const response = await this.client.chat.completions.create(requestOptions) + const response = await this.client.chat.completions.create(requestOptions, createOptions) return response.choices[0]?.message.content || "" } catch (error) { + // Preserve abort identity (series standard): caller-initiated + // cancellations and request timeouts must surface as a + // DOM-standard AbortError, not a wrapped completion error. The + // OpenAI SDK reports both with messages ending in a period + // ("Request was aborted.", "Request timed out."), which would not + // match task-level abort detection (message ending in "aborted"). + if ( + options?.abortSignal?.aborted || + error instanceof APIUserAbortError || + error instanceof APIConnectionTimeoutError || + (error instanceof Error && error.name === "AbortError") + ) { + throw createAbortError("Zoo Gateway") + } try { await surfaceGatewayApiError(error) } catch (surfaceError) {