feat(api): abort signal support for Zoo Gateway completePrompt

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<typeof console.error>
+
 		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 () => ({