diff --git a/src/api/providers/__tests__/zoo-gateway.spec.ts b/src/api/providers/__tests__/zoo-gateway.spec.ts index fe088bb5c9..d5bfd5320f 100644 --- a/src/api/providers/__tests__/zoo-gateway.spec.ts +++ b/src/api/providers/__tests__/zoo-gateway.spec.ts @@ -26,6 +26,7 @@ vitest.mock("../../../i18n", () => ({ t: (key: string) => key, })) +import type { MockInstance } from "vitest" import OpenAI from "openai" import { zooGatewayDefaultModelId, ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types" @@ -500,10 +501,17 @@ describe("ZooGatewayHandler", () => { }) describe("completePrompt", () => { + let consoleError: MockInstance + beforeEach(() => { mockCreate.mockImplementation(async () => ({ choices: [{ message: { role: "assistant", content: "Test completion response" } }], })) + consoleError = vitest.spyOn(console, "error").mockImplementation(() => {}) + }) + + afterEach(() => { + consoleError.mockRestore() }) it("returns completion text from the gateway", async () => { @@ -520,6 +528,7 @@ describe("ZooGatewayHandler", () => { temperature: ZOO_GATEWAY_DEFAULT_TEMPERATURE, max_completion_tokens: 64000, }), + {}, ) }) @@ -534,6 +543,264 @@ describe("ZooGatewayHandler", () => { ) }) + it("forwards a per-request signal and the timeout to the SDK", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + + await handler.completePrompt("Test", { abortSignal: caller.signal, timeoutMs: 5_000 }) + + const requestOptions = mockCreate.mock.calls[0][1] as { signal?: AbortSignal; timeout?: number } + expect(requestOptions.signal).toBeInstanceOf(AbortSignal) + expect(requestOptions.signal).not.toBe(caller.signal) + expect(requestOptions.timeout).toBe(5_000) + }) + + it("sends no signal when the caller gives none", async () => { + const handler = new ZooGatewayHandler(mockOptions) + + await handler.completePrompt("Test") + + expect(mockCreate.mock.calls[0][1]).toEqual({}) + }) + + it.each([ + ["with a session", mockOptions], + ["without a session", { zooGatewayModelId: "anthropic/claude-sonnet-4" }], + ])( + "rejects a pre-aborted signal with an AbortError, %s, before model lookup", + async (_label, options: ApiHandlerOptions) => { + const { getModels } = await import("../fetchers/modelCache") + const handler = new ZooGatewayHandler(options) + const caller = new AbortController() + caller.abort() + + await expect(handler.completePrompt("Test", { abortSignal: caller.signal })).rejects.toMatchObject({ + name: "AbortError", + }) + expect(getModels).not.toHaveBeenCalled() + expect(mockCreate).not.toHaveBeenCalled() + }, + ) + + it("keeps an ordinary model lookup failure while the caller signal is live", async () => { + const { getModels, refreshModels } = await import("../fetchers/modelCache") + const lookupFailure = new Error("model catalog unavailable") + vitest.mocked(getModels).mockRejectedValue(lookupFailure) + vitest.mocked(refreshModels).mockRejectedValue(lookupFailure) + const handler = new ZooGatewayHandler(mockOptions) + + await expect(handler.completePrompt("Test", { abortSignal: new AbortController().signal })).rejects.toBe( + lookupFailure, + ) + }) + + it.each([ + ["a raw AbortError", Object.assign(new Error("The user aborted a request"), { name: "AbortError" })], + ["the SDK abort error", new Error("Request was aborted.")], + ])( + "normalizes %s from the model lookup to the provider AbortError while the caller signal is live", + async (_label, lookupFailure: Error) => { + const { getModels, refreshModels } = await import("../fetchers/modelCache") + vitest.mocked(getModels).mockRejectedValue(lookupFailure) + vitest.mocked(refreshModels).mockRejectedValue(lookupFailure) + const handler = new ZooGatewayHandler(mockOptions) + + await expect( + handler.completePrompt("Test", { abortSignal: new AbortController().signal }), + ).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + expect(mockCreate).not.toHaveBeenCalled() + }, + ) + + it("rejects with an AbortError when the timeout alone expires during the SDK call", async () => { + const handler = new ZooGatewayHandler(mockOptions) + let sdkSignal: AbortSignal | undefined + mockCreate.mockImplementation((_params: unknown, requestOptions: { signal?: AbortSignal }) => { + sdkSignal = requestOptions.signal + return new Promise(() => {}) + }) + + await expect(handler.completePrompt("Test", { timeoutMs: 20 })).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + expect(sdkSignal?.aborted).toBe(true) + expect(showErrorMessage).not.toHaveBeenCalled() + }) + + it("settles on abort while a gateway error notification waits for the user", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + mockCreate.mockRejectedValue(makeApiError(402)) + showErrorMessage.mockImplementation(() => new Promise(() => {})) + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }) + await vitest.waitFor(() => expect(showErrorMessage).toHaveBeenCalled()) + caller.abort() + + await expect(pending).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + // The abort is not a failure of the notification, so nothing is logged. + expect(consoleError).not.toHaveBeenCalled() + }) + + it.each([ + ["with a live caller signal", true, new Error("notification failed"), "notification failed"], + ["without a caller signal", false, "notification dismissed", "notification dismissed"], + ])( + "logs a failed gateway error notification and still rejects with the gateway error, %s", + async (_label, withSignal: boolean, notificationFailure: unknown, logged: string) => { + const handler = new ZooGatewayHandler(mockOptions) + mockCreate.mockRejectedValue(makeApiError(402, { message: "out of credits" })) + showErrorMessage.mockRejectedValue(notificationFailure) + const options = withSignal ? { abortSignal: new AbortController().signal } : undefined + + await expect(handler.completePrompt("Test", options)).rejects.toThrow( + "Zoo Gateway completion error: out of credits", + ) + expect(consoleError).toHaveBeenCalledWith("Failed to surface Zoo Gateway error:", logged) + }, + ) + + it("rethrows an SDK failure that is not an Error unchanged", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const failure = { reason: "upstream failure" } + mockCreate.mockRejectedValue(failure) + + await expect(handler.completePrompt("Test", { abortSignal: new AbortController().signal })).rejects.toBe( + failure, + ) + }) + + it("rejects with an AbortError when the timeout alone expires during model lookup", async () => { + const { getModels } = await import("../fetchers/modelCache") + vitest.mocked(getModels).mockImplementation(() => new Promise(() => {})) + const handler = new ZooGatewayHandler(mockOptions) + + await expect(handler.completePrompt("Test", { timeoutMs: 20 })).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("rejects with an AbortError when the signal aborts during model lookup", async () => { + const { getModels } = await import("../fetchers/modelCache") + vitest.mocked(getModels).mockImplementation(() => new Promise(() => {})) + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }) + caller.abort() + + await expect(pending).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("rejects with an AbortError at once when aborted mid-flight, without a gateway error", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + let sdkSignal: AbortSignal | undefined + // The SDK call never settles, as when the SDK sleeps between retries. + mockCreate.mockImplementation((_params: unknown, requestOptions: { signal?: AbortSignal }) => { + sdkSignal = requestOptions.signal + return new Promise(() => {}) + }) + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }) + await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled()) + caller.abort() + + await expect(pending).rejects.toMatchObject({ + name: "AbortError", + message: "The Zoo Gateway request was aborted", + }) + expect(sdkSignal?.aborted).toBe(true) + expect(showErrorMessage).not.toHaveBeenCalled() + }) + + it("reports an SDK abort error after the caller aborts as an AbortError, not a gateway error", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + mockCreate.mockImplementation(async () => { + caller.abort() + throw new Error("Request was aborted.") + }) + + await expect(handler.completePrompt("Test", { abortSignal: caller.signal })).rejects.toMatchObject({ + name: "AbortError", + }) + expect(showErrorMessage).not.toHaveBeenCalled() + }) + + it("shows no gateway error for a 402 that rejects just before the caller aborts", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + let rejectCompletion: (reason: unknown) => void = () => {} + const completion = new Promise((_resolve, reject) => { + rejectCompletion = reject + }) + mockCreate.mockImplementation(() => completion) + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }) + await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled()) + // This callback runs after the handler's own race settles with the 402 but before completePrompt + // resumes, so only the abort check before the notification keeps the 402 notification away. + completion.catch(() => caller.abort()) + rejectCompletion(makeApiError(402)) + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }) + expect(showErrorMessage).not.toHaveBeenCalled() + }) + + it("does not return a completion that resolves just before the caller aborts", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + let resolveCompletion: (value: unknown) => void = () => {} + const completion = new Promise((resolve) => { + resolveCompletion = resolve + }) + mockCreate.mockImplementation(() => completion) + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }) + await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled()) + // This callback runs after the handler's own race settles but before completePrompt resumes, + // so only the check after the await can drop the late result. + void completion.then(() => caller.abort()) + resolveCompletion({ choices: [{ message: { role: "assistant", content: "late" } }] }) + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }) + }) + + it.each([ + ["a completed request", () => {}], + ["a failed request", () => mockCreate.mockRejectedValue(new Error("upstream failure"))], + ["a request aborted mid-flight", () => mockCreate.mockImplementation(() => new Promise(() => {}))], + ])("removes the same abort listener from the caller's signal after %s", async (_label, setup) => { + setup() + const handler = new ZooGatewayHandler(mockOptions) + const caller = new AbortController() + const addListener = vitest.spyOn(caller.signal, "addEventListener") + const removeListener = vitest.spyOn(caller.signal, "removeEventListener") + + const pending = handler.completePrompt("Test", { abortSignal: caller.signal }).catch(() => undefined) + await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled()) + caller.abort() + await pending + + const abortListeners = addListener.mock.calls.filter(([type]) => type === "abort") + expect(abortListeners).toHaveLength(1) + expect(removeListener).toHaveBeenCalledWith("abort", abortListeners[0][1]) + }) + it("returns an empty string when the model returns no content", async () => { const handler = new ZooGatewayHandler(mockOptions) mockCreate.mockImplementation(async () => ({ diff --git a/src/api/providers/zoo-gateway.ts b/src/api/providers/zoo-gateway.ts index 5f33ee48fe..2171c49dda 100644 --- a/src/api/providers/zoo-gateway.ts +++ b/src/api/providers/zoo-gateway.ts @@ -30,6 +30,18 @@ import type { import { NOT_PROVIDED } from "./constants" import { RouterProvider } from "./router-provider" import { getZooGatewayModels } from "./fetchers/zoo-gateway" +import { + bridgeExternalAbortSignal, + createAbortError, + hasUsableTimeout, + isRequestAborted, + mergeAbortSignalAndTimeout, + rejectOnAbort, + throwIfAborted, + toSdkTimeoutMs, +} from "./utils/abort-signal" + +const ZOO_GATEWAY_PROVIDER_NAME = "Zoo Gateway" function getApiErrorStatus(error: unknown): number | undefined { if (typeof error === "object" && error !== null && "status" in error) { @@ -294,12 +306,33 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { - this.ensureAuthenticated() - - const { id: modelId, info } = await this.fetchModel() + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Same abort-signal contract as the other providers (see requesty.ts): the SDK gets a + // per-request signal bridged from the caller's signal, so its abort listener does not stay on + // the caller's long-lived signal. The `finally` below detaches the bridge on every exit path. + const { signal: bridgedCallerSignal, dispose: removeExternalAbortListener } = bridgeExternalAbortSignal( + options?.abortSignal, + ) + const requestAbortSignal = mergeAbortSignalAndTimeout(bridgedCallerSignal, options?.timeoutMs) try { + // A cancel that came first wins over every other failure, also a missing session. + throwIfAborted(requestAbortSignal) + this.ensureAuthenticated() + + let model: Awaited> + try { + model = requestAbortSignal + ? await rejectOnAbort(this.fetchModel(), requestAbortSignal, ZOO_GATEWAY_PROVIDER_NAME) + : await this.fetchModel() + } catch (error) { + if (isRequestAborted(error, requestAbortSignal)) { + throw createAbortError(ZOO_GATEWAY_PROVIDER_NAME) + } + throw error + } + const { id: modelId, info } = model + const requestOptions: OpenAI.Chat.ChatCompletionCreateParams = { model: modelId, messages: [{ role: "user", content: prompt }], @@ -312,21 +345,54 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio requestOptions.max_completion_tokens = info.maxTokens - const response = await this.client.chat.completions.create(requestOptions) - return response.choices[0]?.message.content || "" - } catch (error) { + const createOptions: OpenAI.RequestOptions = { + ...(requestAbortSignal && { signal: requestAbortSignal }), + ...(hasUsableTimeout(options?.timeoutMs) && { timeout: toSdkTimeoutMs(options.timeoutMs) }), + } + + let response: OpenAI.Chat.ChatCompletion try { - await surfaceGatewayApiError(error) - } catch (surfaceError) { - console.error( - "Failed to surface Zoo Gateway error:", - surfaceError instanceof Error ? surfaceError.message : surfaceError, - ) + const pendingCompletion = this.client.chat.completions.create(requestOptions, createOptions) + // Race the completion against the signal: the SDK does not observe the signal while it + // sleeps between retries, so awaiting the SDK call alone can outlive the cancel. + response = requestAbortSignal + ? await rejectOnAbort(pendingCompletion, requestAbortSignal, ZOO_GATEWAY_PROVIDER_NAME) + : await pendingCompletion + } catch (error) { + // A cancel or timeout is not a gateway failure: do not show a gateway error for it. + if (requestAbortSignal?.aborted) { + throw createAbortError(ZOO_GATEWAY_PROVIDER_NAME) + } + try { + // The notification can wait for the user. A cancel or timeout must still settle at once. + const surfacing = surfaceGatewayApiError(error) + await (requestAbortSignal + ? rejectOnAbort(surfacing, requestAbortSignal, ZOO_GATEWAY_PROVIDER_NAME) + : surfacing) + } catch (surfaceError) { + if (!requestAbortSignal?.aborted) { + console.error( + "Failed to surface Zoo Gateway error:", + surfaceError instanceof Error ? surfaceError.message : surfaceError, + ) + } + } + if (requestAbortSignal?.aborted) { + throw createAbortError(ZOO_GATEWAY_PROVIDER_NAME) + } + if (error instanceof Error) { + throw new Error(`Zoo Gateway completion error: ${error.message}`) + } + throw error } - if (error instanceof Error) { - throw new Error(`Zoo Gateway completion error: ${error.message}`) + + if (requestAbortSignal?.aborted) { + // The response resolved after the request was aborted: do not return the late result. + throw createAbortError(ZOO_GATEWAY_PROVIDER_NAME) } - throw error + return response.choices[0]?.message.content || "" + } finally { + removeExternalAbortListener() } } }