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 0920042586..af4afac8fd 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 @@ -428,6 +436,186 @@ 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 reject immediately with AbortError when the external signal is pre-aborted, before the request is built", async () => { + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("This operation was aborted") + // The fast-fail must run before the client is touched. + expect(handler["client"].models.generateContent).not.toHaveBeenCalled() + }) + + 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 is aborted while the request is in flight", async () => { + const controller = new AbortController() + // The request stays pending until the external signal aborts it; the + // mock rejects with a plain (non-abort) error once the signal has + // aborted, so the catch block must classify it via signal.aborted. + vi.mocked(handler["client"].models.generateContent).mockImplementationOnce(() => { + return new Promise((_resolve, reject) => { + const fail = () => reject(new Error("Gemini API error")) + if (controller.signal.aborted) { + fail() + return + } + controller.signal.addEventListener("abort", fail, { once: true }) + }) + }) + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + // Abort after entry — the pre-abort fast-fail must not apply, and the + // catch block handles the failure. + controller.abort() + const error = await promise.catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("rethrows non-Error completion failures as-is after capturing telemetry", async () => { + // The instanceof branch only covers Error instances; a non-Error rejection + // must be rethrown untouched (wrapping it would lose the original value). + vi.mocked(handler["client"].models.generateContent).mockRejectedValueOnce("Gemini: not an error object") + + const error = await handler.completePrompt("Test prompt").catch((e: unknown) => e) + expect(error).toBe("Gemini: not an error object") + expect(error).not.toBeInstanceOf(Error) + }) + + it("wraps completion failures in a new error before rethrowing", async () => { + // The consumer-facing error must be the provider's wrap (a NEW Error + // built from the localized template), not the raw SDK error instance: + // this is what the instanceof branch is for (and what makes the branch + // observable to the mutation gate). Unit tests cannot assert the i18n + // string itself (i18next is not initialized with resources here). + const mockError = new Error("Gemini completion failure") + vi.mocked(handler["client"].models.generateContent).mockRejectedValueOnce(mockError) + + const error = await handler.completePrompt("Test prompt").catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect(error).not.toBe(mockError) + }) + + it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => { + // The signal flag has not propagated yet, but the SDK rejects with its + // own abort error: isRequestAborted must catch the error-name branch. + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockRejectedValue( + Object.assign(new Error("Request was aborted."), { name: "AbortError" }), + ) + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("should surface the wrapped provider error when the request fails without options", async () => { + vi.mocked(handler["client"].models.generateContent).mockRejectedValue(new Error("Gemini API error")) + + const error = await handler.completePrompt("Test prompt").catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + // The catch must take the non-abort wrapping path (not the abort + // DOMException path) and must not crash reading a missing signal. + // The i18n message itself is asserted by the error telemetry tests. + expect((error as Error).message).not.toContain("Cannot read properties") + }) + + it("should surface the wrapped provider error when the request fails with options but no signal", async () => { + vi.mocked(handler["client"].models.generateContent).mockRejectedValue(new Error("Gemini API error")) + + const error = await handler.completePrompt("Test prompt", { timeoutMs: 0 }).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + // Options exist but no signal: exercises the second optional-chain + // position, which would crash here if the `?.` were removed. + expect((error as Error).message).not.toContain("Cannot read properties") + }) }) describe("getModel", () => { @@ -561,6 +749,275 @@ describe("GeminiHandler", () => { }) }) + describe("completePrompt request options", () => { + it("applies the explicit model temperature when the model supports temperature", async () => { + const temperatureHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + modelTemperature: 0.5, + }) + temperatureHandler["client"] = handler["client"] + const baseInfo = temperatureHandler.getModel() + // Explicit supportsTemperature: true makes the `!== false` check observable — + // with the default model info (undefined) the literal cannot be mutated. + vi.spyOn(temperatureHandler, "getModel").mockReturnValue({ + ...baseInfo, + info: { + ...baseInfo.info, + supportsTemperature: true, + defaultTemperature: 1, + }, + }) + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await temperatureHandler.completePrompt("Test prompt") + + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ temperature: 0.5 }), + }), + ) + }) + + it("ignores the explicit model temperature when the model does not support temperature", async () => { + const noTemperatureHandler = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + modelTemperature: 0.5, + }) + noTemperatureHandler["client"] = handler["client"] + const baseInfo = noTemperatureHandler.getModel() + vi.spyOn(noTemperatureHandler, "getModel").mockReturnValue({ + ...baseInfo, + info: { + ...baseInfo.info, + supportsTemperature: false, + defaultTemperature: 0.3, + }, + }) + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await noTemperatureHandler.completePrompt("Test prompt") + + // supportsTemperature: false forces the fallback to the model default — + // the user's explicit modelTemperature must not be applied. + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ temperature: 0.3 }), + }), + ) + }) + + 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((error as Error).message).toBe("This operation was aborted") + 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() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + 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() + + // Bound the wait so a broken abort bridge fails this test fast (and fails + // the Stryker mutant) instead of hanging until the runner timeout. + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("abort propagation deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + expect(capturedSignal).toBeDefined() + // The in-flight request must run against a request-local merged signal, not + // the external one forwarded by reference. + expect(capturedSignal).not.toBe(controller.signal) + expect(capturedSignal?.aborted).toBe(true) + // The bridge (RequestConfigBuilder.addMergedSignal) uses AbortSignal.any, + // so the external signal must never be managed with manual listeners. + expect(addEventListenerSpy).not.toHaveBeenCalled() + expect(removeEventListenerSpy).not.toHaveBeenCalled() + }) + + it("should set a live request-local 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 + // The request-local controller signal is always present so the provider + // keeps its own abort handle; with no external signal it never aborts. + expect(config.abortSignal).toBeInstanceOf(AbortSignal) + expect(config.abortSignal?.aborted).toBe(false) + }) + + it("should wrap a non-abort stream failure with the i18n message and capture telemetry", async () => { + const mockError = new Error("Gemini stream failure") + handler["client"].models.generateContentStream = vi.fn().mockRejectedValue(mockError) + + const stream = handler.createMessage("You are a helpful assistant", messages, makeCreateMessageMetadata()) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + // The catch must take the non-abort wrapping path, not the abort path + // and not a raw rethrow. i18next is not initialized with resources in + // unit tests (see "wraps stream failures in a new error before + // rethrowing"), so the localized template string is not assertable + // here; pin the wrapping path by its observable shape instead: a NEW + // Error (not the raw SDK instance), a message that is not the raw + // message, and a name that is not "AbortError". A removed wrapper + // fails the instance/message checks; an abort-path misclassification + // fails the name check. + expect(error).not.toBe(mockError) + expect((error as Error).name).not.toBe("AbortError") + expect((error as Error).message).not.toBe("Gemini stream failure") + expect(mockCaptureException).toHaveBeenCalledTimes(1) + expect(mockCaptureException).toHaveBeenCalledWith( + expect.objectContaining({ + message: "Gemini stream failure", + provider: "Gemini", + operation: "createMessage", + }), + ) + }) + + it("should wrap a stream failure when no metadata is provided at all", async () => { + const mockError = new Error("Gemini stream failure") + handler["client"].models.generateContentStream = vi.fn().mockRejectedValue(mockError) + + const stream = handler.createMessage("You are a helpful assistant", messages) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + // No metadata at all: exercises the first optional-chain position, + // which would crash here if the `?.` were removed. Pin the wrapping + // path by its observable shape (see the sibling test above): a NEW + // Error with a non-raw message and a non-AbortError name. + expect(error).not.toBe(mockError) + expect((error as Error).name).not.toBe("AbortError") + expect((error as Error).message).not.toBe("Gemini stream failure") + }) + }) + describe("error telemetry", () => { const mockMessages: Anthropic.Messages.MessageParam[] = [ { @@ -595,6 +1052,46 @@ describe("GeminiHandler", () => { expect(capturedError).toBeInstanceOf(ApiProviderError) }) + it("rethrows non-Error stream failures as-is after capturing telemetry", async () => { + // The instanceof branch only covers Error instances; a non-Error throw + // must be rethrown untouched (wrapping it would lose the original value). + handler["client"].models.generateContentStream = vi.fn().mockReturnValue( + (async function* () { + yield { candidates: [] } + throw "Gemini stream: not an error object" + })(), + ) + + const error = await collectStream(handler.createMessage(systemPrompt, mockMessages)).catch( + (e: unknown) => e, + ) + expect(error).toBe("Gemini stream: not an error object") + expect(error).not.toBeInstanceOf(Error) + // Telemetry still captures the stringified failure. + expect(mockCaptureException).toHaveBeenCalledWith( + expect.objectContaining({ + message: "Gemini stream: not an error object", + operation: "createMessage", + }), + ) + }) + + it("wraps stream failures in a new error before rethrowing", async () => { + // The consumer-facing error must be the provider's wrap (a NEW Error + // built from the localized template), not the raw SDK error instance: + // this is what the instanceof branch is for (and what makes the branch + // observable to the mutation gate). Unit tests cannot assert the i18n + // string itself (i18next is not initialized with resources here). + const mockError = new Error("Gemini stream failure") + handler["client"].models.generateContentStream = vi.fn().mockRejectedValue(mockError) + + const error = await collectStream(handler.createMessage(systemPrompt, mockMessages)).catch( + (e: unknown) => e, + ) + expect(error).toBeInstanceOf(Error) + expect(error).not.toBe(mockError) + }) + it("should capture telemetry on completePrompt error", async () => { const mockError = new Error("Gemini completion error") ;(handler["client"].models.generateContent as any).mockRejectedValue(mockError) @@ -630,4 +1127,159 @@ describe("GeminiHandler", () => { expect(mockCaptureException).toHaveBeenCalled() }) }) + + describe("createMessage responseId capture", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("captures the final response's responseId so api history can store it", async () => { + handler["client"].models.generateContentStream = vi.fn().mockReturnValue( + asyncStreamFrom([ + { + candidates: [ + { + content: { parts: [{ text: "final" }] }, + finishReason: "STOP", + }, + ], + responseId: "resp-123", + }, + ]), + ) + + await collectStream(handler.createMessage("You are a helpful assistant", messages)) + + expect(handler["lastResponseId"]).toBe("resp-123") + }) + }) + + describe("createMessage streaming loop abort defense", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + function makePartsChunk(parts: Array>) { + return { + candidates: [ + { + content: { parts }, + }, + ], + } + } + + function startStream(signal: AbortSignal, generator: AsyncGenerator) { + handler["client"].models.generateContentStream = vi.fn().mockReturnValue(generator) + return handler.createMessage( + "You are a helpful assistant", + messages, + makeCreateMessageMetadata({ abortSignal: signal }), + ) + } + + it("rejects with AbortError instead of yielding text content after a mid-chunk abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([makePartsChunk([{ thought: true, text: "reasoning" }, { text: "content" }])]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("reasoning") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("rejects with AbortError instead of yielding reasoning after a mid-chunk abort (text-first chunk)", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([makePartsChunk([{ text: "content" }, { thought: true, text: "reasoning" }])]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("rejects with AbortError instead of yielding a tool-call name after a mid-chunk abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([makePartsChunk([{ text: "content" }, { functionCall: { name: "f", args: {} } }])]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("rejects with AbortError instead of yielding tool-call arguments after a mid-chunk abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([ + makePartsChunk([ + { functionCall: { name: "f", args: {} } }, + { functionCall: { name: "g", args: {} } }, + ]), + ]), + ) + const name = await stream.next() + expect(name.value?.type).toBe("tool_call_partial") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + }) + + it("does not process or pull chunks beyond the aborted point and surfaces AbortError instead of completing the stream", async () => { + const controller = new AbortController() + let pulls = 0 + const generator = (async function* () { + pulls++ + yield makePartsChunk([{ text: "first chunk" }]) + pulls++ + // Fallback shape (no candidates): its text yield is the first yield + // of the iteration and carries no pre-yield guard, so only the + // top-of-loop break prevents it from leaking. + yield { text: "buffered chunk" } + pulls++ + yield makePartsChunk([{ text: "third chunk" }]) + })() + const stream = startStream(controller.signal, generator) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Gemini request was aborted") + // The for-await mechanism pulls the already-buffered chunk before the + // top-of-loop check runs; the break must ensure it is not processed and + // that no chunk beyond it is pulled. + expect(pulls).toBe(2) + }) + }) }) diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index dd582e5c2c..d651444b6a 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1,10 +1,13 @@ +import OpenAI from "openai" import { Anthropic } from "@anthropic-ai/sdk" import { LiteLLMHandler } from "../lite-llm" +import { getModels } from "../fetchers/modelCache" import { ApiHandlerOptions } from "../../../shared/api" -import { litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types" +import { type ModelRecord, 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", () => ({ @@ -1380,6 +1383,840 @@ describe("LiteLLMHandler", () => { }) }) + 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 reject immediately with AbortError when the signal is pre-aborted, before model discovery", async () => { + const fetchModelSpy = vi.spyOn(handler, "fetchModel").mockResolvedValue({ + id: litellmDefaultModelId, + info: litellmDefaultModelInfo, + }) + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("This operation was aborted") + // The pre-abort fast-fail must run before provider model discovery: + // an already-aborted request must not trigger getModels/refreshModels. + expect(fetchModelSpy).not.toHaveBeenCalled() + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should pass timeout through to client", async () => { + // The SDK call receives the remaining budget, so pin the clock: on a loaded runner + // real elapsed time between the deadline and the SDK call makes it 4999. + vi.useFakeTimers() + try { + 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 }), + ) + } finally { + vi.useRealTimers() + } + }) + + it("should merge signal and timeoutMs together", async () => { + // The SDK call receives the remaining budget, so pin the clock: without + // fake timers instrumentation drift makes the remaining timeout 9999. + vi.useFakeTimers() + try { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + const passed = mockCreate.mock.calls[0][1] as { signal?: AbortSignal; timeout?: number } + // The SDK receives a deadline-bounded signal merged from the caller's + // signal and the absolute deadline — not the caller's raw signal. + expect(passed.signal).toBeDefined() + expect(passed.signal).not.toBe(controller.signal) + expect(passed.timeout).toBe(10000) + // Caller cancellation still propagates through the merged signal. + controller.abort() + expect(passed.signal?.aborted).toBe(true) + } finally { + vi.useRealTimers() + } + }) + + 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("normalizes fractional, non-finite, and overflowing timeoutMs before signaling", async () => { + // AbortSignal.timeout validates its delay as uint32: 1500.5 used to + // throw ERR_OUT_OF_RANGE from the deadline signal, Infinity threw + // outside the completion try, and an overflowing value clamped to a + // 1ms deadline that aborted instantly. All three must now behave. + vi.useFakeTimers() + try { + mockCreate.mockResolvedValue({ choices: [{ message: { content: "response" } }] }) + + // Fractional: rounded up and forwarded whole to the SDK. + const frac = await handler.completePrompt("test prompt", { timeoutMs: 1500.5 }) + expect(frac).toBe("response") + expect(mockCreate).toHaveBeenLastCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 1501 }), + ) + + // Infinity: means "no timeout" - no deadline signal, no option. + const inf = await handler.completePrompt("test prompt", { timeoutMs: Infinity }) + expect(inf).toBe("response") + expect(mockCreate).toHaveBeenLastCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, + ) + + // Overflowing: capped deadline, not an immediate 1ms abort. + const big = await handler.completePrompt("test prompt", { timeoutMs: 2_147_483_648 }) + expect(big).toBe("response") + const passed = mockCreate.mock.calls.at(-1)![1] as { timeout?: number; signal?: AbortSignal } + expect(passed.timeout).toBe(2_147_483_647) + expect(passed.signal?.aborted).toBe(false) + } finally { + vi.useRealTimers() + } + }) + + it("bounds the whole SDK retry sequence with the absolute deadline", async () => { + // Regression: forwarding only the remaining budget as the SDK's + // per-attempt `timeout` left the default two retries enabled, and each + // retry adds a backoff sleep that is not signal-aware — the pinned SDK + // reproduced three attempts and ~1,530ms elapsed for a 30ms timeout. + // The absolute deadline must bound completion AND the retry waits. + const actual = await vi.importActual<{ default: typeof OpenAI }>("openai") + const attemptSignals: Array = [] + const hangingFetch = (_url: string | URL | Request, init?: RequestInit) => + new Promise((_resolve, reject) => { + attemptSignals.push(init?.signal) + const signal = init?.signal + const onAbort = () => + reject(signal?.reason ?? new DOMException("Request was aborted.", "AbortError")) + if (signal?.aborted) { + onAbort() + return + } + signal?.addEventListener("abort", onAbort, { once: true }) + }) + const realClient = new actual.default({ + baseURL: "http://localhost:4000/v1", + apiKey: "test-key", + fetch: hangingFetch as unknown as typeof globalThis.fetch, + }) + Object.assign(handler, { client: realClient }) + + const startedAt = Date.now() + const error = await handler.completePrompt("test prompt", { timeoutMs: 30 }).catch((e: unknown) => e) + const elapsed = Date.now() - startedAt + + expect(error).toBeInstanceOf(Error) + // Deadline expiry is a timeout, not a caller cancellation. + expect((error as Error).name).toBe("TimeoutError") + // The deadline bounds the retry sequence: the first attempt is aborted + // at the deadline and the SDK must not run its default retries past it. + expect(attemptSignals).toHaveLength(1) + // Well inside the old ~1,530ms blow-up, comfortably above the 30ms budget. + expect(elapsed).toBeLessThan(1000) + }) + + it("reports deadline expiry as a timeout, not as caller cancellation, when the SDK aborts at the deadline", async () => { + // The SDK surfaces deadline expiry as an abort-shaped error + // (APIUserAbortError). With no caller abort in play it must surface as + // the distinct TimeoutError, never as "The LiteLLM request was aborted". + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + return new Promise((_resolve, reject) => { + const onAbort = () => + reject(Object.assign(new Error("Request was aborted."), { name: "APIUserAbortError" })) + if (options?.signal?.aborted) { + onAbort() + return + } + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + }) + + const error = await handler.completePrompt("test prompt", { timeoutMs: 30 }).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("TimeoutError") + expect((error as Error).message).toBe("The LiteLLM completion request timed out") + // The SDK call carries a deadline-bounded signal even without a caller + // signal, so in-flight attempts are aborted at the deadline. + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + + it("still reports caller cancellation as AbortError when a deadline is also active", async () => { + const controller = new AbortController() + let requestStartedResolve!: () => void + const requestStarted = new Promise((resolve) => { + requestStartedResolve = resolve + }) + mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { + requestStartedResolve() + return new Promise((_resolve, reject) => { + const onAbort = () => + reject(Object.assign(new Error("Request was aborted."), { name: "APIUserAbortError" })) + if (options?.signal?.aborted) { + onAbort() + return + } + options?.signal?.addEventListener("abort", onAbort, { once: true }) + }) + }) + + const promise = handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 30000 }) + await requestStarted + controller.abort() + const error = await promise.catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("should surface a standard AbortError when the signal is aborted while the request is in flight", async () => { + const controller = new AbortController() + // Readiness barrier: resolves once the SDK call has started, so the + // abort lands while the request is in flight — not during model + // discovery, which now settles via the rejectOnAbort race first. + let requestStartedResolve!: () => void + const requestStarted = new Promise((resolve) => { + requestStartedResolve = resolve + }) + // The request stays pending until the external signal aborts it; the + // mock rejects with a plain (non-abort) error once the signal has + // aborted, so the catch block must classify it via signal.aborted. + mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { + requestStartedResolve() + const signal = options?.signal + return new Promise((_resolve, reject) => { + const fail = () => reject(new Error("LiteLLM API error")) + if (signal?.aborted) { + fail() + return + } + signal?.addEventListener("abort", fail, { once: true }) + }) + }) + const promise = handler.completePrompt("test prompt", { abortSignal: controller.signal }) + // Abort while the (pending) request is in flight — after discovery and + // entry, so the pre-abort fast-fail and the discovery race do not + // settle it first and the catch block handles the SDK rejection. + await requestStarted + controller.abort() + const error = await promise.catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => { + // The signal flag has not propagated yet, but the SDK rejects with its + // own abort error: isRequestAborted must catch the error-name branch. + mockCreate.mockRejectedValueOnce(Object.assign(new Error("Request was aborted."), { name: "AbortError" })) + const controller = new AbortController() + + const error = await handler + .completePrompt("test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("should surface the wrapped provider error when the request fails without options", async () => { + mockCreate.mockRejectedValueOnce(new Error("boom")) + + const error = await handler.completePrompt("test prompt").catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).message).toBe("LiteLLM completion error: boom") + }) + + it("should surface the wrapped provider error when the request fails with options but no signal", async () => { + mockCreate.mockRejectedValueOnce(new Error("boom")) + + const error = await handler.completePrompt("test prompt", { timeoutMs: 0 }).catch((e: unknown) => e) + expect((error as Error).message).toBe("LiteLLM completion error: boom") + }) + }) + + 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 fetchModelSpy = vi.spyOn(handler, "fetchModel").mockResolvedValue({ + id: litellmDefaultModelId, + info: litellmDefaultModelInfo, + }) + 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((error as Error).message).toBe("This operation was aborted") + // The pre-abort fast-fail must run before provider model discovery: + // an already-aborted request must not trigger getModels/refreshModels. + expect(fetchModelSpy).not.toHaveBeenCalled() + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + 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 pendingAbort = new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("aborted", "AbortError")) + if (capturedSignal?.aborted) { + onAbort() + return + } + capturedSignal?.addEventListener("abort", onAbort, { once: true }) + }) + // A correctly aborted request stops pulling at the top-of-loop break + // and never reads this element, so its rejection must be handled + // here to keep it from being reported as an unhandled rejection. + void pendingAbort.catch(() => undefined) + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "partial" } }], + usage: undefined, + }, + pendingAbort, + ]) + 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() + + // Bound the wait so a broken abort bridge fails this test fast (and fails + // the Stryker mutant) instead of hanging until the runner timeout. + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("abort propagation deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + // The streaming bridge (RequestConfigBuilder.addMergedSignal) uses + // AbortSignal.any. The only manual listener on the external signal is + // the transient discovery-race listener (rejectOnAbort); it must be + // detached by the time the request settles — no listener may outlive + // the request. Exactly one registration proves the streaming bridge + // itself adds no manual listeners of its own. + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const registeredListener = addEventListenerSpy.mock.calls[0]?.[1] as EventListener | undefined + expect(typeof registeredListener).toBe("function") + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registeredListener) + }) + + it("should wrap a non-abort stream failure with the i18n-free provider message and no metadata", async () => { + // createMessage awaits create(...).withResponse(), so the failure must + // surface from the withResponse() call, not from create() itself. + mockCreate.mockReturnValueOnce({ + withResponse: vi.fn().mockRejectedValue(new Error("boom")), + }) + + const error = await collectStream(handler.createMessage("system", messages)).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).message).toBe("LiteLLM streaming error: boom") + }) + + it("should wrap a non-abort stream failure when metadata exists without an abort signal", async () => { + mockCreate.mockReturnValueOnce({ + withResponse: vi.fn().mockRejectedValue(new Error("boom")), + }) + + const error = await collectStream( + handler.createMessage("system", messages, makeCreateMessageMetadata()), + ).catch((e: unknown) => e) + expect((error as Error).message).toBe("LiteLLM streaming error: boom") + }) + }) + + describe("model discovery cancellation (rejectOnAbort)", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("settles with AbortError when the external signal aborts during model discovery", async () => { + // The shared single-flight discovery fetch is in flight (a cold-cache + // model fetch that never settles): the request must not wait for it. + const discoveryStarted = new Promise((resolve) => { + vi.mocked(getModels).mockImplementationOnce( + () => + new Promise(() => { + resolve() + }), + ) + }) + const controller = new AbortController() + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const collector = collectStream(stream).catch((e: unknown) => e) + await discoveryStarted + controller.abort() + + // Bound the wait so a broken discovery race fails this test fast + // (and fails the Stryker mutant) instead of hanging. + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("discovery abort deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("preserves the model-fetch error when discovery fails without an abort", async () => { + const boom = new Error( + "Failed to fetch LiteLLM models: No response from server. Check LiteLLM server status and base URL.", + ) + vi.mocked(getModels).mockRejectedValueOnce(boom) + + const error = await collectStream( + handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: new AbortController().signal }), + ), + ).catch((e: unknown) => e) + // Exact identity: a non-abort discovery failure must not be wrapped + // or misclassified as a cancellation. + expect(error).toBe(boom) + }) + + it("preserves the model-fetch error when discovery fails and no metadata is passed", async () => { + const boom = new Error( + "Failed to fetch LiteLLM models: No response from server. Check LiteLLM server status and base URL.", + ) + vi.mocked(getModels).mockRejectedValueOnce(boom) + + const error = await collectStream(handler.createMessage("system", messages)).catch((e: unknown) => e) + // The catch's re-classification must read the signal through the + // optional chain: a bare `metadata.abortSignal` would throw a + // TypeError here instead of surfacing the fetcher's error. + expect(error).toBe(boom) + }) + + it("normalizes a fetcher AbortError-shaped discovery failure to the standard AbortError", async () => { + // A low-level transport can reject with an AbortError-named error that + // is not the standard abort message (the SDK's DOMException shape). + // The catch must normalize it via the isRequestAborted error-name + // branch even though the signal never aborted. + vi.mocked(getModels).mockRejectedValueOnce( + Object.assign(new Error("The operation was aborted"), { name: "AbortError" }), + ) + + const error = await collectStream( + handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: new AbortController().signal }), + ), + ).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("normalizes to AbortError when the signal aborts while discovery is failing", async () => { + // Discovery rejects only once the signal aborts: whichever settles the + // race first (the abort listener or the forwarded failure), the + // surfaced error must be the standard AbortError. + const controller = new AbortController() + const discoveryStarted = new Promise((resolve) => { + vi.mocked(getModels).mockImplementationOnce( + () => + new Promise((_resolve, reject) => { + resolve() + controller.signal.addEventListener( + "abort", + () => reject(new Error("Failed to fetch LiteLLM models: 503 Service Unavailable.")), + { once: true }, + ) + }), + ) + }) + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const collector = collectStream(stream).catch((e: unknown) => e) + await discoveryStarted + controller.abort() + + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("discovery abort deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("does not let one request's abort settle a concurrent request sharing the discovery fetch", async () => { + // One handler: both requests join the same single-flight fetchModel. + let discoveryResolve!: (models: ModelRecord) => void + const discoveryStarted = new Promise((resolve) => { + vi.mocked(getModels).mockImplementationOnce( + () => + new Promise((resolveModels) => { + discoveryResolve = resolveModels + resolve() + }), + ) + }) + const abortedController = new AbortController() + const siblingController = new AbortController() + + const aborted = collectStream( + handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ), + ).catch((e: unknown) => e) + // The sibling request must keep flowing normally once the shared + // discovery settles. + mockCreate.mockImplementationOnce(() => ({ + withResponse: vi.fn().mockResolvedValue({ + data: asyncStreamFrom([{ choices: [{ delta: { content: "sibling" } }], usage: undefined }]), + }), + })) + const sibling = collectStream( + handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: siblingController.signal }), + ), + ).catch((e: unknown) => e) + + await discoveryStarted + abortedController.abort() + + const abortedError = await aborted + expect((abortedError as Error).name).toBe("AbortError") + + // The shared fetch still runs (the discovery promise is still + // pending, which this test owns), so the sibling request cannot + // have settled yet. + let siblingSettled = false + void sibling.then(() => { + siblingSettled = true + }) + // Drain the microtask queue (and one macrotask) before asserting: + // an erroneous early settlement (e.g. the abort leaking into the + // sibling) settles `sibling` in a microtask, so an assertion + // right after attaching the callback could never observe it. + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(siblingSettled).toBe(false) + + // ...and settles normally once the shared discovery resolves. + discoveryResolve({ [litellmDefaultModelId]: litellmDefaultModelInfo }) + const siblingResult = await sibling + expect(siblingResult).toBeInstanceOf(Array) + const siblingStream = siblingResult as Array<{ type: string; text?: string }> + expect(siblingStream.filter((chunk) => chunk.type === "text").map((chunk) => chunk.text)).toEqual([ + "sibling", + ]) + }) + + it("settles completePrompt with AbortError when the signal aborts during model discovery", async () => { + const discoveryStarted = new Promise((resolve) => { + vi.mocked(getModels).mockImplementationOnce( + () => + new Promise(() => { + resolve() + }), + ) + }) + const controller = new AbortController() + + const collector = handler + .completePrompt("test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + await discoveryStarted + controller.abort() + + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("discovery abort deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("settles completePrompt with a timeout error when timeoutMs elapses during model discovery", async () => { + // A per-request timeout bounds discovery as well: a stalled cold-cache + // fetch must not outlive the requested timeout. A timeout is not a + // user cancellation, so it surfaces as a timeout-specific error + // instead of the abort contract error. + vi.mocked(getModels).mockImplementationOnce(() => new Promise(() => {})) + + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("timeout deadline exceeded")), 3000) + handler.completePrompt("test prompt", { timeoutMs: 50 }).then( + () => { + clearTimeout(deadline) + resolve(new Error("completePrompt unexpectedly resolved")) + }, + (e) => { + clearTimeout(deadline) + resolve(e) + }, + ) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("TimeoutError") + expect((error as Error).message).toBe("The LiteLLM model discovery timed out") + }) + + it("passes only the remaining timeout budget to the SDK after discovery", async () => { + // Discovery consumes part of the configured timeout: the SDK call + // must receive the remaining budget, not the full timeoutMs. Pin the clock + // so coverage or Stryker instrumentation drift cannot move the measured budget. + vi.useFakeTimers() + try { + vi.mocked(getModels).mockImplementationOnce( + () => + new Promise((resolve) => { + setTimeout(() => resolve({ [litellmDefaultModelId]: litellmDefaultModelInfo }), 50) + }), + ) + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + + const pending = handler.completePrompt("test prompt", { timeoutMs: 500 }) + await vi.advanceTimersByTimeAsync(50) + const result = await pending + + expect(result).toBe("ok") + const createOptions = mockCreate.mock.calls[0]?.[1] as { timeout?: number } | undefined + expect(createOptions?.timeout).toBe(450) + } finally { + vi.useRealTimers() + } + }) + + it("preserves the model-fetch error when completePrompt's discovery fails without a signal", async () => { + // No abort signal and no timeout: a plain fetcher failure surfaces + // unwrapped (the SDK section's completion-error wrap does not cover + // discovery) and must never be reclassified as a cancellation. + const boom = new Error("Failed to fetch LiteLLM models: No response from server.") + vi.mocked(getModels).mockRejectedValueOnce(boom) + + const error = await handler.completePrompt("test prompt").catch((e: unknown) => e) + expect(error).toBe(boom) + }) + + it("normalizes a fetcher AbortError-shaped discovery failure in completePrompt", async () => { + // As in createMessage: the discovery catch is the only re-classifier + // on this path (the SDK section's catch does not cover discovery), so + // an AbortError-named fetcher failure must normalize to the standard + // AbortError even though the signal never aborted. + vi.mocked(getModels).mockRejectedValueOnce( + Object.assign(new Error("The operation was aborted"), { name: "AbortError" }), + ) + + const error = await handler.completePrompt("test prompt", { timeoutMs: 0 }).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + }) + + describe("createMessage streaming loop abort defense", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + function chunkOf(delta: Record) { + return { choices: [{ delta }] } + } + + function startStream(signal: AbortSignal, generator: AsyncGenerator) { + mockCreate.mockReturnValue({ + withResponse: vi.fn().mockResolvedValue({ data: generator }), + }) + return handler.createMessage("system", messages, makeCreateMessageMetadata({ abortSignal: signal })) + } + + it("rejects with AbortError instead of yielding text content after a mid-chunk abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([chunkOf({ reasoning_content: "reasoning", content: "content" })]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("reasoning") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("rejects with AbortError instead of yielding a tool call after a mid-chunk abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([ + chunkOf({ + content: "content", + tool_calls: [ + { + index: 0, + id: "1", + type: "function", + function: { name: "f", arguments: "{}" }, + }, + ], + }), + ]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + }) + + it("does not process or pull chunks beyond the aborted point and surfaces AbortError instead of completing the stream", async () => { + const controller = new AbortController() + let pulls = 0 + const generator = (async function* () { + pulls++ + yield chunkOf({ content: "first" }) + pulls++ + // Reasoning shape: its reasoning yield is the first yield of the + // iteration and carries no pre-yield guard, so only the top-of-loop + // break prevents it from leaking. + yield chunkOf({ reasoning_content: "buffered" }) + pulls++ + yield chunkOf({ content: "third" }) + })() + const stream = startStream(controller.signal, generator) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The LiteLLM request was aborted") + // The for-await mechanism pulls the already-buffered chunk before the + // top-of-loop check runs; the break must ensure it is not processed and + // that no chunk beyond it is pulled. + expect(pulls).toBe(2) + }) + + it("completes normally for a usage-only chunk with empty choices (delta is undefined)", async () => { + const controller = new AbortController() + // Real OpenAI-shaped streams end with a usage-only chunk whose choices + // array is empty: delta is undefined, so the chunk must be skipped + // without touching the optional-chained delta fields. + const generator = asyncStreamFrom([ + chunkOf({ content: "content" }), + { choices: [], usage: { prompt_tokens: 5, completion_tokens: 7 } }, + ]) + const stream = startStream(controller.signal, generator) + const first = await stream.next() + expect(first.value?.type).toBe("text") + + const chunks: unknown[] = [] + for await (const c of stream) chunks.push(c) + expect(chunks).toHaveLength(1) + // Assert the mapped usage values, not just the type: this branch maps + // prompt_tokens/completion_tokens onto the usage chunk (L365-366). + expect(chunks[0]).toMatchObject({ type: "usage", inputTokens: 5, outputTokens: 7 }) + }) + }) + describe("model catalog", () => { it("reports its cache scope", () => { const handler = new LiteLLMHandler(mockOptions) diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index 10c2d7ed98..bf97a397ba 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -11,7 +11,7 @@ vi.mock("@roo-code/telemetry", () => ({ // Mock Mistral client - must come before other imports const mockCreate = vi.fn() const mockComplete = vi.fn() -import { asyncStreamFrom } from "../../../test-utils/stream" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" vi.mock("@mistralai/mistralai", () => { return { Mistral: vi.fn().mockImplementation(function () { @@ -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 @@ -129,6 +135,7 @@ describe("MistralHandler", () => { tools: expect.any(Array), toolChoice: "any", }), + expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }), ) expect(result.value).toBeDefined() @@ -152,6 +159,14 @@ describe("MistralHandler", () => { it("should handle errors gracefully", async () => { mockCreate.mockRejectedValueOnce(new Error("API Error")) await expect(handler.createMessage(systemPrompt, messages).next()).rejects.toThrow("API Error") + expect(mockCaptureException).toHaveBeenCalledTimes(1) + expect(mockCaptureException).toHaveBeenCalledWith( + expect.objectContaining({ + message: "API Error", + provider: "Mistral", + operation: "createMessage", + }), + ) }) it("should handle thinking content as reasoning chunks", async () => { @@ -233,6 +248,209 @@ describe("MistralHandler", () => { expect(results[1]).toEqual({ type: "reasoning", text: "Some reasoning" }) expect(results[2]).toEqual({ type: "text", text: "Second text" }) }) + + it("should ignore non-string, non-array delta content", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([{ data: { choices: [{ delta: { content: 42 }, index: 0 }] } }]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(0) + }) + + it("should ignore a thinking chunk with no thinking payload", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + data: { + choices: [{ delta: { content: [{ type: "thinking", thinking: undefined }] }, index: 0 }], + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(0) + }) + + it("should ignore non-text thinking parts", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { + content: [ + { + type: "thinking", + thinking: [{ type: "reference", text: "A reference" }], + }, + ], + }, + index: 0, + }, + ], + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(0) + }) + + it("should yield text for a text chunk that also carries a stray thinking payload", async () => { + // Dispatch is on chunk.type, not on payload presence: a text chunk with a + // stray thinking array must be emitted as text, never as reasoning. + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { + content: [ + { + type: "text", + text: "hello", + thinking: [{ type: "text", text: "reason" }], + }, + ], + }, + index: 0, + }, + ], + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ type: "text", text: "hello" }) + }) + + it("should ignore a non-text chunk that carries a text payload", async () => { + // Dispatch is on chunk.type, not on payload presence: an unknown chunk + // type with a truthy text payload must not be emitted as text. + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { data: { choices: [{ delta: { content: [{ type: "other", text: "leak" }] }, index: 0 }] } }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(0) + }) + + it("should ignore empty text chunks and unknown chunk types", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { + content: [{ type: "text", text: "" }, { type: "other" }], + }, + index: 0, + }, + ], + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(0) + }) + + it("should yield a tool call partial without function details when only an id is provided", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + data: { + choices: [{ delta: { toolCalls: [{ id: "t1" }] }, index: 0 }], + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: unknown[] = [] + for await (const chunk of iterator) { + results.push(chunk) + } + expect(results).toHaveLength(1) + expect(results[0]).toEqual({ + type: "tool_call_partial", + index: 0, + id: "t1", + name: undefined, + arguments: undefined, + }) + }) + + 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", () => { @@ -284,6 +502,7 @@ describe("MistralHandler", () => { ]), toolChoice: "any", }), + expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }), ) }) @@ -301,6 +520,7 @@ describe("MistralHandler", () => { tools: expect.any(Array), toolChoice: "any", }), + expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }), ) }) @@ -438,6 +658,7 @@ describe("MistralHandler", () => { expect.objectContaining({ toolChoice: "any", }), + expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }), ) }) }) @@ -447,11 +668,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 +707,386 @@ describe("MistralHandler", () => { mockComplete.mockRejectedValueOnce(new Error("API Error")) await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Mistral completion error: API Error") }) + + it("should wrap a non-abort error when options are provided without an abort signal", async () => { + // options is a defined object lacking abortSignal: the optional-chaining + // guard must not crash reading a missing abortSignal before wrapping the + // provider error. + mockComplete.mockRejectedValueOnce(new Error("boom")) + await expect(handler.completePrompt("Test prompt", { timeoutMs: 5000 })).rejects.toThrow( + "Mistral completion error: boom", + ) + }) + + 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 reject immediately with AbortError when the external signal is pre-aborted, before the request is built", async () => { + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("This operation was aborted") + // The fast-fail must run before the client is touched. + expect(mockComplete).not.toHaveBeenCalled() + }) + + 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 is aborted while the request is in flight", async () => { + const controller = new AbortController() + // The request stays pending until the external signal aborts it; the + // mock rejects with a plain (non-abort) error once the signal has + // aborted, so the catch block must classify it via signal.aborted. + mockComplete.mockImplementationOnce(() => { + return new Promise((_resolve, reject) => { + const fail = () => reject(new Error("API Error")) + if (controller.signal.aborted) { + fail() + return + } + controller.signal.addEventListener("abort", fail, { once: true }) + }) + }) + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + // Abort after entry — the pre-abort fast-fail must not apply, and the + // catch block handles the failure. + controller.abort() + const error = await promise.catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + }) + + it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => { + // The signal flag has not propagated yet, but the SDK rejects with its + // own abort error: isRequestAborted must catch the error-name branch. + mockComplete.mockRejectedValueOnce(Object.assign(new Error("Request was aborted."), { name: "AbortError" })) + const controller = new AbortController() + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was 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((error as Error).message).toBe("This operation was aborted") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + const addEventListenerSpy = vi.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vi.spyOn(controller.signal, "removeEventListener") + 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() + + // Bound the wait so a broken abort bridge fails this test fast (and fails + // the Stryker mutant) instead of hanging until the runner timeout. + const error = await new Promise((resolve) => { + const deadline = setTimeout(() => resolve(new Error("abort propagation deadline exceeded")), 3000) + collector.then((result) => { + clearTimeout(deadline) + resolve(result) + }) + }) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + expect(capturedSignal).toBeDefined() + // The in-flight stream must run against a request-local merged signal, not + // the external one forwarded by reference. + expect(capturedSignal).not.toBe(controller.signal) + expect(capturedSignal?.aborted).toBe(true) + // The bridge (RequestConfigBuilder.addMergedSignal) uses AbortSignal.any, + // so the external signal must never be managed with manual listeners. + expect(addEventListenerSpy).not.toHaveBeenCalled() + expect(removeEventListenerSpy).not.toHaveBeenCalled() + }) + + it("should wrap a non-abort stream failure when metadata is provided without a signal", async () => { + mockCreate.mockRejectedValueOnce(new Error("boom")) + + const error = await collectStream( + handler.createMessage(systemPrompt, messages, makeCreateMessageMetadata()), + ).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).message).toBe("Mistral completion error: boom") + }) + + it("should pass a live request-local 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] + // The request-local controller signal is always present so the provider + // keeps its own abort handle; with no external signal it never aborts. + expect(streamOptions?.fetchOptions?.signal).toBeInstanceOf(AbortSignal) + expect(streamOptions?.fetchOptions?.signal?.aborted).toBe(false) + }) + }) + + describe("createMessage streaming loop abort defense", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: [{ type: "text", text: "Hello!" }], + }, + ] + + const thinkingChunk = { type: "thinking", thinking: [{ type: "text", text: "reasoning" }] } + const textChunk = { type: "text", text: "content" } + + function eventOf(delta: Record, usage?: Record) { + return { + data: { + choices: [ + { + delta, + index: 0, + }, + ], + ...(usage ? { usage } : {}), + }, + } + } + + function startStream(signal: AbortSignal, generator: AsyncGenerator) { + mockCreate.mockImplementationOnce(async () => generator) + return handler.createMessage(systemPrompt, messages, makeCreateMessageMetadata({ abortSignal: signal })) + } + + it("rejects with AbortError instead of yielding text content after a mid-event abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([eventOf({ content: [thinkingChunk, textChunk] })]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("reasoning") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + }) + + it("rejects with AbortError instead of yielding reasoning after a mid-event abort (text-first event)", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([eventOf({ content: [textChunk, thinkingChunk] })]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + }) + + it("rejects with AbortError instead of yielding a tool call after a mid-event abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([ + eventOf({ + content: [thinkingChunk], + toolCalls: [{ id: "1", function: { name: "f", arguments: "{}" } }], + }), + ]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("reasoning") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + }) + + it("rejects with AbortError instead of yielding usage after a mid-event abort", async () => { + const controller = new AbortController() + const stream = startStream( + controller.signal, + asyncStreamFrom([eventOf({ content: [thinkingChunk] }, { promptTokens: 1, completionTokens: 2 })]), + ) + const first = await stream.next() + expect(first.value?.type).toBe("reasoning") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + }) + + it("does not process or pull events beyond the aborted point and surfaces AbortError instead of completing the stream", async () => { + const controller = new AbortController() + let pulls = 0 + const generator = (async function* () { + pulls++ + yield eventOf({ content: [textChunk] }) + pulls++ + // String-content shape: its text yield is the first yield of the + // iteration and carries no pre-yield guard, so only the top-of-loop + // break prevents it from leaking. + yield eventOf({ content: "buffered" }) + pulls++ + yield eventOf({ content: [textChunk] }) + })() + const stream = startStream(controller.signal, generator) + const first = await stream.next() + expect(first.value?.type).toBe("text") + controller.abort() + + const error = await stream.next().catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("The Mistral request was aborted") + // The for-await mechanism pulls the already-buffered event before the + // top-of-loop check runs; the break must ensure it is not processed and + // that no event beyond it is pulled. + expect(pulls).toBe(2) + }) }) }) diff --git a/src/api/providers/__tests__/vertex.spec.ts b/src/api/providers/__tests__/vertex.spec.ts index cdc038cbab..32be5f6bfe 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 2adf2b5416..47717249e2 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -27,6 +27,9 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, Complete import { BaseProvider } from "./base-provider" import { NOT_PROVIDED } from "./constants" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" +import { createAbortError, isRequestAborted, throwIfAborted } from "./utils/abort-signal" +import { RequestConfigBuilder } from "./config-builder/request-config-builder" +import { getRequestTimeoutMs } from "./utils/request-timeout" type GeminiHandlerOptions = ApiHandlerOptions & { isVertex?: boolean @@ -344,7 +347,22 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } - const params: GenerateContentParameters = { model, contents, config } + // Fast-fail if the request was already aborted before building. + throwIfAborted(metadata?.abortSignal) + + // The request-local controller is the provider-owned abort handle; the + // signal the SDK receives merges it with the external Task signal + // (AbortSignal.any inside RequestConfigBuilder), so external aborts + // cancel the in-flight request without manual listener management. + const requestBuilder = new RequestConfigBuilder<{ signal?: AbortSignal }>() + requestBuilder.addMergedSignal(new AbortController(), metadata) + const requestSignal = requestBuilder.getOption("signal") + + const params: GenerateContentParameters = { + model, + contents, + config: requestSignal ? { ...config, abortSignal: requestSignal } : config, + } try { const result = await this.client.models.generateContentStream(params) @@ -360,6 +378,12 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl let toolCallCounter = 0 for await (const chunk of result) { + // Stop consuming buffered chunks once the request is aborted (the + // SDK iterator may keep delivering buffered content after abort). + if (requestSignal?.aborted) { + break + } + // Track the final structured response (per SDK pattern: candidate.finishReason) if (chunk.candidates && chunk.candidates[0]?.finishReason) { finalResponse = chunk as { responseId?: string } @@ -391,6 +415,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl if (part.thought) { // This is a thinking/reasoning part if (part.text) { + // Re-check before emitting: the consumer may have aborted + // while processing a previously yielded part of this chunk. + throwIfAborted(requestSignal) yield { type: "reasoning", text: part.text } } } else if (part.functionCall) { @@ -400,6 +427,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl const args = JSON.stringify(part.functionCall.args) // Emit name first + throwIfAborted(requestSignal) yield { type: "tool_call_partial", index: toolCallCounter, @@ -409,6 +437,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } // Then emit arguments + throwIfAborted(requestSignal) yield { type: "tool_call_partial", index: toolCallCounter, @@ -421,6 +450,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } else { // This is regular content if (part.text) { + throwIfAborted(requestSignal) yield { type: "text", text: part.text } } } @@ -429,6 +459,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } // Fallback to the original text property if no candidates structure + // (no pre-yield guard here: this is the only yield of a fallback + // chunk, and the top-of-loop check runs without a suspension point + // before it). else if (chunk.text) { yield { type: "text", text: chunk.text } } @@ -438,6 +471,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } + // An aborted request must surface as an AbortError, not as a normal + // stream completion: the top-of-loop break (or a swallowed mid-stream + // abort) ends the loop without throwing otherwise. + throwIfAborted(requestSignal) + if (finalResponse?.responseId) { // Capture responseId so Task.addToApiConversationHistory can store it // alongside the assistant message in api_history.json. @@ -473,6 +511,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } } catch (error) { + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, metadata?.abortSignal)) { + throw createAbortError("Gemini") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") TelemetryService.instance.captureException(apiError) @@ -573,7 +616,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return citationLinks.join(", ") } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Fast-fail if the request was already aborted before building. + throwIfAborted(options?.abortSignal) + const { id: model, info } = this.getModel() try { @@ -581,14 +627,29 @@ 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) { + 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 }] }], @@ -609,6 +670,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return text } catch (error) { + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, options?.abortSignal)) { + throw createAbortError("Gemini") + } 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 6b46694126..ea494e26ef 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -29,6 +29,16 @@ import type { import { RouterProvider } from "./router-provider" import { getLiteLLMModels } from "./fetchers/litellm" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { + createAbortError, + isRequestAborted, + mergeAbortSignalAndTimeout, + mergeAbortSignals, + rejectOnAbort, + throwIfAborted, +} from "./utils/abort-signal" +import { RequestConfigBuilder } from "./config-builder/request-config-builder" +import { getRequestTimeoutMs } from "./utils/request-timeout" /** * LiteLLM provider handler @@ -149,7 +159,33 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { id: modelId, info } = await this.fetchModel() + // Fast-fail if the request was already aborted before building, so an + // already-aborted request fails with AbortError before provider model + // discovery (getModels/refreshModels) begins. + throwIfAborted(metadata?.abortSignal) + + // Model discovery is shared single-flight (RouterProvider.fetchModel): + // concurrent callers join one in-flight fetch, so the fetch itself + // carries no request signal — one caller's abort would reject other + // requests' shared discovery. Cancellation during discovery is a + // rejectOnAbort race: this request settles promptly with AbortError + // while the shared fetch keeps running. + let model: Awaited> + try { + if (metadata?.abortSignal) { + model = await rejectOnAbort(this.fetchModel(), metadata.abortSignal, "LiteLLM") + } else { + model = await this.fetchModel() + } + } catch (error) { + // An abort landing while discovery fails must still surface as the + // standard AbortError, not the fetcher's generic model-fetch error. + if (isRequestAborted(error, metadata?.abortSignal)) { + throw createAbortError("LiteLLM") + } + throw error + } + const { id: modelId, info } = model // Models that require reasoning_content to be echoed back during tool-call // continuations (see LITELLM_PRESERVE_REASONING_MODEL_IDS) need convertToR1Format: @@ -283,29 +319,48 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestHeaders["X-Zoo-Session-ID"] = metadata.taskId } + // The request-local controller is the provider-owned abort handle; the + // request signal merges it with the external Task signal (AbortSignal.any + // inside RequestConfigBuilder), so external aborts cancel the in-flight + // request without manual listener management. + const requestBuilder = new RequestConfigBuilder<{ signal?: AbortSignal }>() + requestBuilder.addMergedSignal(new AbortController(), metadata) + const requestSignal = requestBuilder.getOption("signal") + try { const { data: completion } = await this.client.chat.completions - .create(requestOptions, { headers: requestHeaders }) + .create(requestOptions, { headers: requestHeaders, signal: requestSignal }) .withResponse() let lastUsage for await (const chunk of completion) { + // Stop consuming buffered chunks once the request is aborted (the + // OpenAI SDK iterator may keep delivering buffered content after + // abort, and swallows the mid-stream AbortError). + if (requestSignal?.aborted) { + break + } + const delta = chunk.choices[0]?.delta const usage = chunk.usage as LiteLLMUsage const reasoningText = extractReasoningFromDelta(delta) if (reasoningText) { + // No pre-yield guard: this is the first yield of the iteration — + // the top-of-loop check runs without a suspension point before it. yield { type: "reasoning", text: reasoningText } } if (delta?.content) { + throwIfAborted(requestSignal) yield { type: "text", text: delta.content } } // Handle tool calls in stream - emit partial chunks for NativeToolCallParser if (delta?.tool_calls) { for (const toolCall of delta.tool_calls) { + throwIfAborted(requestSignal) yield { type: "tool_call_partial", index: toolCall.index, @@ -321,6 +376,11 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa } } + // An aborted request must surface as an AbortError, not as a normal + // stream completion: the top-of-loop break (or a swallowed mid-stream + // abort) ends the loop without throwing otherwise. + throwIfAborted(requestSignal) + if (lastUsage) { // Extract cache-related information if available // LiteLLM may use different field names for cache tokens @@ -355,6 +415,11 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa yield usageData } } catch (error) { + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, metadata?.abortSignal)) { + throw createAbortError("LiteLLM") + } if (error instanceof Error) { throw new Error(`LiteLLM streaming error: ${error.message}`) } @@ -362,13 +427,90 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { - const { id: modelId, info } = await this.fetchModel() + /** + * Fresh error for the per-request absolute deadline elapsing in a given + * phase. Distinct from the abort contract error on purpose: a timeout is + * not a user cancellation, and callers treat the two differently. + */ + private createTimeoutError(providerName: string, phase: "model discovery" | "completion request"): Error { + const timeoutError = new Error(`The ${providerName} ${phase} timed out`) + timeoutError.name = "TimeoutError" + return timeoutError + } + + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Fast-fail if the request was already aborted before building, so an + // already-aborted request fails with AbortError before provider model + // discovery (getModels/refreshModels) begins. + const callerAbortSignal = options?.abortSignal + throwIfAborted(callerAbortSignal) + + // A single deadline spans model discovery and the completion request: + // discovery consumes part of the configured timeout, so the SDK call + // must receive only the remaining budget, not the full timeoutMs. + // timeoutMs <= 0 means no per-request timeout and no deadline. + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + const deadline = timeoutMs !== undefined ? Date.now() + timeoutMs : undefined + + // Per the abort-signal series contract, the merged signal (external + // abort + per-request timeout, timeoutMs <= 0 disabling the timeout) + // also bounds model discovery: a stalled cold-cache fetch must settle + // promptly when the caller cancels or the timeout elapses. + const discoverySignal = mergeAbortSignalAndTimeout(options?.abortSignal, timeoutMs) + + // Model discovery is shared single-flight (RouterProvider.fetchModel): + // concurrent callers join one in-flight fetch, so the fetch itself + // carries no request signal — one caller's abort would reject other + // requests' shared discovery. Cancellation during discovery is a + // rejectOnAbort race: this request settles promptly while the shared + // fetch keeps running. + let model: Awaited> + try { + if (discoverySignal) { + model = await rejectOnAbort(this.fetchModel(), discoverySignal, "LiteLLM") + } else { + model = await this.fetchModel() + } + } catch (error) { + // An abort landing while discovery fails must still surface as the + // standard AbortError, not the fetcher's generic model-fetch error. + if (isRequestAborted(error, discoverySignal)) { + if (discoverySignal?.aborted) { + // Our merged signal fired: a caller stop surfaces as the + // standard AbortError, while the per-request timeout + // elapsing during discovery is a timeout — it must not be + // misreported as a user cancellation. + throw callerAbortSignal?.aborted + ? createAbortError("LiteLLM") + : this.createTimeoutError("LiteLLM", "model discovery") + } + // The signal never aborted: the fetcher itself failed with an + // AbortError-shaped error; normalize to the standard AbortError. + throw createAbortError("LiteLLM") + } + throw error + } + const { id: modelId, info } = model // Check if this is a GPT-5 model that requires max_completion_tokens instead of max_tokens const usesMaxCompletionTokens = this.isGpt5(modelId) || info.requiresResponsesApi const reasoningEffort = this.getReasoningEffort(info) + // The absolute deadline is enforced THROUGH the SDK call, not just as a + // per-attempt timeout: the SDK's `timeout` option bounds one attempt, + // and its default two retries each add a backoff sleep that is not + // signal-aware (a 30ms budget otherwise elapsed ~1.5s across three + // attempts). The deadline signal merged with the caller's signal aborts + // in-flight attempts and makes the SDK's next attempt short-circuit at + // its pre-flight `signal.aborted` check, and the rejectOnAbort race + // below settles the caller at the deadline even while the SDK is + // mid-retry-sleep. Retries stay enabled within the remaining budget. + const remainingMs = + deadline !== undefined && Number.isFinite(deadline) + ? Math.min(2_147_483_647, Math.max(1, Math.ceil(deadline - Date.now()))) + : undefined + const deadlineSignal = remainingMs !== undefined ? AbortSignal.timeout(remainingMs) : undefined + try { const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = { model: modelId, @@ -389,9 +531,44 @@ 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 = {} + // mergeAbortSignals returns the caller's signal directly when no deadline + // exists, so one expression covers all four caller/deadline combinations. + const requestSignal = callerAbortSignal + ? mergeAbortSignals(callerAbortSignal, deadlineSignal) + : deadlineSignal + if (requestSignal) { + createOptions.signal = requestSignal + } + if (remainingMs !== undefined) { + // Only the remaining budget reaches the SDK: model discovery has + // already consumed part of the configured timeout. + createOptions.timeout = remainingMs + } + + const request = this.client.chat.completions.create( + requestOptions, + Object.keys(createOptions).length > 0 ? createOptions : undefined, + ) + const response = deadlineSignal ? await rejectOnAbort(request, deadlineSignal, "LiteLLM") : await request return response.choices[0]?.message.content || "" } catch (error) { + // Deadline expiry is distinct from caller cancellation: the SDK reports + // the deadline as an abort-shaped error (APIUserAbortError), so check + // the deadline first — with the caller still running it must surface as + // a timeout, never masquerade as the cancellation AbortError. + if (deadlineSignal?.aborted && !callerAbortSignal?.aborted) { + throw this.createTimeoutError("LiteLLM", "completion request") + } + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, callerAbortSignal)) { + throw createAbortError("LiteLLM") + } 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 e7023513b0..5893b796ae 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -15,6 +15,8 @@ import { ApiHandlerOptions } from "../../shared/api" import { convertToMistralMessages } from "../transform/mistral-format" import { ApiStream } from "../transform/stream" +import { createAbortError, isRequestAborted, mergeAbortSignalAndTimeout, throwIfAborted } from "./utils/abort-signal" +import { RequestConfigBuilder } from "./config-builder/request-config-builder" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -100,66 +102,101 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand // Temporary debug log for QA // console.log("[MISTRAL DEBUG] Raw API request body:", requestOptions) + // Fast-fail if the request was already aborted before building. + throwIfAborted(metadata?.abortSignal) + + // The request-local controller is the provider-owned abort handle; the + // fetch signal merges it with the external Task signal (AbortSignal.any + // inside RequestConfigBuilder), so external aborts cancel the in-flight + // request without manual listener management. + const requestBuilder = new RequestConfigBuilder<{ signal?: AbortSignal }>() + requestBuilder.addMergedSignal(new AbortController(), metadata) + const requestSignal = requestBuilder.getOption("signal") + 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}`) - } + response = await this.client.chat.stream(requestOptions, { + fetchOptions: { signal: requestSignal }, + }) + + for await (const event of response) { + // Stop consuming buffered events once the request is aborted (the + // SDK iterator may keep delivering buffered content after abort). + if (requestSignal?.aborted) { + break + } + + 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 (no pre-yield guard: this is + // the first yield of the iteration — the top-of-loop check + // runs without a suspension point before it). + 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) { + throwIfAborted(requestSignal) + yield { type: "reasoning", text: thinkingPart.text } + } } + } else if (chunk.type === "text" && chunk.text) { + // Handle text content normally + throwIfAborted(requestSignal) + 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] + // 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] + throwIfAborted(requestSignal) + yield { + type: "tool_call_partial", + index: i, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } + } + } + + if (event.data.usage) { + throwIfAborted(requestSignal) yield { - type: "tool_call_partial", - index: i, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + 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, - } + // An aborted request must surface as an AbortError, not as a normal + // stream completion: the top-of-loop break (or a swallowed mid-stream + // abort) ends the loop without throwing otherwise. + throwIfAborted(requestSignal) + } catch (error) { + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, metadata?.abortSignal)) { + throw createAbortError("Mistral") } + 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}`) } } @@ -191,15 +228,30 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand return { id, info, maxTokens, temperature } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Fast-fail if the request was already aborted before building. + throwIfAborted(options?.abortSignal) + 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 @@ -213,6 +265,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand return content || "" } catch (error) { + // Aborted request: covers both the external signal and SDK-native + // abort errors that may surface before the signal flag propagates. + if (isRequestAborted(error, options?.abortSignal)) { + throw createAbortError("Mistral") + } 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__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index f779d48e1f..3cd64e8fd1 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -275,6 +275,56 @@ describe("abort-signal utilities", () => { await vi.waitFor(() => expect(result?.aborted).toBe(true)) }) + + it("treats a non-finite timeoutMs as no timeout", () => { + const controller = new AbortController() + + // Infinity must not reach AbortSignal.timeout(), which throws + // ERR_OUT_OF_RANGE for non-integer delays. + expect(() => mergeAbortSignalAndTimeout(undefined, Infinity)).not.toThrow() + expect(mergeAbortSignalAndTimeout(undefined, Infinity)).toBeUndefined() + expect(mergeAbortSignalAndTimeout(controller.signal, Infinity)).toBe(controller.signal) + }) + + it("rounds a fractional timeoutMs up before scheduling", async () => { + // Node validates AbortSignal.timeout delays as uint32: 1.5 or 0.5 + // would throw ERR_OUT_OF_RANGE. Pin the rounded delay that actually + // reaches AbortSignal.timeout (1.5 -> 2, 0.5 -> 1, never 0) so a + // round-down mutation reddens this test instead of still aborting + // eventually. + const timeoutSpy = vi.spyOn(AbortSignal, "timeout") + let result: AbortSignal | undefined + let tiny: AbortSignal | undefined + try { + expect(() => { + result = mergeAbortSignalAndTimeout(undefined, 1.5) + }).not.toThrow() + expect(timeoutSpy).toHaveBeenCalledWith(2) + expect(() => { + tiny = mergeAbortSignalAndTimeout(undefined, 0.5) + }).not.toThrow() + expect(timeoutSpy).toHaveBeenCalledWith(1) + } finally { + timeoutSpy.mockRestore() + } + expect(result).toBeInstanceOf(AbortSignal) + expect(result?.aborted).toBe(false) + expect(tiny).toBeInstanceOf(AbortSignal) + await vi.waitFor(() => expect(result?.aborted).toBe(true)) + }) + + it("caps an overflowing timeoutMs instead of aborting immediately", async () => { + // Node clamps AbortSignal.timeout delays above 2^31-1 ms to 1 ms; + // the helper caps the delay so an overflowing timeout does not abort + // the request within the first millisecond. + const result = mergeAbortSignalAndTimeout(undefined, 2_147_483_648) + + expect(result).toBeInstanceOf(AbortSignal) + expect(result?.aborted).toBe(false) + + await new Promise((resolve) => setTimeout(resolve, 20)) + expect(result?.aborted).toBe(false) + }) }) describe("mergeAbortSignals", () => { 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..f38a6b039f --- /dev/null +++ b/src/api/providers/utils/__tests__/request-timeout.spec.ts @@ -0,0 +1,38 @@ +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("rounds fractional values up and rejects non-finite values", () => { + // AbortSignal.timeout validates its delay as uint32: fractions must + // round up (never down to 0), and Infinity must mean "no timeout". + expect(getRequestTimeoutMs(1500.5)).toBe(1501) + expect(getRequestTimeoutMs(0.5)).toBe(1) + expect(getRequestTimeoutMs(Infinity)).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/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index cbb7a13aaf..763a97c4c0 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -1,18 +1,24 @@ /** * Merge an optional external abort signal with an optional timeout. * - * Timeout values <= 0 are treated as disabled. The timeout is created via the + * Timeout values <= 0 and non-finite values are treated as disabled; finite + * delays above 2^31-1 ms are capped (Node clamps them to 1 ms otherwise). The timeout is created via the * native AbortSignal.timeout() API, which self-manages its timer — callers do * not need to (and cannot) clear it manually. */ export function mergeAbortSignalAndTimeout(externalSignal?: AbortSignal, timeoutMs?: number): AbortSignal | undefined { - const hasTimeout = typeof timeoutMs === "number" && timeoutMs > 0 + // Node's AbortSignal.timeout() throws ERR_OUT_OF_RANGE for non-integer + // delays (Infinity) and clamps delays above 2^31-1 ms to 1 ms, which would + // abort almost immediately; treat non-finite values as "no timeout", + // round fractional delays up to whole milliseconds (uint32 validation + // throws for 1500.5), and cap the rest. + const hasTimeout = typeof timeoutMs === "number" && Number.isFinite(timeoutMs) && timeoutMs > 0 if (!hasTimeout) { return externalSignal } - const timeoutSignal = AbortSignal.timeout(timeoutMs) + const timeoutSignal = AbortSignal.timeout(Math.min(Math.max(1, Math.ceil(timeoutMs)), 2_147_483_647)) if (!externalSignal) { return timeoutSignal diff --git a/src/api/providers/utils/request-timeout.ts b/src/api/providers/utils/request-timeout.ts new file mode 100644 index 0000000000..6fd1d8b592 --- /dev/null +++ b/src/api/providers/utils/request-timeout.ts @@ -0,0 +1,15 @@ +/** + * Returns the value to pass as a client/SDK request timeout option, or undefined. + * + * Per the abort-signal series contract, timeoutMs <= 0, non-finite values + * (or undefined) mean '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. Finite values are rounded up to whole milliseconds: + * downstream AbortSignal.timeout() validates its delay as uint32 and throws + * for fractions. + */ +export function getRequestTimeoutMs(timeoutMs?: number): number | undefined { + return typeof timeoutMs === "number" && Number.isFinite(timeoutMs) && timeoutMs > 0 + ? Math.ceil(timeoutMs) + : undefined +}