From 2a58311c0b33b27843544491a04e4c535709395a Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 10:11:42 +0800 Subject: [PATCH 1/3] feat(api): abort signal support for gemini, mistral, lite-llm (completePrompt + createMessage) --- .../__tests__/gemini-handler.spec.ts | 38 ++++ src/api/providers/__tests__/gemini.spec.ts | 209 ++++++++++++++++++ src/api/providers/__tests__/lite-llm.spec.ts | 104 +++++++++ src/api/providers/__tests__/mistral.spec.ts | 145 +++++++++++- src/api/providers/__tests__/vertex.spec.ts | 56 +++++ src/api/providers/gemini.ts | 56 ++++- src/api/providers/lite-llm.ts | 54 ++++- src/api/providers/mistral.ts | 151 ++++++++----- 8 files changed, 751 insertions(+), 62 deletions(-) diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 110f60289c..364f62e23d 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -55,6 +55,44 @@ describe("GeminiHandler backend support", () => { expect(promptConfig.tools).toBeUndefined() }) + it("completePrompt should pass abort signal through to client via httpOptions", async () => { + const options = { + apiProvider: "gemini", + enableUrlContext: false, + enableGrounding: false, + } as ApiHandlerOptions + const handler = new GeminiHandler(options) + + const controller = new AbortController() + const stub = vi.fn().mockResolvedValue({ text: "response" }) + handler["client"].models.generateContent = stub + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + + expect(stub).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + abortSignal: controller.signal, + }), + }), + ) + }) + + it("completePrompt should work without options (backward compatible)", async () => { + const options = { + apiProvider: "gemini", + enableUrlContext: false, + enableGrounding: false, + } as ApiHandlerOptions + const handler = new GeminiHandler(options) + + const stub = vi.fn().mockResolvedValue({ text: "response" }) + handler["client"].models.generateContent = stub + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + describe("error scenarios", () => { it("should handle grounding metadata extraction failure gracefully", async () => { const options = { diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 701c453e3e..59968060dc 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -12,14 +12,22 @@ vitest.mock("@roo-code/telemetry", () => ({ import { Anthropic } from "@anthropic-ai/sdk" +import type { GenerateContentResponse } from "@google/genai" + import { type ModelInfo, geminiDefaultModelId, ApiProviderError } from "@roo-code/types" import { t } from "i18next" import { GeminiHandler } from "../gemini" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" const GEMINI_MODEL_NAME = geminiDefaultModelId +// @google/genai's GenerateContentResponse exposes `text` via a getter backed by +// `candidates`, so the stub only carries the field the provider reads; the double +// cast is the least-friction way to satisfy the class type in mocks. +const stubGenerateContentResponse = (text: string) => ({ text }) as unknown as GenerateContentResponse + describe("GeminiHandler", () => { let handler: GeminiHandler @@ -342,6 +350,71 @@ describe("GeminiHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should pass abort signal through to client via config.abortSignal", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + abortSignal: controller.signal, + httpOptions: undefined, + temperature: 1, + }, + }) + }) + + it("should work without options (backward compatible)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: undefined, + temperature: 1, + }, + }) + }) + + it("should pass timeoutMs through to client via httpOptions with abortSignal on config", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + abortSignal: controller.signal, + httpOptions: { timeout: 10000 }, + temperature: 1, + }, + }) + }) + + it("should pass only timeoutMs when no signal is provided", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: { timeout: 5000 }, + temperature: 1, + }, + }) + }) }) describe("getModel", () => { @@ -475,6 +548,142 @@ describe("GeminiHandler", () => { }) }) + describe("completePrompt request options", () => { + it("should pass timeout and baseUrl through httpOptions", async () => { + const handlerWithBaseUrl = new GeminiHandler({ + apiKey: "test-key", + apiModelId: GEMINI_MODEL_NAME, + geminiApiKey: "test-key", + googleGeminiBaseUrl: "https://gemini.example.test", + }) + handlerWithBaseUrl["client"] = handler["client"] + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + const result = await handlerWithBaseUrl.completePrompt("Test prompt", { timeoutMs: 1234 }) + + expect(result).toBe("Response") + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + httpOptions: { + timeout: 1234, + baseUrl: "https://gemini.example.test", + }, + }), + }), + ) + }) + + it("should pass abortSignal on config instead of httpOptions", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + abortSignal: controller.signal, + httpOptions: undefined, + }), + }), + ) + }) + + it("should omit httpOptions when timeoutMs and baseUrl are not provided", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("Response"), + ) + + await handler.completePrompt("Test prompt") + + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + config: expect.objectContaining({ + httpOptions: undefined, + }), + }), + ) + }) + }) + + describe("createMessage abort signal (bridging)", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "You are a helpful assistant", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(handler["client"].models.generateContentStream).not.toHaveBeenCalled() + }) + + it("should abort the in-flight request when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + const stub = vi.fn().mockImplementation(async (params: { config?: { abortSignal?: AbortSignal } }) => { + capturedSignal = params.config?.abortSignal + return (async function* () { + yield { text: "partial" } + if (capturedSignal?.aborted) { + throw new DOMException("aborted", "AbortError") + } + await new Promise((_resolve, reject) => { + capturedSignal?.addEventListener( + "abort", + () => reject(new DOMException("aborted", "AbortError")), + { once: true }, + ) + }) + })() + }) + handler["client"].models.generateContentStream = stub + + const stream = handler.createMessage( + "You are a helpful assistant", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + + it("should not set config.abortSignal when no external signal is provided", async () => { + const stub = vi.fn().mockReturnValue((async function* () {})()) + handler["client"].models.generateContentStream = stub + + await collectStream(handler.createMessage("You are a helpful assistant", messages)) + + const config = stub.mock.calls[0][0].config + expect(config.abortSignal).toBeUndefined() + }) + }) + describe("error telemetry", () => { const mockMessages: Anthropic.Messages.MessageParam[] = [ { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index eee5cf52bb..badab1e1a7 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -6,6 +6,7 @@ import { ApiHandlerOptions } from "../../../shared/api" import { litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" +import { makeCreateMessageMetadata } from "../../../test-utils/api" // Mock vscode first to avoid import errors vi.mock("vscode", () => ({ @@ -1235,4 +1236,107 @@ describe("LiteLLMHandler", () => { expect(requestHeaders).not.toHaveProperty("X-Zoo-Session-ID") }) }) + + describe("completePrompt", () => { + it("should pass abort signal through to client", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + const controller = new AbortController() + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal }), + ) + }) + + it("should pass timeout through to client", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), + ) + }) + + it("should merge signal and timeoutMs together", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ signal: controller.signal, timeout: 10000 }), + ) + }) + + it("should work without options (backward compatible)", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + }) + + describe("createMessage abort signal (bridging)", () => { + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello", + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + // The stream is built inside the mock implementation so that capturedSignal + // is already set before the abort-aware chunk is created. + mockCreate.mockImplementationOnce((_body: unknown, options?: { signal?: AbortSignal }) => { + capturedSignal = options?.signal + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "partial" } }], + usage: undefined, + }, + new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("aborted", "AbortError")) + if (capturedSignal?.aborted) { + onAbort() + return + } + capturedSignal?.addEventListener("abort", onAbort, { once: true }) + }), + ]) + return { withResponse: vi.fn().mockResolvedValue({ data: mockStream }) } + }) + + const stream = handler.createMessage( + "system", + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + }) }) diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index f2a7591bd8..aff33b6030 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -54,6 +54,7 @@ import { MistralHandler } from "../mistral" import type { ApiHandlerOptions } from "../../../shared/api" import type { ApiHandlerCreateMessageMetadata } from "../../index" import type { ApiStreamTextChunk, ApiStreamReasoningChunk, ApiStreamToolCallPartialChunk } from "../../transform/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" describe("MistralHandler", () => { let handler: MistralHandler @@ -447,11 +448,14 @@ describe("MistralHandler", () => { const prompt = "Test prompt" const result = await handler.completePrompt(prompt) - expect(mockComplete).toHaveBeenCalledWith({ - model: mockOptions.apiModelId, - messages: [{ role: "user", content: prompt }], - temperature: 0, - }) + expect(mockComplete).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: prompt }], + temperature: 0, + }, + undefined, + ) expect(result).toBe("Test response") }) @@ -483,5 +487,136 @@ describe("MistralHandler", () => { mockComplete.mockRejectedValueOnce(new Error("API Error")) await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Mistral completion error: API Error") }) + + it("should pass abort signal through to client", async () => { + const controller = new AbortController() + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + fetchOptions: { signal: controller.signal }, + }) + }) + + it("should work without options (backward compatible)", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) + }) + + it("should pass timeout through to client", async () => { + const controller = new AbortController() + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + fetchOptions: { signal: controller.signal }, + timeoutMs: 5000, + }) + }) + + it("should pass only timeoutMs when no signal provided", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 3000 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeoutMs: 3000, + }) + }) + + it("should still forward timeoutMs=0 (uses !== undefined check, not truthy check)", async () => { + mockComplete.mockResolvedValueOnce({ + choices: [{ message: { content: "response" } }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeoutMs: 0, + }) + }) + }) + + describe("createMessage abort signal bridging", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: [{ type: "text", text: "Hello!" }], + }, + ] + + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight stream when the external signal is triggered", async () => { + const controller = new AbortController() + let capturedSignal: AbortSignal | undefined + mockCreate.mockImplementationOnce( + async (_options: unknown, requestOptions?: { fetchOptions?: { signal?: AbortSignal } }) => { + capturedSignal = requestOptions?.fetchOptions?.signal + return asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { content: "partial" }, + index: 0, + }, + ], + }, + }, + new Promise((_resolve, reject) => { + const onAbort = () => reject(new DOMException("aborted", "AbortError")) + if (capturedSignal?.aborted) { + onAbort() + return + } + capturedSignal?.addEventListener("abort", onAbort, { once: true }) + }), + ]) + }, + ) + + const stream = handler.createMessage( + systemPrompt, + messages, + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const collector = collectStream(stream).catch((e: unknown) => e) + await new Promise((resolve) => setTimeout(resolve, 10)) + controller.abort() + + const error = await collector + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(capturedSignal).toBeDefined() + expect(capturedSignal?.aborted).toBe(true) + }) + + it("should not pass a signal to the stream call when no external signal is provided", async () => { + const stream = handler.createMessage(systemPrompt, messages) + await collectStream(stream) + const streamOptions = mockCreate.mock.calls[0][1] + expect(streamOptions).toBeUndefined() + }) }) }) diff --git a/src/api/providers/__tests__/vertex.spec.ts b/src/api/providers/__tests__/vertex.spec.ts index a304518ca7..99c5787e0e 100644 --- a/src/api/providers/__tests__/vertex.spec.ts +++ b/src/api/providers/__tests__/vertex.spec.ts @@ -21,10 +21,19 @@ vitest.mock("@roo-code/telemetry", () => ({ import { Anthropic } from "@anthropic-ai/sdk" +import type { GenerateContentResponse } from "@google/genai" + import { ApiStreamChunk } from "../../transform/stream" import { t } from "i18next" import { VertexHandler } from "../vertex" +import { collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" + +// @google/genai's GenerateContentResponse exposes `text` via a getter backed by +// `candidates`, so the stub only carries the field the provider reads; the double +// cast is the least-friction way to satisfy the class type in mocks. +const stubGenerateContentResponse = (text: string) => ({ text }) as unknown as GenerateContentResponse describe("VertexHandler", () => { let handler: VertexHandler @@ -137,6 +146,53 @@ describe("VertexHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should pass abort signal through to client via config.abortSignal", async () => { + const controller = new AbortController() + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: expect.any(String), + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: expect.objectContaining({ + abortSignal: controller.signal, + httpOptions: undefined, + temperature: 1, + }), + }), + ) + }) + + it("should work without options (backward compatible)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + }) + }) + + describe("createMessage abort signal (inherited from GeminiHandler)", () => { + it("should reject immediately with AbortError when the external signal is pre-aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "You are a helpful assistant", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const error = await collectStream(stream).catch((e: unknown) => e) + expect(error).toBeInstanceOf(Error) + expect((error as Error).name).toBe("AbortError") + expect(handler["client"].models.generateContentStream).not.toHaveBeenCalled() + }) }) describe("getModel", () => { diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index ec0d14e4c9..bc6bcfe2f1 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -344,7 +344,30 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } - const params: GenerateContentParameters = { model, contents, config } + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight generateContentStream request + // can be cancelled. The @google/genai SDK merges this signal with its own + // timeout handling, which is preserved rather than replaced. + // A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("Gemini request aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + + const params: GenerateContentParameters = { + model, + contents, + config: requestAbortController ? { ...config, abortSignal: requestAbortController.signal } : config, + } try { const result = await this.client.models.generateContentStream(params) @@ -477,6 +500,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } } } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("Gemini request aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") TelemetryService.instance.captureException(apiError) @@ -486,6 +514,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl } throw error + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -585,14 +617,25 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl const temperatureConfig: number | undefined = supportsTemperature ? (this.options.modelTemperature ?? info.defaultTemperature ?? 1) : info.defaultTemperature + const httpOpts: { timeout?: number; baseUrl?: string } = {} + if (options?.timeoutMs !== undefined) { + httpOpts.timeout = options.timeoutMs + } + if (this.options.googleGeminiBaseUrl) { + httpOpts.baseUrl = this.options.googleGeminiBaseUrl + } const promptConfig: GenerateContentConfig = { - httpOptions: this.options.googleGeminiBaseUrl - ? { baseUrl: this.options.googleGeminiBaseUrl } - : undefined, + httpOptions: Object.keys(httpOpts).length > 0 ? httpOpts : undefined, temperature: temperatureConfig, } + // @google/genai expects request cancellation on config.abortSignal + // (not httpOptions.signal), so the signal is passed directly to the config. + if (options?.abortSignal) { + promptConfig.abortSignal = options.abortSignal + } + const request = { model, contents: [{ role: "user", parts: [{ text: prompt }] }], @@ -613,6 +656,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return text } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("Gemini completion aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "completePrompt") TelemetryService.instance.captureException(apiError) diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index 8cfe2d0a19..da82dd19cb 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -246,9 +246,31 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestHeaders["X-Zoo-Session-ID"] = metadata.taskId } + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight streaming request can be + // cancelled. A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("LiteLLM streaming aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + try { const { data: completion } = await this.client.chat.completions - .create(requestOptions, { headers: requestHeaders }) + .create( + requestOptions, + requestAbortController + ? { headers: requestHeaders, signal: requestAbortController.signal } + : { headers: requestHeaders }, + ) .withResponse() let lastUsage @@ -315,10 +337,19 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa yield usageData } } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("LiteLLM streaming aborted", "AbortError") + } if (error instanceof Error) { throw new Error(`LiteLLM streaming error: ${error.message}`) } throw error + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -345,9 +376,28 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestOptions.max_tokens = info.maxTokens } - const response = await this.client.chat.completions.create(requestOptions) + // Build request options with abortSignal and/or timeout. The OpenAI SDK + // treats a timeout of 0 as an immediate timeout, so non-positive timeoutMs + // values disable the timeout instead of being forwarded. + const createOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + createOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + createOptions.timeout = options.timeoutMs + } + + const response = await this.client.chat.completions.create( + requestOptions, + Object.keys(createOptions).length > 0 ? createOptions : undefined, + ) return response.choices[0]?.message.content || "" } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("LiteLLM completion aborted", "AbortError") + } if (error instanceof Error) { throw new Error(`LiteLLM completion error: ${error.message}`) } diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index c7816feaa2..9c304f7eb0 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -101,66 +101,98 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand // Temporary debug log for QA // console.log("[MISTRAL DEBUG] Raw API request body:", requestOptions) + // Bridge the external abort signal from Task (metadata.abortSignal) into a + // request-local controller so the in-flight streaming request can be + // cancelled. A pre-aborted signal rejects immediately with an AbortError. + const externalAbortSignal = metadata?.abortSignal + let requestAbortController: AbortController | undefined + let externalAbortListener: (() => void) | undefined + if (externalAbortSignal) { + if (externalAbortSignal.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } + const controller = new AbortController() + requestAbortController = controller + const onExternalAbort = () => controller.abort() + externalAbortListener = onExternalAbort + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + let response try { - response = await this.client.chat.stream(requestOptions) - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error) - const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") - TelemetryService.instance.captureException(apiError) - throw new Error(`Mistral completion error: ${errorMessage}`) - } + if (requestAbortController) { + response = await this.client.chat.stream(requestOptions, { + fetchOptions: { signal: requestAbortController.signal }, + }) + } else { + response = await this.client.chat.stream(requestOptions) + } - for await (const event of response) { - const delta = event.data.choices[0]?.delta - - if (delta?.content) { - if (typeof delta.content === "string") { - // Handle string content as text - yield { type: "text", text: delta.content } - } else if (Array.isArray(delta.content)) { - // Handle array of content chunks - // The SDK v1.9.18 supports ThinkChunk with type "thinking" - for (const chunk of delta.content as ContentChunkWithThinking[]) { - if (chunk.type === "thinking" && chunk.thinking) { - // Handle thinking content as reasoning chunks - // ThinkChunk has a 'thinking' property that contains an array of text/reference chunks - for (const thinkingPart of chunk.thinking) { - if (thinkingPart.type === "text" && thinkingPart.text) { - yield { type: "reasoning", text: thinkingPart.text } + for await (const event of response) { + const delta = event.data.choices[0]?.delta + + if (delta?.content) { + if (typeof delta.content === "string") { + // Handle string content as text + yield { type: "text", text: delta.content } + } else if (Array.isArray(delta.content)) { + // Handle array of content chunks + // The SDK v1.9.18 supports ThinkChunk with type "thinking" + for (const chunk of delta.content as ContentChunkWithThinking[]) { + if (chunk.type === "thinking" && chunk.thinking) { + // Handle thinking content as reasoning chunks + // ThinkChunk has a 'thinking' property that contains an array of text/reference chunks + for (const thinkingPart of chunk.thinking) { + if (thinkingPart.type === "text" && thinkingPart.text) { + yield { type: "reasoning", text: thinkingPart.text } + } } + } else if (chunk.type === "text" && chunk.text) { + // Handle text content normally + yield { type: "text", text: chunk.text } } - } else if (chunk.type === "text" && chunk.text) { - // Handle text content normally - yield { type: "text", text: chunk.text } } } } - } - // Handle tool calls in stream - // Mistral SDK provides tool_calls in delta similar to OpenAI format - const toolCalls = (delta as { toolCalls?: MistralToolCall[] })?.toolCalls - if (toolCalls) { - for (let i = 0; i < toolCalls.length; i++) { - const toolCall = toolCalls[i] - yield { - type: "tool_call_partial", - index: i, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + // Handle tool calls in stream + // Mistral SDK provides tool_calls in delta similar to OpenAI format + const toolCalls = (delta as { toolCalls?: MistralToolCall[] })?.toolCalls + if (toolCalls) { + for (let i = 0; i < toolCalls.length; i++) { + const toolCall = toolCalls[i] + yield { + type: "tool_call_partial", + index: i, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, + } } } - } - if (event.data.usage) { - yield { - type: "usage", - inputTokens: event.data.usage.promptTokens || 0, - outputTokens: event.data.usage.completionTokens || 0, + if (event.data.usage) { + yield { + type: "usage", + inputTokens: event.data.usage.promptTokens || 0, + outputTokens: event.data.usage.completionTokens || 0, + } } } + } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (metadata?.abortSignal?.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } + const errorMessage = error instanceof Error ? error.message : String(error) + const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage") + TelemetryService.instance.captureException(apiError) + throw new Error(`Mistral completion error: ${errorMessage}`) + } finally { + if (externalAbortSignal && externalAbortListener) { + externalAbortSignal.removeEventListener("abort", externalAbortListener) + } } } @@ -196,11 +228,23 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand const { id: model, temperature } = this.getModel() try { - const response = await this.client.chat.complete({ - model, - messages: [{ role: "user", content: prompt }], - temperature, - }) + // Build Mistral SDK RequestOptions + const requestOptions: Parameters[1] = {} + if (options?.abortSignal) { + requestOptions.fetchOptions = { signal: options.abortSignal } + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeoutMs = options.timeoutMs + } + + const response = await this.client.chat.complete( + { + model, + messages: [{ role: "user", content: prompt }], + temperature, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = response.choices?.[0]?.message.content @@ -214,6 +258,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand return content || "" } catch (error) { + // User-initiated abort: surface a standard AbortError rather than the + // wrapped provider error so callers can distinguish cancellation. + if (options?.abortSignal?.aborted) { + throw new DOMException("Mistral completion aborted", "AbortError") + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model, "completePrompt") TelemetryService.instance.captureException(apiError) From f6eba43d3992f48e657feea481333ee6976d48f0 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 10:58:21 +0800 Subject: [PATCH 2/3] fix(api): unify timeoutMs:0 handling across gemini/mistral/lite-llm + fix test title --- .../__tests__/gemini-handler.spec.ts | 2 +- src/api/providers/__tests__/gemini.spec.ts | 15 ++++++++++ src/api/providers/__tests__/lite-llm.spec.ts | 6 ++++ src/api/providers/__tests__/mistral.spec.ts | 6 ++-- src/api/providers/gemini.ts | 9 ++++-- src/api/providers/lite-llm.ts | 13 ++++---- src/api/providers/mistral.ts | 8 +++-- .../utils/__tests__/request-timeout.spec.ts | 30 +++++++++++++++++++ src/api/providers/utils/request-timeout.ts | 10 +++++++ 9 files changed, 85 insertions(+), 14 deletions(-) create mode 100644 src/api/providers/utils/__tests__/request-timeout.spec.ts create mode 100644 src/api/providers/utils/request-timeout.ts diff --git a/src/api/providers/__tests__/gemini-handler.spec.ts b/src/api/providers/__tests__/gemini-handler.spec.ts index 364f62e23d..232849a091 100644 --- a/src/api/providers/__tests__/gemini-handler.spec.ts +++ b/src/api/providers/__tests__/gemini-handler.spec.ts @@ -55,7 +55,7 @@ describe("GeminiHandler backend support", () => { expect(promptConfig.tools).toBeUndefined() }) - it("completePrompt should pass abort signal through to client via httpOptions", async () => { + it("completePrompt should pass abort signal through to client via config.abortSignal", async () => { const options = { apiProvider: "gemini", enableUrlContext: false, diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 59968060dc..6350f3d395 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -415,6 +415,21 @@ describe("GeminiHandler", () => { }, }) }) + + it("should omit httpOptions entirely for timeoutMs=0 (0 disables the timeout)", async () => { + vi.mocked(handler["client"].models.generateContent).mockResolvedValue( + stubGenerateContentResponse("response"), + ) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(handler["client"].models.generateContent).toHaveBeenCalledWith({ + model: GEMINI_MODEL_NAME, + contents: [{ role: "user", parts: [{ text: "test prompt" }] }], + config: { + httpOptions: undefined, + temperature: 1, + }, + }) + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index badab1e1a7..324c532735 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1272,6 +1272,12 @@ describe("LiteLLMHandler", () => { const result = await handler.completePrompt("test prompt") expect(result).toBe("response") }) + + it("should omit the timeout option for timeoutMs=0 (0 would abort immediately in the OpenAI SDK)", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) + }) }) describe("createMessage abort signal (bridging)", () => { diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index aff33b6030..2fb0c43f22 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -530,14 +530,12 @@ describe("MistralHandler", () => { }) }) - it("should still forward timeoutMs=0 (uses !== undefined check, not truthy check)", async () => { + it("should omit the timeout option for timeoutMs=0 (0 disables the timeout)", async () => { mockComplete.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }], }) await handler.completePrompt("test prompt", { timeoutMs: 0 }) - expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { - timeoutMs: 0, - }) + expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) }) diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index bc6bcfe2f1..434cb06a27 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -27,6 +27,7 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, Complete import { BaseProvider } from "./base-provider" import { NOT_PROVIDED } from "./constants" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" +import { getRequestTimeoutMs } from "./utils/request-timeout" type GeminiHandlerOptions = ApiHandlerOptions & { isVertex?: boolean @@ -618,8 +619,12 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl ? (this.options.modelTemperature ?? info.defaultTemperature ?? 1) : info.defaultTemperature const httpOpts: { timeout?: number; baseUrl?: string } = {} - if (options?.timeoutMs !== undefined) { - httpOpts.timeout = options.timeoutMs + // Per the abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely (some SDKs treat 0 as an + // immediate timeout). + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + httpOpts.timeout = timeoutMs } if (this.options.googleGeminiBaseUrl) { httpOpts.baseUrl = this.options.googleGeminiBaseUrl diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index da82dd19cb..af5152ef70 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -16,6 +16,7 @@ import { sanitizeOpenAiCallId } from "../../utils/tool-id" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { getRequestTimeoutMs } from "./utils/request-timeout" /** * LiteLLM provider handler @@ -376,15 +377,17 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa requestOptions.max_tokens = info.maxTokens } - // Build request options with abortSignal and/or timeout. The OpenAI SDK - // treats a timeout of 0 as an immediate timeout, so non-positive timeoutMs - // values disable the timeout instead of being forwarded. + // Build request options with abortSignal and/or timeout. Per the + // abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely, because the OpenAI SDK treats + // a timeout of 0 as an immediate timeout. const createOptions: OpenAI.RequestOptions = {} if (options?.abortSignal) { createOptions.signal = options.abortSignal } - if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { - createOptions.timeout = options.timeoutMs + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + createOptions.timeout = timeoutMs } const response = await this.client.chat.completions.create( diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index 9c304f7eb0..80748fda9f 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -16,6 +16,7 @@ import { ApiHandlerOptions } from "../../shared/api" import { convertToMistralMessages } from "../transform/mistral-format" import { ApiStream } from "../transform/stream" import { handleProviderError } from "./utils/error-handler" +import { getRequestTimeoutMs } from "./utils/request-timeout" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -233,8 +234,11 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand if (options?.abortSignal) { requestOptions.fetchOptions = { signal: options.abortSignal } } - if (options?.timeoutMs !== undefined) { - requestOptions.timeoutMs = options.timeoutMs + // Per the abort-signal series contract, timeoutMs <= 0 means 'no per-request + // timeout': the option is omitted entirely. + const timeoutMs = getRequestTimeoutMs(options?.timeoutMs) + if (timeoutMs !== undefined) { + requestOptions.timeoutMs = timeoutMs } const response = await this.client.chat.complete( diff --git a/src/api/providers/utils/__tests__/request-timeout.spec.ts b/src/api/providers/utils/__tests__/request-timeout.spec.ts new file mode 100644 index 0000000000..5e969a1036 --- /dev/null +++ b/src/api/providers/utils/__tests__/request-timeout.spec.ts @@ -0,0 +1,30 @@ +import { getRequestTimeoutMs } from "../request-timeout" + +describe("getRequestTimeoutMs", () => { + it("forwards positive timeout values unchanged", () => { + expect(getRequestTimeoutMs(5000)).toBe(5000) + expect(getRequestTimeoutMs(1)).toBe(1) + expect(getRequestTimeoutMs(1234)).toBe(1234) + }) + + it("returns undefined for zero (timeout disabled, not an immediate abort)", () => { + expect(getRequestTimeoutMs(0)).toBeUndefined() + }) + + it("returns undefined for negative values", () => { + expect(getRequestTimeoutMs(-1)).toBeUndefined() + expect(getRequestTimeoutMs(-5000)).toBeUndefined() + }) + + it("returns undefined when no value is provided", () => { + expect(getRequestTimeoutMs()).toBeUndefined() + expect(getRequestTimeoutMs(undefined)).toBeUndefined() + }) + + it("guards against non-number input at the runtime boundary", () => { + expect(getRequestTimeoutMs(NaN)).toBeUndefined() + // Non-number values can only reach this helper through untyped callers + // (e.g. user settings); the double cast exercises the typeof guard. + expect(getRequestTimeoutMs("5000" as unknown as number)).toBeUndefined() + }) +}) diff --git a/src/api/providers/utils/request-timeout.ts b/src/api/providers/utils/request-timeout.ts new file mode 100644 index 0000000000..a3d1e56d4c --- /dev/null +++ b/src/api/providers/utils/request-timeout.ts @@ -0,0 +1,10 @@ +/** + * Returns the value to pass as a client/SDK request timeout option, or undefined. + * + * Per the abort-signal series contract, timeoutMs <= 0 (or undefined) means + * 'no per-request timeout': the option is omitted entirely, because some SDKs + * (e.g. the OpenAI Node SDK) treat timeout: 0 as an IMMEDIATE timeout. + */ +export function getRequestTimeoutMs(timeoutMs?: number): number | undefined { + return typeof timeoutMs === "number" && timeoutMs > 0 ? timeoutMs : undefined +} From 4a1307248c2e8ace2d83bcc5a7558f75d58066fb Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Thu, 20 Aug 2026 19:32:03 +0800 Subject: [PATCH 3/3] test(api): close changed-line coverage gaps in gemini, mistral, lite-llm --- src/api/providers/__tests__/gemini.spec.ts | 13 +++++ src/api/providers/__tests__/lite-llm.spec.ts | 13 +++++ src/api/providers/__tests__/mistral.spec.ts | 56 +++++++++++++++++++- 3 files changed, 81 insertions(+), 1 deletion(-) diff --git a/src/api/providers/__tests__/gemini.spec.ts b/src/api/providers/__tests__/gemini.spec.ts index 6350f3d395..f3b497c73f 100644 --- a/src/api/providers/__tests__/gemini.spec.ts +++ b/src/api/providers/__tests__/gemini.spec.ts @@ -430,6 +430,19 @@ describe("GeminiHandler", () => { }, }) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + const controller = new AbortController() + controller.abort() + vi.mocked(handler["client"].models.generateContent).mockRejectedValue(new Error("Gemini API error")) + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("Gemini completion aborted") + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index 324c532735..ee701b1134 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1278,6 +1278,19 @@ describe("LiteLLMHandler", () => { await handler.completePrompt("test prompt", { timeoutMs: 0 }) expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + mockCreate.mockRejectedValueOnce(new Error("LiteLLM API error")) + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("LiteLLM completion aborted") + }) }) describe("createMessage abort signal (bridging)", () => { diff --git a/src/api/providers/__tests__/mistral.spec.ts b/src/api/providers/__tests__/mistral.spec.ts index 2fb0c43f22..a5bfc75fd0 100644 --- a/src/api/providers/__tests__/mistral.spec.ts +++ b/src/api/providers/__tests__/mistral.spec.ts @@ -53,7 +53,12 @@ import type OpenAI from "openai" import { MistralHandler } from "../mistral" import type { ApiHandlerOptions } from "../../../shared/api" import type { ApiHandlerCreateMessageMetadata } from "../../index" -import type { ApiStreamTextChunk, ApiStreamReasoningChunk, ApiStreamToolCallPartialChunk } from "../../transform/stream" +import type { + ApiStreamTextChunk, + ApiStreamReasoningChunk, + ApiStreamToolCallPartialChunk, + ApiStreamUsageChunk, +} from "../../transform/stream" import { makeCreateMessageMetadata } from "../../../test-utils/api" describe("MistralHandler", () => { @@ -234,6 +239,42 @@ describe("MistralHandler", () => { expect(results[1]).toEqual({ type: "reasoning", text: "Some reasoning" }) expect(results[2]).toEqual({ type: "text", text: "Second text" }) }) + + it("should yield a usage chunk when the stream event carries usage data", async () => { + // The final event carries usage without any delta content; the handler + // must translate it into a usage chunk with the reported token counts. + mockCreate.mockImplementationOnce(async (_options) => + asyncStreamFrom([ + { + data: { + choices: [ + { + delta: { content: "Test response" }, + index: 0, + }, + ], + }, + }, + { + data: { + choices: [], + usage: { promptTokens: 12, completionTokens: 34 }, + }, + }, + ]), + ) + + const iterator = handler.createMessage(systemPrompt, messages) + const results: (ApiStreamTextChunk | ApiStreamUsageChunk)[] = [] + + for await (const chunk of iterator) { + results.push(chunk as ApiStreamTextChunk | ApiStreamUsageChunk) + } + + expect(results).toHaveLength(2) + expect(results[0]).toEqual({ type: "text", text: "Test response" }) + expect(results[1]).toEqual({ type: "usage", inputTokens: 12, outputTokens: 34 }) + }) }) describe("native tool calling", () => { @@ -537,6 +578,19 @@ describe("MistralHandler", () => { await handler.completePrompt("test prompt", { timeoutMs: 0 }) expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) }) + + it("should surface a standard AbortError when the signal was aborted and the request fails", async () => { + mockComplete.mockRejectedValueOnce(new Error("API Error")) + const controller = new AbortController() + controller.abort() + + const error = await handler + .completePrompt("Test prompt", { abortSignal: controller.signal }) + .catch((e: unknown) => e) + expect(error).toBeInstanceOf(DOMException) + expect((error as Error).name).toBe("AbortError") + expect((error as Error).message).toBe("Mistral completion aborted") + }) }) describe("createMessage abort signal bridging", () => {