feat(api): abort signal support for openai, openai-compatible base, zai, kimi-code (round 2)

src/api/providers/__tests__/base-openai-compatible-provider.spec.ts

@@ -1,7 +1,7 @@
 // npx vitest run api/providers/__tests__/base-openai-compatible-provider.spec.ts
 
 import { Anthropic } from "@anthropic-ai/sdk"
-import OpenAI from "openai"
+import OpenAI, { APIUserAbortError } from "openai"
 
 import type { ModelInfo } from "@roo-code/types"
 
@@ -14,6 +14,8 @@ const mockCreate = vi.fn()
 
 // Mock OpenAI module
 vi.mock("openai", () => ({
+	// Named export consumed by the provider for abort-error normalization
+	APIUserAbortError: class extends Error {},
 	default: vi.fn(function () {
 		return {
 			chat: {
@@ -49,6 +51,20 @@ class TestOpenAiCompatibleProvider extends BaseOpenAiCompatibleProvider<"test-mo
 	}
 }
 
+/**
+ * Captures the rejection of an operation as an Error. The abort contract always
+ * throws an Error; the guard keeps strict typing without a cast. Fails the test
+ * if the operation resolves.
+ */
+async function captureError(operation: Promise<unknown>): Promise<Error> {
+	try {
+		await operation
+	} catch (error) {
+		return error instanceof Error ? error : new Error(String(error))
+	}
+	throw new Error("Expected the operation to reject")
+}
+
 describe("BaseOpenAiCompatibleProvider", () => {
 	let handler: TestOpenAiCompatibleProvider
 
@@ -278,6 +294,171 @@ describe("BaseOpenAiCompatibleProvider", () => {
 		})
 	})
 
+	describe("abort signal wiring", () => {
+		it("should pass the metadata abort signal to the client request", async () => {
+			const controller = new AbortController()
+			mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))
+
+			const stream = handler.createMessage("system prompt", [], {
+				taskId: "test-task",
+				abortSignal: controller.signal,
+			})
+			await stream.next()
+
+			expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), {
+				signal: controller.signal,
+			})
+		})
+
+		it("should reject before issuing any request when the abort signal is already aborted", async () => {
+			const controller = new AbortController()
+			controller.abort()
+
+			await expect(async () => {
+				for await (const _ of handler.createMessage("system prompt", [], {
+					taskId: "test-task",
+					abortSignal: controller.signal,
+				})) {
+					// consume
+				}
+			}).rejects.toMatchObject({ name: "AbortError", message: "This operation was aborted" })
+			expect(mockCreate).not.toHaveBeenCalled()
+		})
+
+		it("should normalize the SDK APIUserAbortError from the stream path into the abort contract", async () => {
+			mockCreate.mockImplementationOnce(() => {
+				throw new APIUserAbortError()
+			})
+
+			const result = await captureError(
+				(async () => {
+					for await (const _ of handler.createMessage("system prompt", [])) {
+						// consume
+					}
+				})(),
+			)
+
+			// The SDK error has name "Error" and a message ending in a period; the
+			// provider must rethrow the Task.ts contract shape instead.
+			expect(result.name).toBe("AbortError")
+			expect(result.message).toBe("TestProvider request aborted")
+			expect(result.message.endsWith("aborted")).toBe(true)
+		})
+
+		it("should normalize an abort error raised during stream iteration into the abort contract", async () => {
+			mockCreate.mockImplementationOnce(() =>
+				(async function* () {
+					yield { choices: [{ delta: { content: "partial" } }] }
+					throw new APIUserAbortError()
+				})(),
+			)
+
+			const result = await captureError(
+				(async () => {
+					for await (const _ of handler.createMessage("system prompt", [])) {
+						// consume
+					}
+				})(),
+			)
+
+			// The iterator rejects after the first chunk; the provider must normalize
+			// it to the Task.ts contract shape (name + message ending in "aborted").
+			expect(result.name).toBe("AbortError")
+			expect(result.message).toBe("TestProvider request aborted")
+			expect(result.message.endsWith("aborted")).toBe(true)
+		})
+
+		it("should wrap a non-abort base_resp stream error with the provider prefix through the iteration wrapper", async () => {
+			mockCreate.mockImplementationOnce(() =>
+				asyncStreamFrom([
+					{
+						choices: [{ delta: { content: "partial" } }],
+						base_resp: { status_code: 1041, status_msg: "Invalid token" },
+					},
+				]),
+			)
+
+			const result = await captureError(
+				(async () => {
+					for await (const _ of handler.createMessage("system prompt", [])) {
+						// consume
+					}
+				})(),
+			)
+
+			expect(result.message).toBe("TestProvider completion error: TestProvider API Error (1041): Invalid token")
+		})
+		it("should still wrap non-abort request errors with the provider prefix", async () => {
+			mockCreate.mockImplementationOnce(() => {
+				throw new Error("boom")
+			})
+
+			const result = await captureError(
+				(async () => {
+					for await (const _ of handler.createMessage("system prompt", [])) {
+						// consume
+					}
+				})(),
+			)
+
+			expect(result.message).toBe("TestProvider completion error: boom")
+		})
+
+		it("should pass the completePrompt abort signal to the client request", async () => {
+			const controller = new AbortController()
+			mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })
+
+			const result = await handler.completePrompt("test prompt", { abortSignal: controller.signal })
+
+			expect(result).toBe("response")
+			expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), {
+				signal: controller.signal,
+			})
+		})
+
+		it("should merge completePrompt timeoutMs into the request signal", async () => {
+			mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })
+
+			await handler.completePrompt("test prompt", { timeoutMs: 5000 })
+
+			const requestOptions = mockCreate.mock.calls.at(-1)?.[1]
+			expect(requestOptions?.signal).toBeInstanceOf(AbortSignal)
+			expect(requestOptions?.signal.aborted).toBe(false)
+		})
+
+		it("should not set a request signal for zero completePrompt timeoutMs", async () => {
+			mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })
+
+			await handler.completePrompt("test prompt", { timeoutMs: 0 })
+
+			expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), undefined)
+		})
+
+		it("should reject before any request when the completePrompt signal is already aborted", async () => {
+			const controller = new AbortController()
+			controller.abort()
+
+			await expect(
+				handler.completePrompt("test prompt", { abortSignal: controller.signal }),
+			).rejects.toMatchObject({
+				name: "AbortError",
+				message: "This operation was aborted",
+			})
+			expect(mockCreate).not.toHaveBeenCalled()
+		})
+
+		it("should normalize the SDK APIUserAbortError from completePrompt", async () => {
+			mockCreate.mockImplementationOnce(() => {
+				throw new APIUserAbortError()
+			})
+
+			const result = await captureError(handler.completePrompt("test prompt"))
+
+			expect(result.name).toBe("AbortError")
+			expect(result.message).toBe("TestProvider request aborted")
+		})
+	})
+
 	describe("Tool call handling", () => {
 		it("should yield tool_call_end events when finish_reason is tool_calls", async () => {
 			mockCreate.mockImplementationOnce(() =>

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()
+	})
+})

src/api/providers/__tests__/fireworks.spec.ts

@@ -23,6 +23,8 @@ vi.mock("openai", () => ({
 			},
 		}
 	}),
