diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 4f2ec12295..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 @@ -56,12 +47,39 @@ describe("GeminiHandler backend support", () => { expect(promptConfig.tools).toBeUndefined() }) + it("completePrompt should pass abort signal through to client via config.abortSignal", async () => { + const options = makeApiHandlerOptions() + 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 = makeApiHandlerOptions() + 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 = { - apiProvider: providerIdentifiers.gemini, - enableGrounding: true, - } as ApiHandlerOptions + const options = makeApiHandlerOptions() const handler = new GeminiHandler(options) const mockStream = async function* () { @@ -93,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* () { @@ -144,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") @@ -192,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 @@ -213,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 @@ -236,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 @@ -267,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 @@ -304,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 @@ -327,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 @@ -351,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 @@ -374,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 @@ -435,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 @@ -496,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 @@ -541,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 @@ -590,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 @@ -632,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 @@ -684,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 2f19028eb7..00775f5c97 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -12,15 +12,23 @@ 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 type { ApiHandlerCreateMessageMetadata } from "../../index" 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 let mockGenerateContentStream: ReturnType @@ -381,6 +389,99 @@ 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, + }, + }) + }) + + 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, + }, + }) + }) + + 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", () => { @@ -514,6 +615,242 @@ 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("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)", () => { + 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() + // 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) + }) + + 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 20bd3b1be2..dd43f68a84 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", () => ({ @@ -1239,4 +1240,133 @@ 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") + }) + + 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) + }) + + 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)", () => { + 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 + // 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" } }], + 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 requestStarted + 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..de7e0c908a 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -53,7 +53,13 @@ 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", () => { let handler: MistralHandler @@ -233,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", () => { @@ -447,11 +489,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 +528,177 @@ 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 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 }) + 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 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 }) + 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) + }) + }) + + 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) }), 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", () => { + 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() + // 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) + }) + + 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..222fc3bb64 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 @@ -172,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 @@ -296,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, @@ -344,7 +398,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 +554,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 +568,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } throw error + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -585,14 +671,30 @@ 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 } = {} + // 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) { + assertSecureGeminiBaseUrl(this.options.googleGeminiBaseUrl, model, "completePrompt") + 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 +715,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..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 @@ -246,9 +247,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 +338,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 +377,30 @@ 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. 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 + } + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + createOptions.timeout = 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..93e40b54f2 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 { mergeAbortSignalAndTimeout } from "./utils/abort-signal" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -101,66 +102,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 - 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 } + 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 +229,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] = {} + // 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( + { + model, + messages: [{ role: "user", content: prompt }], + temperature, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = response.choices?.[0]?.message.content @@ -214,6 +259,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) 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 +}