From 2a58311c0b33b27843544491a04e4c535709395a Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 10:11:42 +0800 Subject: [PATCH 1/6] feat(api): abort signal support for gemini, mistral, lite-llm (completePrompt + createMessage) --- .../__tests__/gemini-handler.spec.ts | 38 ++++ src/api/providers/__tests__/gemini.spec.ts | 209 ++++++++++++++++++ src/api/providers/__tests__/lite-llm.spec.ts | 104 +++++++++ src/api/providers/__tests__/mistral.spec.ts | 145 +++++++++++- src/api/providers/__tests__/vertex.spec.ts | 56 +++++ src/api/providers/gemini.ts | 56 ++++- src/api/providers/lite-llm.ts | 54 ++++- src/api/providers/mistral.ts | 151 ++++++++----- 8 files changed, 751 insertions(+), 62 deletions(-) diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 110f60289c..364f62e23d 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -55,6 +55,44 @@ describe("GeminiHandler backend support", () => { expect(promptConfig.tools).toBeUndefined() }) + it("completePrompt should pass abort signal through to client via httpOptions", async () => { + const options = { + apiProvider: "gemini", + enableUrlContext: false, + enableGrounding: false, + } as ApiHandlerOptions + const handler = new GeminiHandler(options) + + const controller = new AbortController() + const stub = vi.fn().mockResolvedValue({ text: "response" }) + handler["client"].models.generateContent = stub + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + + expect(stub).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + abortSignal: controller.signal, + }), + }), + ) + }) + + it("completePrompt should work without options (backward compatible)", async () => { + const options = { + apiProvider: "gemini", + enableUrlContext: false, + enableGrounding: false, + } as ApiHandlerOptions + const handler = new GeminiHandler(options) + + const stub = vi.fn().mockResolvedValue({ text: "response" }) + handler["client"].models.generateContent = stub + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + describe("error scenarios", () => { it("should handle grounding metadata extraction failure gracefully", async () => { const options = { diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 701c453e3e..59968060dc 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -12,14 +12,22 @@ vitest.mock("@roo-code/telemetry", () => ({ import { Anthropic } from "@anthropic-ai/sdk" +import type { GenerateContentResponse } from "@google/genai" + import { type ModelInfo, geminiDefaultModelId, ApiProviderError } from "@roo-code/types" import { t } from "i18next" import { GeminiHandler } from "../gemini" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" const GEMINI_MODEL_NAME = geminiDefaultModelId +// @google/genai's GenerateContentResponse exposes `text` via a getter backed by +// `candidates`, so the stub only carries the field the provider reads; the double +// cast is the least-friction way to satisfy the class type in mocks. +const stubGenerateContentResponse = (text: string) => ({ text }) as unknown as GenerateContentResponse + describe("GeminiHandler", () => { let handler: GeminiHandler @@ -342,6 +350,71 @@ describe("GeminiHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should pass abort signal through to client via config.abortSignal", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + abortSignal: controller.signal, + httpOptions: undefined, + temperature: 1, + }, + }) + }) + + it("should work without options (backward compatible)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: undefined, + temperature: 1, + }, + }) + }) + + it("should pass timeoutMs through to client via httpOptions with abortSignal on config", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + abortSignal: controller.signal, + httpOptions: { timeout: 10000 }, + temperature: 1, + }, + }) + }) + + it("should pass only timeoutMs when no signal is provided", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: { timeout: 5000 }, + temperature: 1, + }, + }) + }) }) describe("getModel", () => { @@ -475,6 +548,142 @@ describe("GeminiHandler", () => { }) }) + describe("completePrompt request options", () => { + it("should pass timeout and baseUrl through httpOptions", async () => { + const handlerWithBaseUrl = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "https://gemini.example.test", + }) + handlerWithBaseUrl["client"] = handler["client"] + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + const result = await handlerWithBaseUrl.completePrompt("Test prompt", { timeoutMs: 1234 }) + + expect(result).toBe("Response") + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + httpOptions: { + timeout: 1234, + baseUrl: "https://gemini.example.test", + }, + }), + }), + ) + }) + + it("should pass abortSignal on config instead of httpOptions", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + abortSignal: controller.signal, + httpOptions: undefined, + }), + }), + ) + }) + + it("should omit httpOptions when timeoutMs and baseUrl are not provided", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await handler.completePrompt("Test prompt") + + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + httpOptions: undefined, + }), + }), + ) + }) + }) + + describe("createMessage abort signal (bridging)", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "You are a helpful assistant", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(handler["client"].models.generateContentStream).not.toHaveBeenCalled() + }) + + it("should abort the in-flight request when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + const stub = vi.fn().mockImplementation(async (params: { config?: { abortSignal?: AbortSignal } }) => { + capturedSignal = params.config?.abortSignal + return (async function* () { + yield { text: "partial" } + if (capturedSignal?.aborted) { + throw new DOMException("aborted", "AbortError") + } + await new Promise((_resolve, reject) => { + capturedSignal?.addEventListener( + "abort", + () => reject(new DOMException("aborted", "AbortError")), + { once: true }, + ) + }) + })() + }) + handler["client"].models.generateContentStream = stub + + const stream = handler.createMessage( + "You are a helpful assistant", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + + it("should not set config.abortSignal when no external signal is provided", async () => { + const stub = vi.fn().mockReturnValue((async function* () {})()) + handler["client"].models.generateContentStream = stub + + await collectStream(handler.createMessage("You are a helpful assistant", messages)) + + const config = stub.mock.calls[0][0].config + expect(config.abortSignal).toBeUndefined() + }) + }) + describe("error telemetry", () => { const mockMessages: Anthropic.Messages.MessageParam[] = [ { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index eee5cf52bb..badab1e1a7 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -6,6 +6,7 @@ import { ApiHandlerOptions } from "../../../shared/api" import { litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" +import { makeCreateMessageMetadata } from "../../../test-utils/api" // Mock vscode first to avoid import errors vi.mock("vscode", () => ({ @@ -1235,4 +1236,107 @@ describe("LiteLLMHandler", () => { expect(requestHeaders).not.toHaveProperty("X-Zoo-Session-ID") }) }) + + describe("completePrompt", () => { + it("should pass abort signal through to client", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + const controller = new AbortController() + 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 () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), + ) + }) + + it("should merge signal and timeoutMs together", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal, timeout: 10000 }), + ) + }) + + it("should work without options (backward compatible)", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + }) + + describe("createMessage abort signal (bridging)", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + // The stream is built inside the mock implementation so that capturedSignal + // is already set before the abort-aware chunk is created. + mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "partial" } }], + usage: undefined, + }, + new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("aborted", "AbortError")) + if (capturedSignal?.aborted) { + onAbort() + return + } + capturedSignal?.addEventListener("abort", onAbort, { once: true }) + }), + ]) + return { withResponse: vi.fn().mockResolvedValue({ data: mockStream }) } + }) + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + }) }) diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index f2a7591bd8..aff33b6030 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -54,6 +54,7 @@ import { MistralHandler } from "../mistral" import type { ApiHandlerOptions } from "../../../shared/api" import type { ApiHandlerCreateMessageMetadata } from "../../index" import type { ApiStreamTextChunk, ApiStreamReasoningChunk, ApiStreamToolCallPartialChunk } from "../../transform/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" describe("MistralHandler", () => { let handler: MistralHandler @@ -447,11 +448,14 @@ describe("MistralHandler", () => { const prompt = "Test prompt" const result = await handler.completePrompt(prompt) - expect(mockComplete).toHaveBeenCalledWith({ - model: mockOptions.apiModelId, - messages: [{ role: "user", content: prompt }], - temperature: 0, - }) + expect(mockComplete).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: prompt }], + temperature: 0, + }, + undefined, + ) expect(result).toBe("Test response") }) @@ -483,5 +487,136 @@ describe("MistralHandler", () => { mockComplete.mockRejectedValueOnce(new Error("API Error")) await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Mistral completion error: API Error") }) + + it("should pass abort signal through to client", async () => { + const controller = new AbortController() + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + fetchOptions: { signal: controller.signal }, + }) + }) + + it("should work without options (backward compatible)", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) + }) + + it("should pass timeout through to client", async () => { + const controller = new AbortController() + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + fetchOptions: { signal: controller.signal }, + timeoutMs: 5000, + }) + }) + + it("should pass only timeoutMs when no signal provided", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 3000 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeoutMs: 3000, + }) + }) + + it("should still forward timeoutMs=0 (uses !== undefined check, not truthy check)", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeoutMs: 0, + }) + }) + }) + + describe("createMessage abort signal bridging", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: [{ type: "text", text: "Hello!" }], + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementationOnce( + async (_options: unknown, requestOptions?: { fetchOptions?: { signal?: AbortSignal } }) => { + capturedSignal = requestOptions?.fetchOptions?.signal + return asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { content: "partial" }, + index: 0, + }, + ], + }, + }, + new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("aborted", "AbortError")) + if (capturedSignal?.aborted) { + onAbort() + return + } + capturedSignal?.addEventListener("abort", onAbort, { once: true }) + }), + ]) + }, + ) + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + + it("should not pass a signal to the stream call when no external signal is provided", async () => { + const stream = handler.createMessage(systemPrompt, messages) + await collectStream(stream) + const streamOptions = mockCreate.mock.calls[0][1] + expect(streamOptions).toBeUndefined() + }) }) }) diff --git a/src/api/providers/__tests__/vertex.spec.ts b/src/api/providers/__tests__/vertex.spec.ts index a304518ca7..99c5787e0e 100644 --- a/src/api/providers/__tests__/vertex.spec.ts +++ b/src/api/providers/__tests__/vertex.spec.ts @@ -21,10 +21,19 @@ vitest.mock("@roo-code/telemetry", () => ({ import { Anthropic } from "@anthropic-ai/sdk" +import type { GenerateContentResponse } from "@google/genai" + import { ApiStreamChunk } from "../../transform/stream" import { t } from "i18next" import { VertexHandler } from "../vertex" +import { collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" + +// @google/genai's GenerateContentResponse exposes `text` via a getter backed by +// `candidates`, so the stub only carries the field the provider reads; the double +// cast is the least-friction way to satisfy the class type in mocks. +const stubGenerateContentResponse = (text: string) => ({ text }) as unknown as GenerateContentResponse describe("VertexHandler", () => { let handler: VertexHandler @@ -137,6 +146,53 @@ describe("VertexHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should pass abort signal through to client via config.abortSignal", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: expect.any(String), + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: expect.objectContaining({ + abortSignal: controller.signal, + httpOptions: undefined, + temperature: 1, + }), + }), + ) + }) + + it("should work without options (backward compatible)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + }) + + describe("createMessage abort signal (inherited from GeminiHandler)", () => { + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "You are a helpful assistant", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(handler["client"].models.generateContentStream).not.toHaveBeenCalled() + }) }) describe("getModel", () => { diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index ec0d14e4c9..bc6bcfe2f1 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -344,7 +344,30 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } - const params: GenerateContentParameters = { model, contents, config } + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight generateContentStream request + // can be cancelled. The @google/genai SDK merges this signal with its own + // timeout handling, which is preserved rather than replaced. + // A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("Gemini request aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + + const params: GenerateContentParameters = { + model, + contents, + config: requestAbortController ? { ...config, abortSignal: requestAbortController.signal } : config, + } try { const result = await this.client.models.generateContentStream(params) @@ -477,6 +500,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("Gemini request aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") TelemetryService.instance.captureException(apiError) @@ -486,6 +514,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } throw error + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -585,14 +617,25 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl const temperatureConfig: number | undefined = supportsTemperature ? (this.options.modelTemperature ?? info.defaultTemperature ?? 1) : info.defaultTemperature + const httpOpts: { timeout?: number; baseUrl?: string } = {} + if (options?.timeoutMs !== undefined) { + httpOpts.timeout = options.timeoutMs + } + if (this.options.googleGeminiBaseUrl) { + httpOpts.baseUrl = this.options.googleGeminiBaseUrl + } const promptConfig: GenerateContentConfig = { - httpOptions: this.options.googleGeminiBaseUrl - ? { baseUrl: this.options.googleGeminiBaseUrl } - : undefined, + httpOptions: Object.keys(httpOpts).length > 0 ? httpOpts : undefined, temperature: temperatureConfig, } + // @google/genai expects request cancellation on config.abortSignal + // (not httpOptions.signal), so the signal is passed directly to the config. + if (options?.abortSignal) { + promptConfig.abortSignal = options.abortSignal + } + const request = { model, contents: [{ role: "user", parts: [{ text: prompt }] }], @@ -613,6 +656,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return text } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("Gemini completion aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "completePrompt") TelemetryService.instance.captureException(apiError) diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index 8cfe2d0a19..da82dd19cb 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -246,9 +246,31 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestHeaders["X-Zoo-Session-ID"] = metadata.taskId } + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight streaming request can be + // cancelled. A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("LiteLLM streaming aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + try { const { data: completion } = await this.client.chat.completions - .create(requestOptions, { headers: requestHeaders }) + .create( + requestOptions, + requestAbortController + ? { headers: requestHeaders, signal: requestAbortController.signal } + : { headers: requestHeaders }, + ) .withResponse() let lastUsage @@ -315,10 +337,19 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa yield usageData } } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("LiteLLM streaming aborted", "AbortError") + } if (error instanceof Error) { throw new Error(`LiteLLM streaming error: ${error.message}`) } throw error + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -345,9 +376,28 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestOptions.max_tokens = info.maxTokens } - const response = await this.client.chat.completions.create(requestOptions) + // Build request options with abortSignal and/or timeout. The OpenAI SDK + // treats a timeout of 0 as an immediate timeout, so non-positive timeoutMs + // values disable the timeout instead of being forwarded. + 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, + Object.keys(createOptions).length > 0 ? createOptions : undefined, + ) return response.choices[0]?.message.content || "" } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("LiteLLM completion aborted", "AbortError") + } if (error instanceof Error) { throw new Error(`LiteLLM completion error: ${error.message}`) } diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index c7816feaa2..9c304f7eb0 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -101,66 +101,98 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand // Temporary debug log for QA // console.log("[MISTRAL DEBUG] Raw API request body:", requestOptions) + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight streaming request can be + // cancelled. A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + let response try { - response = await this.client.chat.stream(requestOptions) - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error) - const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") - TelemetryService.instance.captureException(apiError) - throw new Error(`Mistral completion error: ${errorMessage}`) - } + if (requestAbortController) { + response = await this.client.chat.stream(requestOptions, { + fetchOptions: { signal: requestAbortController.signal }, + }) + } else { + response = await this.client.chat.stream(requestOptions) + } - for await (const event of response) { - const delta = event.data.choices[0]?.delta - - if (delta?.content) { - if (typeof delta.content === "string") { - // Handle string content as text - yield { type: "text", text: delta.content } - } else if (Array.isArray(delta.content)) { - // Handle array of content chunks - // The SDK v1.9.18 supports ThinkChunk with type "thinking" - for (const chunk of delta.content as ContentChunkWithThinking[]) { - if (chunk.type === "thinking" && chunk.thinking) { - // Handle thinking content as reasoning chunks - // ThinkChunk has a 'thinking' property that contains an array of text/reference chunks - for (const thinkingPart of chunk.thinking) { - if (thinkingPart.type === "text" && thinkingPart.text) { - yield { type: "reasoning", text: thinkingPart.text } + for await (const event of response) { + const delta = event.data.choices[0]?.delta + + if (delta?.content) { + if (typeof delta.content === "string") { + // Handle string content as text + yield { type: "text", text: delta.content } + } else if (Array.isArray(delta.content)) { + // Handle array of content chunks + // The SDK v1.9.18 supports ThinkChunk with type "thinking" + for (const chunk of delta.content as ContentChunkWithThinking[]) { + if (chunk.type === "thinking" && chunk.thinking) { + // Handle thinking content as reasoning chunks + // ThinkChunk has a 'thinking' property that contains an array of text/reference chunks + for (const thinkingPart of chunk.thinking) { + if (thinkingPart.type === "text" && thinkingPart.text) { + yield { type: "reasoning", text: thinkingPart.text } + } } + } else if (chunk.type === "text" && chunk.text) { + // Handle text content normally + yield { type: "text", text: chunk.text } } - } else if (chunk.type === "text" && chunk.text) { - // Handle text content normally - yield { type: "text", text: chunk.text } } } } - } - // Handle tool calls in stream - // Mistral SDK provides tool_calls in delta similar to OpenAI format - const toolCalls = (delta as { toolCalls?: MistralToolCall[] })?.toolCalls - if (toolCalls) { - for (let i = 0; i < toolCalls.length; i++) { - const toolCall = toolCalls[i] - yield { - type: "tool_call_partial", - index: i, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + // Handle tool calls in stream + // Mistral SDK provides tool_calls in delta similar to OpenAI format + const toolCalls = (delta as { toolCalls?: MistralToolCall[] })?.toolCalls + if (toolCalls) { + for (let i = 0; i < toolCalls.length; i++) { + const toolCall = toolCalls[i] + yield { + type: "tool_call_partial", + index: i, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } } } - } - if (event.data.usage) { - yield { - type: "usage", - inputTokens: event.data.usage.promptTokens || 0, - outputTokens: event.data.usage.completionTokens || 0, + if (event.data.usage) { + yield { + type: "usage", + inputTokens: event.data.usage.promptTokens || 0, + outputTokens: event.data.usage.completionTokens || 0, + } } } + } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } + const errorMessage = error instanceof Error ? error.message : String(error) + const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") + TelemetryService.instance.captureException(apiError) + throw new Error(`Mistral completion error: ${errorMessage}`) + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -196,11 +228,23 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand const { id: model, temperature } = this.getModel() try { - const response = await this.client.chat.complete({ - model, - messages: [{ role: "user", content: prompt }], - temperature, - }) + // Build Mistral SDK RequestOptions + const requestOptions: Parameters[1] = {} + if (options?.abortSignal) { + requestOptions.fetchOptions = { signal: options.abortSignal } + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeoutMs = options.timeoutMs + } + + const response = await this.client.chat.complete( + { + model, + messages: [{ role: "user", content: prompt }], + temperature, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = response.choices?.[0]?.message.content @@ -214,6 +258,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand return content || "" } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "completePrompt") TelemetryService.instance.captureException(apiError) From f6eba43d3992f48e657feea481333ee6976d48f0 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 10:58:21 +0800 Subject: [PATCH 2/6] fix(api): unify timeoutMs:0 handling across gemini/mistral/lite-llm + fix test title --- .../__tests__/gemini-handler.spec.ts | 2 +- src/api/providers/__tests__/gemini.spec.ts | 15 ++++++++++ src/api/providers/__tests__/lite-llm.spec.ts | 6 ++++ src/api/providers/__tests__/mistral.spec.ts | 6 ++-- src/api/providers/gemini.ts | 9 ++++-- src/api/providers/lite-llm.ts | 13 ++++---- src/api/providers/mistral.ts | 8 +++-- .../utils/__tests__/request-timeout.spec.ts | 30 +++++++++++++++++++ src/api/providers/utils/request-timeout.ts | 10 +++++++ 9 files changed, 85 insertions(+), 14 deletions(-) create mode 100644 src/api/providers/utils/__tests__/request-timeout.spec.ts create mode 100644 src/api/providers/utils/request-timeout.ts diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 364f62e23d..232849a091 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -55,7 +55,7 @@ describe("GeminiHandler backend support", () => { expect(promptConfig.tools).toBeUndefined() }) - it("completePrompt should pass abort signal through to client via httpOptions", async () => { + it("completePrompt should pass abort signal through to client via config.abortSignal", async () => { const options = { apiProvider: "gemini", enableUrlContext: false, diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 59968060dc..6350f3d395 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -415,6 +415,21 @@ describe("GeminiHandler", () => { }, }) }) + + it("should omit httpOptions entirely for timeoutMs=0 (0 disables the timeout)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: undefined, + temperature: 1, + }, + }) + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index badab1e1a7..324c532735 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1272,6 +1272,12 @@ describe("LiteLLMHandler", () => { const result = await handler.completePrompt("test prompt") expect(result).toBe("response") }) + + it("should omit the timeout option for timeoutMs=0 (0 would abort immediately in the OpenAI SDK)", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) + }) }) describe("createMessage abort signal (bridging)", () => { diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index aff33b6030..2fb0c43f22 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -530,14 +530,12 @@ describe("MistralHandler", () => { }) }) - it("should still forward timeoutMs=0 (uses !== undefined check, not truthy check)", async () => { + it("should omit the timeout option for timeoutMs=0 (0 disables the timeout)", async () => { mockComplete.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }], }) await handler.completePrompt("test prompt", { timeoutMs: 0 }) - expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { - timeoutMs: 0, - }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) }) diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index bc6bcfe2f1..434cb06a27 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -27,6 +27,7 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, Complete import { BaseProvider } from "./base-provider" import { NOT_PROVIDED } from "./constants" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" +import { getRequestTimeoutMs } from "./utils/request-timeout" type GeminiHandlerOptions = ApiHandlerOptions & { isVertex?: boolean @@ -618,8 +619,12 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl ? (this.options.modelTemperature ?? info.defaultTemperature ?? 1) : info.defaultTemperature const httpOpts: { timeout?: number; baseUrl?: string } = {} - if (options?.timeoutMs !== undefined) { - httpOpts.timeout = options.timeoutMs + // Per the abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely (some SDKs treat 0 as an + // immediate timeout). + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + httpOpts.timeout = timeoutMs } if (this.options.googleGeminiBaseUrl) { httpOpts.baseUrl = this.options.googleGeminiBaseUrl diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index da82dd19cb..af5152ef70 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -16,6 +16,7 @@ import { sanitizeOpenAiCallId } from "../../utils/tool-id" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { getRequestTimeoutMs } from "./utils/request-timeout" /** * LiteLLM provider handler @@ -376,15 +377,17 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestOptions.max_tokens = info.maxTokens } - // Build request options with abortSignal and/or timeout. The OpenAI SDK - // treats a timeout of 0 as an immediate timeout, so non-positive timeoutMs - // values disable the timeout instead of being forwarded. + // Build request options with abortSignal and/or timeout. Per the + // abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely, because the OpenAI SDK treats + // a timeout of 0 as an immediate timeout. const createOptions: OpenAI.RequestOptions = {} if (options?.abortSignal) { createOptions.signal = options.abortSignal } - if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { - createOptions.timeout = options.timeoutMs + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + createOptions.timeout = timeoutMs } const response = await this.client.chat.completions.create( diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index 9c304f7eb0..80748fda9f 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -16,6 +16,7 @@ import { ApiHandlerOptions } from "../../shared/api" import { convertToMistralMessages } from "../transform/mistral-format" import { ApiStream } from "../transform/stream" import { handleProviderError } from "./utils/error-handler" +import { getRequestTimeoutMs } from "./utils/request-timeout" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -233,8 +234,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand if (options?.abortSignal) { requestOptions.fetchOptions = { signal: options.abortSignal } } - if (options?.timeoutMs !== undefined) { - requestOptions.timeoutMs = options.timeoutMs + // Per the abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely. + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + requestOptions.timeoutMs = timeoutMs } const response = await this.client.chat.complete( diff --git a/src/api/providers/utils/__tests__/request-timeout.spec.ts b/src/api/providers/utils/__tests__/request-timeout.spec.ts new file mode 100644 index 0000000000..5e969a1036 --- /dev/null +++ b/src/api/providers/utils/__tests__/request-timeout.spec.ts @@ -0,0 +1,30 @@ +import { getRequestTimeoutMs } from "../request-timeout" + +describe("getRequestTimeoutMs", () => { + it("forwards positive timeout values unchanged", () => { + expect(getRequestTimeoutMs(5000)).toBe(5000) + expect(getRequestTimeoutMs(1)).toBe(1) + expect(getRequestTimeoutMs(1234)).toBe(1234) + }) + + it("returns undefined for zero (timeout disabled, not an immediate abort)", () => { + expect(getRequestTimeoutMs(0)).toBeUndefined() + }) + + it("returns undefined for negative values", () => { + expect(getRequestTimeoutMs(-1)).toBeUndefined() + expect(getRequestTimeoutMs(-5000)).toBeUndefined() + }) + + it("returns undefined when no value is provided", () => { + expect(getRequestTimeoutMs()).toBeUndefined() + expect(getRequestTimeoutMs(undefined)).toBeUndefined() + }) + + it("guards against non-number input at the runtime boundary", () => { + expect(getRequestTimeoutMs(NaN)).toBeUndefined() + // Non-number values can only reach this helper through untyped callers + // (e.g. user settings); the double cast exercises the typeof guard. + expect(getRequestTimeoutMs("5000" as unknown as number)).toBeUndefined() + }) +}) diff --git a/src/api/providers/utils/request-timeout.ts b/src/api/providers/utils/request-timeout.ts new file mode 100644 index 0000000000..a3d1e56d4c --- /dev/null +++ b/src/api/providers/utils/request-timeout.ts @@ -0,0 +1,10 @@ +/** + * Returns the value to pass as a client/SDK request timeout option, or undefined. + * + * Per the abort-signal series contract, timeoutMs <= 0 (or undefined) means + * 'no per-request timeout': the option is omitted entirely, because some SDKs + * (e.g. the OpenAI Node SDK) treat timeout: 0 as an IMMEDIATE timeout. + */ +export function getRequestTimeoutMs(timeoutMs?: number): number | undefined { + return typeof timeoutMs === "number" && timeoutMs > 0 ? timeoutMs : undefined +} From 4a1307248c2e8ace2d83bcc5a7558f75d58066fb Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 19:32:03 +0800 Subject: [PATCH 3/6] test(api): close changed-line coverage gaps in gemini, mistral, lite-llm --- src/api/providers/__tests__/gemini.spec.ts | 13 +++++ src/api/providers/__tests__/lite-llm.spec.ts | 13 +++++ src/api/providers/__tests__/mistral.spec.ts | 56 +++++++++++++++++++- 3 files changed, 81 insertions(+), 1 deletion(-) diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 6350f3d395..f3b497c73f 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -430,6 +430,19 @@ describe("GeminiHandler", () => { }, }) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + const controller = new AbortController() + controller.abort() + vi.mocked(handler["client"].models.generateContent).mockRejectedValue(new Error("Gemini API error")) + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("Gemini completion aborted") + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index 324c532735..ee701b1134 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1278,6 +1278,19 @@ describe("LiteLLMHandler", () => { await handler.completePrompt("test prompt", { timeoutMs: 0 }) expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + mockCreate.mockRejectedValueOnce(new Error("LiteLLM API error")) + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("LiteLLM completion aborted") + }) }) describe("createMessage abort signal (bridging)", () => { diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index 2fb0c43f22..a5bfc75fd0 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -53,7 +53,12 @@ import type OpenAI from "openai" import { MistralHandler } from "../mistral" import type { ApiHandlerOptions } from "../../../shared/api" import type { ApiHandlerCreateMessageMetadata } from "../../index" -import type { ApiStreamTextChunk, ApiStreamReasoningChunk, ApiStreamToolCallPartialChunk } from "../../transform/stream" +import type { + ApiStreamTextChunk, + ApiStreamReasoningChunk, + ApiStreamToolCallPartialChunk, + ApiStreamUsageChunk, +} from "../../transform/stream" import { makeCreateMessageMetadata } from "../../../test-utils/api" describe("MistralHandler", () => { @@ -234,6 +239,42 @@ describe("MistralHandler", () => { expect(results[1]).toEqual({ type: "reasoning", text: "Some reasoning" }) expect(results[2]).toEqual({ type: "text", text: "Second text" }) }) + + it("should yield a usage chunk when the stream event carries usage data", async () => { + // The final event carries usage without any delta content; the handler + // must translate it into a usage chunk with the reported token counts. + mockCreate.mockImplementationOnce(async (_options) => + asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { content: "Test response" }, + index: 0, + }, + ], + }, + }, + { + data: { + choices: [], + usage: { promptTokens: 12, completionTokens: 34 }, + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: (ApiStreamTextChunk | ApiStreamUsageChunk)[] = [] + + for await (const chunk of iterator) { + results.push(chunk as ApiStreamTextChunk | ApiStreamUsageChunk) + } + + expect(results).toHaveLength(2) + expect(results[0]).toEqual({ type: "text", text: "Test response" }) + expect(results[1]).toEqual({ type: "usage", inputTokens: 12, outputTokens: 34 }) + }) }) describe("native tool calling", () => { @@ -537,6 +578,19 @@ describe("MistralHandler", () => { await handler.completePrompt("test prompt", { timeoutMs: 0 }) expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + mockComplete.mockRejectedValueOnce(new Error("API Error")) + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("Mistral completion aborted") + }) }) describe("createMessage abort signal bridging", () => { From 73af57e9cdce07943100fa3d68a6f9ca44e26954 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 3 Sep 2026 04:08:15 +0800 Subject: [PATCH 4/6] test(api): use provider identifiers in gemini-handler spec Replace raw gemini apiProvider literals in the two abort-signal spec cases with providerIdentifiers.gemini, matching the rest of the file and the zoo/no-raw-provider-identifiers rule that CI lint enforces. --- src/api/providers/__tests__/gemini-handler.spec.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 62de561492..d7594a01c5 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -58,7 +58,7 @@ describe("GeminiHandler backend support", () => { it("completePrompt should pass abort signal through to client via config.abortSignal", async () => { const options = { - apiProvider: "gemini", + apiProvider: providerIdentifiers.gemini, enableUrlContext: false, enableGrounding: false, } as ApiHandlerOptions @@ -81,7 +81,7 @@ describe("GeminiHandler backend support", () => { it("completePrompt should work without options (backward compatible)", async () => { const options = { - apiProvider: "gemini", + apiProvider: providerIdentifiers.gemini, enableUrlContext: false, enableGrounding: false, } as ApiHandlerOptions From fefffcd4092e1f7fdefc62810293484cccfc8293 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 3 Sep 2026 06:41:23 +0800 Subject: [PATCH 5/6] fix(api): harden gemini/mistral abort, timeout and base-url handling Address CodeRabbit findings on the abort-signal series: - gemini: reject non-HTTPS (non-loopback) googleGeminiBaseUrl before requests so API keys are never sent over cleartext (CWE-319) - mistral: route completePrompt timeout through mergeAbortSignalAndTimeout so the timeout actually cancels the request, instead of a dead timeoutMs field - gemini/mistral/lite-llm specs: request-local signal identity assertions, readiness barrier instead of fixed delay, makeApiHandlerOptions over casts --- .../__tests__/gemini-handler.spec.ts | 89 ++++------------ src/api/providers/__tests__/gemini.spec.ts | 100 ++++++++++++++++++ src/api/providers/__tests__/lite-llm.spec.ts | 9 +- src/api/providers/__tests__/mistral.spec.ts | 46 ++++++-- src/api/providers/gemini.ts | 54 ++++++++++ src/api/providers/mistral.ts | 15 ++- 6 files changed, 228 insertions(+), 85 deletions(-) diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index d7594a01c5..5b1b6c91c4 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -12,8 +12,7 @@ vi.mock("@roo-code/telemetry", () => ({ })) import { GeminiHandler } from "../gemini" -import type { ApiHandlerOptions } from "../../../shared/api" -import { providerIdentifiers } from "@roo-code/types/provider-identifiers" +import { makeApiHandlerOptions } from "../../../test-utils/api" describe("GeminiHandler backend support", () => { beforeEach(() => { @@ -24,11 +23,7 @@ describe("GeminiHandler backend support", () => { // URL context and grounding are mutually exclusive with function declarations // in Gemini API, so createMessage only uses function declarations. // URL context/grounding are only added in completePrompt. - const options = { - apiProvider: providerIdentifiers.gemini, - enableUrlContext: true, - enableGrounding: true, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -41,11 +36,7 @@ describe("GeminiHandler backend support", () => { }) it("completePrompt passes config overrides without tools when URL context and grounding disabled", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableUrlContext: false, - enableGrounding: false, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockResolvedValue({ text: "ok" }) // @ts-ignore access private client @@ -57,11 +48,7 @@ describe("GeminiHandler backend support", () => { }) it("completePrompt should pass abort signal through to client via config.abortSignal", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableUrlContext: false, - enableGrounding: false, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const controller = new AbortController() @@ -80,11 +67,7 @@ describe("GeminiHandler backend support", () => { }) it("completePrompt should work without options (backward compatible)", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableUrlContext: false, - enableGrounding: false, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockResolvedValue({ text: "response" }) @@ -96,10 +79,7 @@ describe("GeminiHandler backend support", () => { describe("error scenarios", () => { it("should handle grounding metadata extraction failure gracefully", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableGrounding: true, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const mockStream = async function* () { @@ -131,10 +111,7 @@ describe("GeminiHandler backend support", () => { }) it("should handle malformed grounding metadata", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableGrounding: true, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const mockStream = async function* () { @@ -182,11 +159,7 @@ describe("GeminiHandler backend support", () => { }) it("should handle API errors when tools are enabled", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - enableUrlContext: true, - enableGrounding: true, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const mockError = new Error("API rate limit exceeded") @@ -230,9 +203,7 @@ describe("GeminiHandler backend support", () => { ] it("should ignore allowedFunctionNames because Gemini rejects larger restriction lists", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -251,9 +222,7 @@ describe("GeminiHandler backend support", () => { }) it("should include all tools when allowedFunctionNames is provided", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -274,9 +243,7 @@ describe("GeminiHandler backend support", () => { }) it("should not pass large allowedFunctionNames lists to Gemini", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -305,9 +272,7 @@ describe("GeminiHandler backend support", () => { }) it("should not pass allowedFunctionNames even when history includes tool calls", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -342,9 +307,7 @@ describe("GeminiHandler backend support", () => { }) it("should fall back to tool_choice when allowedFunctionNames is provided", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -365,9 +328,7 @@ describe("GeminiHandler backend support", () => { }) it("should fall back to tool_choice when allowedFunctionNames is empty", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -389,9 +350,7 @@ describe("GeminiHandler backend support", () => { }) it("should not set toolConfig when allowedFunctionNames is undefined and no tool_choice", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -412,9 +371,7 @@ describe("GeminiHandler backend support", () => { describe("Gemini schema compatibility", () => { it("should strip broad JSON Schema metadata from function declarations", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -473,9 +430,7 @@ describe("GeminiHandler backend support", () => { }) it("should collapse composition and type arrays in function declaration schemas", async () => { - const options = { - apiProvider: providerIdentifiers.gemini, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -534,7 +489,7 @@ describe("GeminiHandler backend support", () => { }) it("should deep-merge allOf fragments instead of overwriting earlier properties", async () => { - const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -579,7 +534,7 @@ describe("GeminiHandler backend support", () => { }) it("should resolve $ref entries before dropping $defs", async () => { - const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -628,7 +583,7 @@ describe("GeminiHandler backend support", () => { }) it("should preserve top-level properties and required entries when allOf is also present", async () => { - const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -670,7 +625,7 @@ describe("GeminiHandler backend support", () => { }) it("should stop recursive $ref expansion before the sanitized schema becomes cyclic", async () => { - const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client @@ -722,7 +677,7 @@ describe("GeminiHandler backend support", () => { }) it("should preserve parameter names that collide with stripped schema keywords", async () => { - const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const stub = vi.fn().mockReturnValue((async function* () {})()) // @ts-ignore access private client diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 21a4ab9174..00775f5c97 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -676,6 +676,103 @@ describe("GeminiHandler", () => { }), ) }) + describe("googleGeminiBaseUrl security (CWE-319)", () => { + it("should reject a non-HTTPS non-loopback googleGeminiBaseUrl in completePrompt", async () => { + const insecureHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "http://gemini.example.test", + }) + insecureHandler["client"] = handler["client"] + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await expect(insecureHandler.completePrompt("Test prompt")).rejects.toThrow( + t("common:errors.gemini.generate_complete_prompt", { + error: "Google Gemini base URL must use HTTPS (or a loopback HTTP endpoint for local test proxies)", + }), + ) + expect(handler["client"].models.generateContent).not.toHaveBeenCalled() + }) + + it("should allow a loopback HTTP googleGeminiBaseUrl in completePrompt", async () => { + const loopbackHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "http://127.0.0.1:8080", + }) + loopbackHandler["client"] = handler["client"] + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + const result = await loopbackHandler.completePrompt("Test prompt") + + expect(result).toBe("Response") + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + httpOptions: { + baseUrl: "http://127.0.0.1:8080", + }, + }), + }), + ) + }) + + it("should reject a non-HTTPS non-loopback googleGeminiBaseUrl in createMessage", async () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + const stub = vi.fn().mockReturnValue((async function* () {})()) + const insecureHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "http://insecure.example.com", + }) + insecureHandler["client"] = handler["client"] + handler["client"].models.generateContentStream = stub + + const stream = insecureHandler.createMessage("You are a helpful assistant", messages) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(ApiProviderError) + expect((error as ApiProviderError).message).toBe( + "Google Gemini base URL must use HTTPS (or a loopback HTTP endpoint for local test proxies)", + ) + expect(stub).not.toHaveBeenCalled() + }) + + it("should allow a loopback HTTP googleGeminiBaseUrl in createMessage", async () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + const stub = vi.fn().mockReturnValue((async function* () {})()) + const loopbackHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "http://127.0.0.1:8080", + }) + loopbackHandler["client"] = handler["client"] + handler["client"].models.generateContentStream = stub + + await collectStream(loopbackHandler.createMessage("You are a helpful assistant", messages)) + + const config = stub.mock.calls[0][0].config + expect(config.httpOptions).toEqual({ baseUrl: "http://127.0.0.1:8080" }) + }) + }) }) describe("createMessage abort signal (bridging)", () => { @@ -737,6 +834,9 @@ describe("GeminiHandler", () => { expect(error).toBeInstanceOf(Error) expect((error as Error).name).toBe("AbortError") expect(capturedSignal).toBeDefined() + // The in-flight request must run against a request-local signal, not the + // external one forwarded by reference. + expect(capturedSignal).not.toBe(controller.signal) expect(capturedSignal?.aborted).toBe(true) }) diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index 451f9f9ad8..dd43f68a84 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1324,10 +1324,17 @@ describe("LiteLLMHandler", () => { it("should abort the in-flight stream when the external signal is triggered", async () => { const controller = new AbortController() let capturedSignal: AbortSignal | undefined + // Readiness barrier: resolves once the request-local signal is captured and + // the (mocked) request has started, instead of guessing a fixed delay. + let requestStartedResolve!: () => void + const requestStarted = new Promise((resolve) => { + requestStartedResolve = resolve + }) // The stream is built inside the mock implementation so that capturedSignal // is already set before the abort-aware chunk is created. mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { capturedSignal = options?.signal + requestStartedResolve() const mockStream = asyncStreamFrom([ { choices: [{ delta: { content: "partial" } }], @@ -1352,7 +1359,7 @@ describe("LiteLLMHandler", () => { ) const collector = collectStream(stream).catch((e: unknown) => e) - await new Promise((resolve) => setTimeout(resolve, 10)) + await requestStarted controller.abort() const error = await collector diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index a5bfc75fd0..de7e0c908a 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -549,25 +549,52 @@ describe("MistralHandler", () => { expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) - it("should pass timeout through to client", async () => { + it("should pass a composite abort+timeout signal through to client", async () => { const controller = new AbortController() mockComplete.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }], }) await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) - expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { - fetchOptions: { signal: controller.signal }, - timeoutMs: 5000, - }) + const callArgs = mockComplete.mock.calls[0][1] as { fetchOptions?: { signal?: AbortSignal } } | undefined + expect(callArgs).toBeDefined() + expect(callArgs?.fetchOptions?.signal).toBeInstanceOf(AbortSignal) + // A fresh composite signal — not the external signal forwarded by reference. + expect(callArgs?.fetchOptions?.signal).not.toBe(controller.signal) + expect(callArgs).not.toHaveProperty("timeoutMs") + // The external side is bridged into the composite. + controller.abort() + expect(callArgs?.fetchOptions?.signal?.aborted).toBe(true) }) - it("should pass only timeoutMs when no signal provided", async () => { + it("should pass a timeout signal through to client when no external signal is provided", async () => { mockComplete.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }], }) await handler.completePrompt("test prompt", { timeoutMs: 3000 }) - expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { - timeoutMs: 3000, + const callArgs = mockComplete.mock.calls[0][1] as { fetchOptions?: { signal?: AbortSignal } } | undefined + expect(callArgs).toBeDefined() + expect(callArgs?.fetchOptions?.signal).toBeInstanceOf(AbortSignal) + expect(callArgs).not.toHaveProperty("timeoutMs") + }) + + it("should bridge the per-request timeout into the composite signal", async () => { + let capturedSignal: AbortSignal | undefined + mockComplete.mockImplementationOnce( + (_options: unknown, requestOptions?: { fetchOptions?: { signal?: AbortSignal } }) => { + capturedSignal = requestOptions?.fetchOptions?.signal + return Promise.resolve({ + choices: [{ message: { content: "response" } }], + }) + }, + ) + + await handler.completePrompt("test prompt", { timeoutMs: 200 }) + + expect(capturedSignal).toBeInstanceOf(AbortSignal) + // The timeout side fires on its own: the self-managed AbortSignal.timeout + // aborts the captured signal after the 200ms per-request deadline. + await vi.waitFor(() => { + expect(capturedSignal?.aborted).toBe(true) }) }) @@ -661,6 +688,9 @@ describe("MistralHandler", () => { expect(error).toBeInstanceOf(Error) expect((error as Error).name).toBe("AbortError") expect(capturedSignal).toBeDefined() + // The in-flight stream must run against a request-local signal, not the + // external one forwarded by reference. + expect(capturedSignal).not.toBe(controller.signal) expect(capturedSignal?.aborted).toBe(true) }) diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index 434cb06a27..222fc3bb64 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -173,6 +173,53 @@ function sanitizeSchemaForGemini( return result } +// googleGeminiBaseUrl is user-editable and can reach non-HTTPS values (settings, +// imported profiles). The @google/genai client keeps API-key authentication for +// custom endpoints, so reject cleartext base URLs before any request — with a +// narrow loopback exception for local test proxies. +function isLoopbackUrl(value: string): boolean { + try { + const parsed = new URL(value) + if (parsed.protocol !== "http:" && parsed.protocol !== "https:") { + return false + } + return ( + parsed.hostname === "localhost" || + parsed.hostname === "::1" || + parsed.hostname === "[::1]" || + /^127\./.test(parsed.hostname) + ) + } catch { + return false + } +} + +// Throws an ApiProviderError when baseUrl is not HTTPS (loopback HTTP is the +// narrow exception, for local test proxies). The provider/model/operation +// arguments keep the structured error context consistent with the request-path +// ApiProviderError instances in this file. +function assertSecureGeminiBaseUrl(baseUrl: string, modelId: string, operation: string): void { + let parsed: URL + try { + parsed = new URL(baseUrl) + } catch { + throw new ApiProviderError("Invalid Google Gemini base URL (not a valid URL)", "Gemini", modelId, operation) + } + if (parsed.protocol === "https:") { + return + } + if (parsed.protocol === "http:" && isLoopbackUrl(baseUrl)) { + // Loopback endpoints (localhost/127.x/::1) are allowed for local test proxies. + return + } + throw new ApiProviderError( + "Google Gemini base URL must use HTTPS (or a loopback HTTP endpoint for local test proxies)", + "Gemini", + modelId, + operation, + ) +} + export class GeminiHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions @@ -297,6 +344,12 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl ? (this.options.modelTemperature ?? info.defaultTemperature ?? 1) : info.defaultTemperature + // Reject cleartext (non-loopback) base URLs before building the request so the + // API key is never sent over an insecure endpoint. + if (this.options.googleGeminiBaseUrl) { + assertSecureGeminiBaseUrl(this.options.googleGeminiBaseUrl, model, "createMessage") + } + const config: GenerateContentConfig = { systemInstruction, httpOptions: this.options.googleGeminiBaseUrl ? { baseUrl: this.options.googleGeminiBaseUrl } : undefined, @@ -627,6 +680,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl httpOpts.timeout = timeoutMs } if (this.options.googleGeminiBaseUrl) { + assertSecureGeminiBaseUrl(this.options.googleGeminiBaseUrl, model, "completePrompt") httpOpts.baseUrl = this.options.googleGeminiBaseUrl } diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index 80748fda9f..93e40b54f2 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -16,7 +16,7 @@ import { ApiHandlerOptions } from "../../shared/api" import { convertToMistralMessages } from "../transform/mistral-format" import { ApiStream } from "../transform/stream" import { handleProviderError } from "./utils/error-handler" -import { getRequestTimeoutMs } from "./utils/request-timeout" +import { mergeAbortSignalAndTimeout } from "./utils/abort-signal" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -231,14 +231,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand try { // Build Mistral SDK RequestOptions const requestOptions: Parameters[1] = {} - if (options?.abortSignal) { - requestOptions.fetchOptions = { signal: options.abortSignal } - } - // Per the abort-signal series contract, timeoutMs <= 0 means 'no per-request - // timeout': the option is omitted entirely. - const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) - if (timeoutMs !== undefined) { - requestOptions.timeoutMs = timeoutMs + // Build a single signal that combines the external abort with the per-request + // timeout (timeoutMs <= 0 disables the timeout; see mergeAbortSignalAndTimeout). + const signal = mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs) + if (signal) { + requestOptions.fetchOptions = { signal } } const response = await this.client.chat.complete( From 9b8033a611e8c93716486663c044eb71c00496ca Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 3 Sep 2026 09:45:42 +0800 Subject: [PATCH 6/6] chore: retrigger CodeRabbit review (no-op) The incremental review for the previous head was stuck in a phantom "review finished" state on the CodeRabbit side (the review object never materialized), so this no-op commit moves the head to a fresh sha and forces a new incremental review. No code changes.