+	// Named export consumed by the provider for abort-error normalization
+	APIUserAbortError: class extends Error {},
 }))
 
 describe("FireworksHandler", () => {

src/api/providers/__tests__/kimi-code.spec.ts

@@ -18,6 +18,32 @@ vi.mock("../../../integrations/kimi-code/oauth", () => ({
 
 vi.mock("../fetchers/modelCache", () => ({ getModels: mockGetModels }))
 
+/**
+ * Spies on the inherited OpenAI client's chat.completions.create. `client` is
+ * protected on the OpenAiHandler base (not on the public interface), so it is
+ * reached through a documented `as unknown as` double assertion (AGENTS.md
+ * last resort; no `as any`).
+ */
+function completionsCreate(handler: KimiCodeHandler): ReturnType<typeof vi.fn> {
+	const client = (
+		handler as unknown as {
+			client: { chat: { completions: Record<string, (...args: never[]) => never> } }
+		}
+	).client
+	return vi.spyOn(client.chat.completions, "create") as unknown as ReturnType<typeof vi.fn>
+}
+
+/** Captures the rejection of an operation; fails the test if it resolves. */
+async function captureError(operation: Promise<unknown>): Promise<Error> {
+	try {
+		await operation
+	} catch (error) {
+		// Abort normalization always throws an Error; the guard keeps strict typing without casts.
+		return error instanceof Error ? error : new Error(String(error))
+	}
+	throw new Error("Expected the operation to reject")
+}
+
 describe("KimiCodeHandler", () => {
 	beforeEach(() => {
 		clearAllMocks()
@@ -117,8 +143,7 @@ describe("KimiCodeHandler", () => {
 	it("force-refreshes and retries exactly once after a non-streaming OAuth 401", async () => {
 		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
 		const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 })
-		const createCompletion = vi
-			.spyOn((handler as any).client.chat.completions, "create")
+		const createCompletion = completionsCreate(handler)
 			.mockRejectedValueOnce(unauthorized)
 			.mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] })
 
@@ -235,4 +260,89 @@ describe("KimiCodeHandler", () => {
 		})
 		expect(handler.getModel().reasoning).toEqual({ reasoning_effort: "max" })
 	})
+
+	it("forwards the metadata abort signal to the inherited OpenAI SDK request", async () => {
+		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
+		const controller = new AbortController()
+		const streamChunks = (async function* () {
+			yield { choices: [{ delta: { content: "hi" } }] }
+		})()
+		const createCompletion = completionsCreate(handler).mockResolvedValueOnce(streamChunks)
+
+		const gen = handler.createMessage("system", [{ role: "user", content: "test" }], {
+			taskId: "test-task",
+			abortSignal: controller.signal,
+		})
+		const first = await gen.next()
+
+		expect(first.value).toEqual({ type: "text", text: "hi" })
+		expect(createCompletion).toHaveBeenCalledWith(expect.anything(), { signal: controller.signal })
+	})
+
+	it("rejects before any request when the createMessage abort signal is already aborted", async () => {
+		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
+		const controller = new AbortController()
+		controller.abort()
+		const createCompletion = completionsCreate(handler)
+
+		const gen = handler.createMessage("system", [{ role: "user", content: "test" }], {
+			taskId: "test-task",
+			abortSignal: controller.signal,
+		})
+
+		await expect(async () => {
+			for await (const _ of gen) {
+				// consume
+			}
+		}).rejects.toMatchObject({ name: "AbortError", message: "This operation was aborted" })
+		expect(createCompletion).not.toHaveBeenCalled()
+	})
+
+	it("forwards completePrompt abort options through the override on both 401 retry attempts", async () => {
+		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
+		const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 })
+		const createCompletion = completionsCreate(handler)
+			.mockRejectedValueOnce(unauthorized)
+			.mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] })
+		const controller = new AbortController()
+
+		await expect(handler.completePrompt("test", { abortSignal: controller.signal })).resolves.toBe("retried")
+		expect(mockForceRefreshAccessToken).toHaveBeenCalledOnce()
+		expect(createCompletion).toHaveBeenCalledTimes(2)
+		for (const call of createCompletion.mock.calls) {
+			expect(call[1]).toEqual({ signal: controller.signal })
+		}
+	})
+
+	it("rejects before any request when the completePrompt signal is already aborted", async () => {
+		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
+		const controller = new AbortController()
+		controller.abort()
+		const createCompletion = completionsCreate(handler)
+
+		await expect(handler.completePrompt("test", { abortSignal: controller.signal })).rejects.toMatchObject({
+			name: "AbortError",
+			message: "This operation was aborted",
+		})
+		expect(createCompletion).not.toHaveBeenCalled()
+	})
+
+	it("surfaces a normalized AbortError when the SDK aborts a streaming request", async () => {
+		const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
+		// The real SDK class: this spec does not mock the openai module.
+		const { APIUserAbortError } = await import("openai")
+		completionsCreate(handler).mockRejectedValueOnce(new APIUserAbortError())
+
+		const gen = handler.createMessage("system", [{ role: "user", content: "test" }], { taskId: "test-task" })
+		const result = await captureError(
+			(async () => {
+				for await (const _ of gen) {
+					// consume
+				}
+			})(),
+		)
+
+		expect(result.name).toBe("AbortError")
+		expect(result.message).toBe("OpenAI request aborted")
+	})
 })