diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts index 339fa36a07..b5aa5cb4af 100644 --- a/src/api/providers/__tests__/anthropic-vertex.spec.ts +++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts @@ -8,6 +8,7 @@ import { VERTEX_1M_CONTEXT_MODEL_IDS } from "@roo-code/types" import { AnthropicVertexHandler } from "../anthropic-vertex" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" vitest.mock("../utils/timeout-config", () => ({ getApiRequestTimeout: vitest.fn().mockReturnValue(300_000), @@ -746,6 +747,198 @@ describe("VertexHandler", () => { expect(calledMessages).toHaveLength(2) // Only the two user messages expect(calledMessages.every((m: any) => m.role === "user")).toBe(true) }) + + it("should reject with AbortError when createMessage is called with an already-aborted signal", async () => { + const abortedController = new AbortController() + abortedController.abort() + + vitest + .spyOn(handler["client"].messages, "create") + .mockImplementation((_params: unknown, options?: { signal?: AbortSignal | null }) => { + if (options?.signal?.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + throw error + } + return asyncStreamFrom([]) as never + }) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ) + + await expect(stream.next()).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the request when the external signal aborts mid-flight", async () => { + const controller = new AbortController() + + vitest.spyOn(handler["client"].messages, "create").mockImplementation( + (_params: unknown, options?: { signal?: AbortSignal | null }) => + new Promise((_resolve, reject) => { + const signal = options?.signal + if (!signal) { + return + } + if (signal.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + return + } + signal.addEventListener( + "abort", + () => { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + }, + { once: true }, + ) + }) as never, + ) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const promise = stream.next() + controller.abort() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should remove the external abort listener when the stream completes", async () => { + const handlerWithSignal = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + const stream = handlerWithSignal.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await collectStream(stream) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached once the stream is done. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when message creation fails", async () => { + const handlerWithSignal = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + handlerWithSignal["client"].messages.create = vitest + .fn() + .mockRejectedValue(new Error("vertex create error")) + + const stream = handlerWithSignal.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toThrow("vertex create error") + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The inline finally must detach the listener on the failure path too, + // not only on normal completion. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when the consumer returns early", async () => { + const handlerWithSignal = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + handlerWithSignal["client"].messages.create = vitest.fn().mockImplementation(async () => + asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 1, output_tokens: 0 } } }, + { type: "content_block_delta", delta: { type: "text_delta", text: "first" } }, + { type: "content_block_delta", delta: { type: "text_delta", text: "second" } }, + ]), + ) + + const stream = handlerWithSignal.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const iterator = stream[Symbol.asyncIterator]() + + await iterator.next() + await iterator.return(undefined) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // Abandoning the generator mid-stream runs the finally block too. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should default message_start outputTokens to zero when output_tokens is omitted", async () => { + handler = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const mockCreate = vitest.fn().mockImplementation(async () => + asyncStreamFrom([ + { + type: "message_start", + message: { + usage: { + input_tokens: 10, + }, + }, + }, + ]), + ) + handler["client"].messages.create = mockCreate + + const stream = handler.createMessage(systemPrompt, [{ role: "user", content: "Hello" }]) + const chunks = await collectStream(stream) + + expect(chunks[0]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 0, + cacheWriteTokens: undefined, + cacheReadTokens: undefined, + }) + }) }) describe("completePrompt", () => { @@ -758,18 +951,22 @@ describe("VertexHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("Test response") - expect(handler["client"].messages.create).toHaveBeenCalledWith({ - model: "claude-3-5-sonnet-v2@20241022", - max_tokens: 8192, - temperature: 0, - messages: [ - { - role: "user", - content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }], - }, - ], - stream: false, - }) + expect(handler["client"].messages.create).toHaveBeenCalledWith( + { + model: "claude-3-5-sonnet-v2@20241022", + max_tokens: 8192, + temperature: 0, + messages: [ + { + role: "user", + content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }], + }, + ], + stream: false, + thinking: undefined, + }, + undefined, + ) }) it("should handle API errors for Claude", async () => { @@ -820,6 +1017,98 @@ describe("VertexHandler", () => { expect(result).toBe("") }) + it("should pass abort signal through to client", async () => { + handler = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const controller = new AbortController() + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + + const [, requestOptions] = mockCreate.mock.calls[0] + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.any(Object), + ) + expect(requestOptions?.signal).toBe(controller.signal) + }) + + it("should work without options (backward compatible)", async () => { + handler = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), undefined) + }) + + it("completePrompt should pass signal through to client", async () => { + handler = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const controller = new AbortController() + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + + const [, requestOptions] = mockCreate.mock.calls[0] + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 5000 }), + ) + expect(requestOptions?.signal).toBe(controller.signal) + }) + + it("completePrompt should pass timeoutMs when provided", async () => { + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + await handler.completePrompt("test prompt", { timeoutMs: 3000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 3000 }), + ) + }) + + it("completePrompt should pass timeout when timeoutMs=0 (defined check)", async () => { + handler = new AnthropicVertexHandler({ + apiModelId: "claude-3-5-sonnet-v2@20241022", + vertexProjectId: "test-project", + vertexRegion: "us-central1", + }) + + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + // 0 is a defined value: it must reach the client as `timeout: 0`, + // not be dropped by a truthiness check. + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: 0 }), + ) + }) + it("should handle empty content array for Claude", async () => { handler = new AnthropicVertexHandler({ apiModelId: "claude-3-5-sonnet-v2@20241022", diff --git a/src/api/providers/__tests__/anthropic.spec.ts b/src/api/providers/__tests__/anthropic.spec.ts index 30fc120684..d1a171ba83 100644 --- a/src/api/providers/__tests__/anthropic.spec.ts +++ b/src/api/providers/__tests__/anthropic.spec.ts @@ -3,6 +3,7 @@ import { AnthropicHandler } from "../anthropic" import { ApiHandlerOptions } from "../../../shared/api" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import type { ApiStreamChunk } from "../../../api/transform/stream" import { clearAllMocks } from "../../../test-utils/reset" @@ -767,20 +768,273 @@ describe("AnthropicHandler", () => { expect(deltaUsage?.outputTokens).toBe(200) expect(deltaUsage?.reasoningTokens).toBe(0) }) + + it("should reject with AbortError when createMessage is called with an already-aborted signal", async () => { + const abortedController = new AbortController() + abortedController.abort() + + mockCreate.mockImplementation(async (_params: unknown, options?: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + throw error + } + return asyncStreamFrom([]) + }) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ) + + await expect(stream.next()).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the request when the external signal aborts mid-flight", async () => { + const controller = new AbortController() + + mockCreate.mockImplementation((_params: unknown, options?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = options?.signal + if (!signal) { + return + } + if (signal.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + return + } + signal.addEventListener( + "abort", + () => { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + }, + { once: true }, + ) + }) + }) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const promise = stream.next() + controller.abort() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should not pre-abort the bridged signal for a pending external signal", async () => { + const controller = new AbortController() + + mockCreate.mockImplementationOnce(async () => asyncStreamFrom([])) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await collectStream(stream) + + const requestOptions = mockCreate.mock.calls[mockCreate.mock.calls.length - 1]?.[1] + expect(requestOptions?.signal?.aborted).toBe(false) + }) + + it("should register the external abort listener with the once option", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + + mockCreate.mockImplementationOnce(async () => asyncStreamFrom([])) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await collectStream(stream) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + // The listener is registered once so it detaches itself when the signal aborts. + expect(addEventListenerSpy.mock.calls[0][2]).toEqual({ once: true }) + }) + + it("should remove the external abort listener when the stream completes", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockCreate.mockImplementationOnce(async () => asyncStreamFrom([])) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await collectStream(stream) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached once the stream is done. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when message creation fails", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockCreate.mockRejectedValueOnce(new Error("Anthropic createMessage error")) + + const stream = handler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toThrow("Anthropic createMessage error") + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached when creation fails. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should propagate the creation error when no external abort signal is provided", async () => { + mockCreate.mockRejectedValueOnce(new Error("Anthropic createMessage error")) + + const stream = handler.createMessage(systemPrompt, [{ role: "user", content: "Hello" }]) + + await expect(collectStream(stream)).rejects.toThrow("Anthropic createMessage error") + }) + + it("should remove the external abort listener when message creation fails for a non-cached model", async () => { + const customHandler = new AnthropicHandler({ + apiKey: "test-api-key", + apiModelId: "claude-sonnet-5-bf", + }) + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockCreate.mockRejectedValueOnce(new Error("Anthropic createMessage error")) + + const stream = customHandler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toThrow("Anthropic createMessage error") + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached when creation fails. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should reject with AbortError when the external signal aborts a non-cached model mid-flight", async () => { + const customHandler = new AnthropicHandler({ + apiKey: "test-api-key", + apiModelId: "claude-sonnet-5-bf", + }) + const controller = new AbortController() + + mockCreate.mockImplementation((_params: unknown, options?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = options?.signal + if (!signal) { + return + } + if (signal.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + return + } + signal.addEventListener( + "abort", + () => { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + }, + { once: true }, + ) + }) + }) + + const stream = customHandler.createMessage( + systemPrompt, + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const promise = stream.next() + controller.abort() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + + // The non-cached path must hand the bridged signal to the SDK. + const requestOptions = mockCreate.mock.calls[mockCreate.mock.calls.length - 1]?.[1] + expect(requestOptions?.signal).toBeDefined() + expect(requestOptions?.signal?.aborted).toBe(true) + }) + + it("should propagate the creation error for a non-cached model when no external abort signal is provided", async () => { + const customHandler = new AnthropicHandler({ + apiKey: "test-api-key", + apiModelId: "claude-sonnet-5-bf", + }) + + mockCreate.mockRejectedValueOnce(new Error("Anthropic createMessage error")) + + const stream = customHandler.createMessage(systemPrompt, [{ role: "user", content: "Hello" }]) + + await expect(collectStream(stream)).rejects.toThrow("Anthropic createMessage error") + }) + + it("should join beta headers with a comma for the default model", async () => { + mockCreate.mockImplementationOnce(async () => asyncStreamFrom([])) + + const stream = handler.createMessage(systemPrompt, [{ role: "user", content: "Hello" }]) + + await collectStream(stream) + + const requestOptions = mockCreate.mock.calls[mockCreate.mock.calls.length - 1]?.[1] + expect(requestOptions?.headers?.["anthropic-beta"]).toBe( + "fine-grained-tool-streaming-2025-05-14,prompt-caching-2024-07-31", + ) + }) }) describe("completePrompt", () => { it("should complete prompt successfully", async () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("Test response") - expect(mockCreate).toHaveBeenCalledWith({ - model: mockOptions.apiModelId, - messages: [{ role: "user", content: "Test prompt" }], - max_tokens: 8192, - temperature: 0, - thinking: undefined, - stream: false, - }) + expect(mockCreate).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: "Test prompt" }], + max_tokens: 8192, + temperature: 0, + thinking: undefined, + stream: false, + }, + undefined, + ) }) it("should handle API errors", async () => { @@ -803,6 +1057,95 @@ describe("AnthropicHandler", () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("") }) + + it("should pass abort signal through to client", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: "test prompt" }], + max_tokens: 8192, + temperature: 0, + thinking: undefined, + stream: false, + }, + { signal: controller.signal }, + ) + }) + + it("should work without options (backward compatible)", async () => { + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockCreate).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: "test prompt" }], + max_tokens: 8192, + temperature: 0, + thinking: undefined, + stream: false, + }, + undefined, + ) + }) + + it("should merge signal and timeout together", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 }) + expect(mockCreate).toHaveBeenCalledWith( + { + model: mockOptions.apiModelId, + messages: [{ role: "user", content: "test prompt" }], + max_tokens: 8192, + temperature: 0, + thinking: undefined, + stream: false, + }, + expect.objectContaining({ signal: controller.signal, timeout: 10000 }), + ) + }) + + it("should pass timeoutMs through to client alongside abortSignal", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: mockOptions.apiModelId }), + expect.objectContaining({ signal: controller.signal, timeout: 5000 }), + ) + }) + + it("should pass the same signal instance", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.any(Object), + expect.objectContaining({ signal: controller.signal }), + ) + // Verify it's the exact same instance, not just equal + const callOptions = mockCreate.mock.calls[0][1] + expect(callOptions?.signal).toBe(controller.signal) + }) + + it("should not include signal-related options when not provided", async () => { + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt") + expect(mockCreate).toHaveBeenCalledWith(expect.any(Object), undefined) + }) + + it("should pass timeout when timeoutMs=0 (defined check)", async () => { + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + // timeoutMs=0 must be forwarded as an explicit 0 timeout, not dropped as if unset + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: mockOptions.apiModelId }), { + timeout: 0, + }) + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/minimax.spec.ts b/src/api/providers/__tests__/minimax.spec.ts index 063932b303..be2343f804 100644 --- a/src/api/providers/__tests__/minimax.spec.ts +++ b/src/api/providers/__tests__/minimax.spec.ts @@ -13,7 +13,8 @@ import { Anthropic } from "@anthropic-ai/sdk" import { type MinimaxModelId, minimaxDefaultModelId, minimaxModels } from "@roo-code/types" import { MiniMaxHandler } from "../minimax" -import { asyncStreamFrom } from "../../../test-utils/stream" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" vitest.mock("@anthropic-ai/sdk", () => { @@ -239,6 +240,86 @@ describe("MiniMaxHandler", () => { await expect(handler.completePrompt("test prompt")).rejects.toThrow() }) + it("should pass abort signal through to client", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + { signal: controller.signal }, // second arg (options) + ) + }) + + it("should work without options (backward compatible)", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, // second arg (options) + ) + }) + + it("should pass timeout through to client", async () => { + const controller = new AbortController() + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + signal: controller.signal, + timeout: 5000, + }) + }) + + it("should pass only timeoutMs when no signal provided", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 3000 }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeout: 3000, + }) + }) + + it("should pass timeout when timeoutMs=0 (defined check)", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + { timeout: 0 }, // !== undefined check means 0 is passed through + ) + }) + + it("completePrompt should send the configured temperature, user message and stream flag", async () => { + const temperatureHandler = new MiniMaxHandler({ + minimaxApiKey: "test-minimax-api-key", + modelTemperature: 0.5, + }) + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + + await temperatureHandler.completePrompt("test prompt") + + expect(mockCreate).toHaveBeenCalledWith( + { + model: minimaxDefaultModelId, + max_tokens: 16_384, + temperature: 0.5, + messages: [{ role: "user", content: "test prompt" }], + stream: false, + }, + undefined, + ) + }) + it("createMessage should yield text content from stream", async () => { const testContent = "This is test content from MiniMax stream" @@ -306,6 +387,7 @@ describe("MiniMaxHandler", () => { messages: expect.any(Array), stream: true, }), + undefined, ) }) @@ -319,6 +401,7 @@ describe("MiniMaxHandler", () => { expect.objectContaining({ temperature: 1, }), + undefined, ) }) @@ -375,6 +458,159 @@ describe("MiniMaxHandler", () => { arguments: undefined, }) }) + + it("should reject with AbortError when createMessage is called with an already-aborted signal", async () => { + const abortedController = new AbortController() + abortedController.abort() + + mockCreate.mockImplementation(async (_params: unknown, options?: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + throw error + } + return asyncStreamFrom([]) + }) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ) + + await expect(stream.next()).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the request when the external signal aborts mid-flight", async () => { + const controller = new AbortController() + + mockCreate.mockImplementation((_params: unknown, options?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = options?.signal + if (!signal) { + return + } + if (signal.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + return + } + signal.addEventListener( + "abort", + () => { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + }, + { once: true }, + ) + }) + }) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const promise = stream.next() + controller.abort() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should not pre-abort the bridged signal for a pending external signal", async () => { + const controller = new AbortController() + let observedSignal: AbortSignal | undefined + + mockCreate.mockImplementation(async (_params: unknown, options?: { signal?: AbortSignal }) => { + observedSignal = options?.signal + return asyncStreamFrom([]) + }) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).resolves.toEqual([]) + + expect(observedSignal?.aborted).toBe(false) + }) + + it("should register the external abort listener with the once option", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + + mockCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).resolves.toEqual([]) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + // The listener is registered once so it detaches itself when the signal aborts. + expect(addEventListenerSpy.mock.calls[0][2]).toEqual({ once: true }) + }) + + it("should remove the external abort listener when the stream completes", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).resolves.toEqual([]) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached once the stream is done. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when message creation fails", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockCreate.mockRejectedValueOnce(new Error("MiniMax createMessage error")) + + const stream = handler.createMessage( + "system prompt", + [{ role: "user", content: "Hello" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(collectStream(stream)).rejects.toThrow("MiniMax createMessage error") + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The same retained callback must be detached when creation fails. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should propagate the creation error when no external abort signal is provided", async () => { + mockCreate.mockRejectedValueOnce(new Error("MiniMax createMessage error")) + + const stream = handler.createMessage("system prompt", [{ role: "user", content: "Hello" }]) + + await expect(collectStream(stream)).rejects.toThrow("MiniMax createMessage error") + }) }) describe("Model Configuration", () => { diff --git a/src/api/providers/__tests__/stream-cleanup.spec.ts b/src/api/providers/__tests__/stream-cleanup.spec.ts new file mode 100644 index 0000000000..677ef62c8f --- /dev/null +++ b/src/api/providers/__tests__/stream-cleanup.spec.ts @@ -0,0 +1,52 @@ +// npx vitest run src/api/providers/__tests__/stream-cleanup.spec.ts + +import { withFinallyCleanup } from "../stream-cleanup" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" + +describe("withFinallyCleanup", () => { + it("forwards every chunk and runs the cleanup after the stream completes", async () => { + const cleanup = vitest.fn() + + const chunks = await collectStream(withFinallyCleanup(asyncStreamFrom(["a", "b", "c"]), cleanup)) + + expect(chunks).toEqual(["a", "b", "c"]) + expect(cleanup).toHaveBeenCalledTimes(1) + }) + + it("runs the cleanup when the consumer stops early", async () => { + const cleanup = vitest.fn() + const stream = withFinallyCleanup(asyncStreamFrom(["a", "b", "c"]), cleanup) + const iterator = stream[Symbol.asyncIterator]() + + expect((await iterator.next()).value).toBe("a") + await iterator.return(undefined) + + expect(cleanup).toHaveBeenCalledTimes(1) + }) + + it("runs the cleanup when the underlying stream fails", async () => { + const cleanup = vitest.fn() + async function* failing() { + yield "a" + throw new Error("boom") + } + + await expect(collectStream(withFinallyCleanup(failing(), cleanup))).rejects.toThrow("boom") + expect(cleanup).toHaveBeenCalledTimes(1) + }) + + it("runs the cleanup at most once when the consumer stops after a failure", async () => { + const cleanup = vitest.fn() + async function* failing() { + yield* [] + throw new Error("boom") + } + const stream = withFinallyCleanup(failing(), cleanup) + const iterator = stream[Symbol.asyncIterator]() + + await expect(iterator.next()).rejects.toThrow("boom") + await iterator.return(undefined) + + expect(cleanup).toHaveBeenCalledTimes(1) + }) +}) diff --git a/src/api/providers/__tests__/xai.spec.ts b/src/api/providers/__tests__/xai.spec.ts index d0524ff0e4..cbff98e0dd 100644 --- a/src/api/providers/__tests__/xai.spec.ts +++ b/src/api/providers/__tests__/xai.spec.ts @@ -14,15 +14,22 @@ const mockResponsesCreate = vitest.hoisted(() => vitest.fn()) vitest.mock("openai", async () => { const { mockOpenAiResponsesClient } = await import("../../../test-utils/api") - return mockOpenAiResponsesClient(mockResponsesCreate) + const actual = await vi.importActual("openai") + return { + ...mockOpenAiResponsesClient(mockResponsesCreate), + // Expose the real SDK error class so the mock transport can throw it exactly + // the way the OpenAI SDK does when a request signal aborts. + APIUserAbortError: actual.APIUserAbortError, + } }) -import OpenAI from "openai" +import OpenAI, { APIUserAbortError } from "openai" import { xaiDefaultModelId, xaiModels } from "@roo-code/types" import { XAIHandler } from "../xai" -import { asyncStreamFrom } from "../../../test-utils/stream" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" describe("XAIHandler", () => { @@ -97,6 +104,7 @@ describe("XAIHandler", () => { store: false, include: ["reasoning.encrypted_content"], }), + undefined, ) }) @@ -226,9 +234,66 @@ describe("XAIHandler", () => { tool_choice: "auto", parallel_tool_calls: true, }), + undefined, ) }) + it("createMessage should honor an explicit parallelToolCalls false", async () => { + const testTools = [ + { + type: "function" as const, + function: { + name: "test_tool", + description: "A test Tool", + parameters: { type: "object", properties: { arg1: { type: "string" } }, required: ["arg1"] }, + }, + }, + ] + + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", [], { + taskId: "test-task-id", + tools: testTools, + parallelToolCalls: false, + }) + await stream.next() + + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ + parallel_tool_calls: false, + }), + undefined, + ) + }) + + it("createMessage should send max_output_tokens and temperature for the default model", async () => { + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", []) + await stream.next() + + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ + max_output_tokens: 65_536, + temperature: 0, + }), + undefined, + ) + }) + + it("createMessage should omit tool parameters when no tools are provided", async () => { + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", []) + await stream.next() + + const callArgs = mockResponsesCreate.mock.calls[mockResponsesCreate.mock.calls.length - 1][0] + expect(callArgs).not.toHaveProperty("tools") + expect(callArgs).not.toHaveProperty("tool_choice") + expect(callArgs).not.toHaveProperty("parallel_tool_calls") + }) + it("completePrompt should return text from Responses API", async () => { const expectedResponse = "This is a test response" mockResponsesCreate.mockResolvedValueOnce({ @@ -246,6 +311,101 @@ describe("XAIHandler", () => { await expect(handler.completePrompt("test prompt")).rejects.toThrow(`xAI completion error: ${errorMessage}`) }) + it("completePrompt should normalize the SDK APIUserAbortError to the contract AbortError on abort", async () => { + const controller = new AbortController() + controller.abort() + const sdkAbortError = new APIUserAbortError() + mockResponsesCreate.mockRejectedValueOnce(sdkAbortError) + + // The raw SDK abort error (name "Error", message "Request was aborted.") does not + // satisfy the caller's AbortError contract; it must be normalized, not wrapped by + // handleOpenAIError. + await expect(handler.completePrompt("test prompt", { abortSignal: controller.signal })).rejects.toMatchObject({ + name: "AbortError", + message: "The xAI request was aborted", + }) + }) + + it("completePrompt should normalize a native AbortError to the contract error on abort", async () => { + const controller = new AbortController() + controller.abort() + const abortError = new Error("This operation was aborted") + abortError.name = "AbortError" + mockResponsesCreate.mockRejectedValueOnce(abortError) + + // Normalized to the branded contract message while keeping the contract name. + await expect(handler.completePrompt("test prompt", { abortSignal: controller.signal })).rejects.toMatchObject({ + name: "AbortError", + message: "The xAI request was aborted", + }) + }) + + it("completePrompt should pass abort signal through to client", async () => { + const controller = new AbortController() + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + expect(mockResponsesCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + signal: controller.signal, + }) + }) + + it("completePrompt should work without options (backward compatible)", async () => { + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + const result = await handler.completePrompt("test prompt") + expect(result).toBe("response") + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, + ) + }) + + it("completePrompt should send the full Responses API request body", async () => { + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt") + + expect(mockResponsesCreate).toHaveBeenCalledWith( + { + model: xaiDefaultModelId, + input: [{ role: "user", content: [{ type: "input_text", text: "test prompt" }] }], + store: false, + }, + undefined, + ) + }) + + it("completePrompt should pass timeout through to client", async () => { + const controller = new AbortController() + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 }) + expect(mockResponsesCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + signal: controller.signal, + timeout: 5000, + }) + }) + + it("completePrompt should pass only timeoutMs when no signal provided", async () => { + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt", { timeoutMs: 3000 }) + expect(mockResponsesCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeout: 3000, + }) + }) + + it("completePrompt should pass timeout when timeoutMs=0 (defined check)", async () => { + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + { timeout: 0 }, // !== undefined check means 0 is passed through + ) + }) + it("should include reasoning effort for mini models in Responses API format", async () => { const miniModelHandler = new XAIHandler({ apiModelId: "grok-3-mini", @@ -263,6 +423,7 @@ describe("XAIHandler", () => { effort: "high", }), }), + undefined, ) }) @@ -283,6 +444,7 @@ describe("XAIHandler", () => { effort: "high", }), }), + undefined, ) }) @@ -304,6 +466,7 @@ describe("XAIHandler", () => { effort: "low", }), }), + undefined, ) }) @@ -329,4 +492,219 @@ describe("XAIHandler", () => { const stream = handler.createMessage("test prompt", []) await expect(stream.next()).rejects.toThrow(`xAI completion error: ${errorMessage}`) }) + + it("should capture createMessage failures in telemetry with operation context", async () => { + mockResponsesCreate.mockRejectedValueOnce(new Error("Stream error")) + + const stream = handler.createMessage("test prompt", []) + await expect(stream.next()).rejects.toThrow("xAI completion error: Stream error") + + expect(mockCaptureException).toHaveBeenCalledWith( + expect.objectContaining({ + name: "ApiProviderError", + provider: "xAI", + modelId: xaiDefaultModelId, + operation: "createMessage", + }), + ) + }) + + it("createMessage should normalize the SDK APIUserAbortError to the contract AbortError on abort", async () => { + const abortedController = new AbortController() + abortedController.abort() + const sdkAbortError = new APIUserAbortError() + + mockResponsesCreate.mockImplementation(async (_params: unknown, options?: { signal?: AbortSignal }) => { + // Mimic the OpenAI SDK: reject with its own APIUserAbortError when the + // request signal is aborted. + if (options?.signal?.aborted) { + throw sdkAbortError + } + return asyncStreamFrom([]) + }) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ) + + // The raw SDK abort error does not satisfy the caller's AbortError contract: it + // must be normalized, not wrapped by handleOpenAIError. + await expect(stream.next()).rejects.toMatchObject({ + name: "AbortError", + message: "The xAI request was aborted", + }) + }) + + it("createMessage should normalize an SDK abort raised while iterating the response stream", async () => { + const controller = new AbortController() + const sdkAbortError = new APIUserAbortError() + + // responses.create resolves; the cancellation lands while the stream is + // being consumed, so the SDK abort error escapes from the iterator. + mockResponsesCreate.mockResolvedValueOnce({ + [Symbol.asyncIterator]: () => ({ + next: async () => { + controller.abort() + throw sdkAbortError + }, + }), + }) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + // The raw SDK abort error (name "Error") must be normalized to the + // contract AbortError even though it escapes outside the create call. + await expect(stream.next()).rejects.toMatchObject({ + name: "AbortError", + message: "The xAI request was aborted", + }) + }) + + it("should reject with AbortError when createMessage is called with an already-aborted signal", async () => { + const abortedController = new AbortController() + abortedController.abort() + + mockResponsesCreate.mockImplementation(async (_params: unknown, options?: { signal?: AbortSignal }) => { + if (options?.signal?.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + throw error + } + return asyncStreamFrom([]) + }) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: abortedController.signal }), + ) + + await expect(stream.next()).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the request when the external signal aborts mid-flight", async () => { + const controller = new AbortController() + + mockResponsesCreate.mockImplementation((_params: unknown, options?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = options?.signal + if (!signal) { + return + } + if (signal.aborted) { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + return + } + signal.addEventListener( + "abort", + () => { + const error = new Error("The operation was aborted") + error.name = "AbortError" + reject(error) + }, + { once: true }, + ) + }) + }) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const promise = stream.next() + controller.abort() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should remove the external abort listener when the stream completes", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockResponsesCreate.mockResolvedValueOnce( + asyncStreamFrom([{ type: "response.output_text.delta", delta: "done" }]), + ) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const chunks = await collectStream(stream) + expect(chunks).toEqual([{ type: "text", text: "done" }]) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The listener is registered once so it detaches itself when the signal aborts. + expect(addEventListenerSpy.mock.calls[0][2]).toEqual({ once: true }) + // The same retained callback must be detached once the stream is done. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when the SDK call rejects", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockResponsesCreate.mockRejectedValueOnce(new Error("xai createMessage error")) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(stream.next()).rejects.toThrow() + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // The finally block must detach the listener when the request fails, not + // only on normal completion. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) + + it("should remove the external abort listener when the consumer returns early", async () => { + const controller = new AbortController() + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + + mockResponsesCreate.mockResolvedValueOnce( + asyncStreamFrom([ + { type: "response.output_text.delta", delta: "first" }, + { type: "response.output_text.delta", delta: "second" }, + ]), + ) + + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const iterator = stream[Symbol.asyncIterator]() + + expect((await iterator.next()).value).toEqual({ type: "text", text: "first" }) + await iterator.return(undefined) + + expect(addEventListenerSpy).toHaveBeenCalledTimes(1) + const [event, listener] = addEventListenerSpy.mock.calls[0] + expect(event).toBe("abort") + // Abandoning the generator mid-stream runs the finally block too. + expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) + }) }) diff --git a/src/api/providers/anthropic-vertex.ts b/src/api/providers/anthropic-vertex.ts index 6ceaad2256..bdc473f69a 100644 --- a/src/api/providers/anthropic-vertex.ts +++ b/src/api/providers/anthropic-vertex.ts @@ -24,6 +24,7 @@ import { } from "../../core/prompts/tools/native-tools/converters" import { BaseProvider } from "./base-provider" +import { bridgeExternalAbortSignal } from "./utils/abort-signal" import { NOT_PROVIDED } from "./constants" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -86,6 +87,15 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple tool_choice: convertOpenAIToolChoiceToAnthropic(metadata?.tool_choice, metadata?.parallelToolCalls), } + // Bridge the external abort signal from request metadata into a per-request + // controller so the SDK call is cancelled when the owning request is aborted + // (or when the signal is already aborted). Without an external signal the + // client-level timeout configured in the constructor remains the only + // cancellation mechanism, preserving the existing behavior. + const externalAbortSignal = metadata?.abortSignal + const { signal: abortSignal, dispose: removeExternalAbortListener } = + bridgeExternalAbortSignal(externalAbortSignal) + /** * Vertex API has specific limitations for prompt caching: * 1. Maximum of 4 blocks can have cache_control @@ -114,109 +124,128 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple } as Anthropic.Messages.MessageCreateParamsStreaming // and prompt caching - const requestOptions = betas?.length ? { headers: { "anthropic-beta": betas.join(",") } } : undefined - - const stream = await this.client.messages.create(params, requestOptions) - - for await (const chunk of stream) { - switch (chunk.type) { - case "message_start": { - const usage = chunk.message!.usage - - // `output_tokens_details.thinking_tokens` decomposes the billed - // output into internal reasoning tokens; surfaced for per-turn - // thinking telemetry (billing already includes it in output_tokens). - yield { - type: "usage", - inputTokens: usage.input_tokens || 0, - outputTokens: usage.output_tokens || 0, - cacheWriteTokens: usage.cache_creation_input_tokens || undefined, - cacheReadTokens: usage.cache_read_input_tokens || undefined, - ...(typeof usage.output_tokens_details?.thinking_tokens === "number" - ? { reasoningTokens: usage.output_tokens_details!.thinking_tokens } - : {}), - } + const requestOptions: Anthropic.RequestOptions = {} + if (betas?.length) { + // Join rather than take the first entry: the header accepts a comma list, so a + // second beta added to the builder later is not silently dropped. + requestOptions.headers = { "anthropic-beta": betas.join(",") } + } + if (abortSignal) { + requestOptions.signal = abortSignal + } - break - } - case "message_delta": { - const deltaThinkingTokens = chunk.usage!.output_tokens_details?.thinking_tokens - - yield { - type: "usage", - inputTokens: 0, - outputTokens: chunk.usage!.output_tokens || 0, - ...(typeof deltaThinkingTokens === "number" ? { reasoningTokens: deltaThinkingTokens } : {}), + try { + const stream = await this.client.messages.create( + params, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) + + for await (const chunk of stream) { + switch (chunk.type) { + case "message_start": { + const usage = chunk.message!.usage + + // `output_tokens_details.thinking_tokens` decomposes the billed + // output into internal reasoning tokens; surfaced for per-turn + // thinking telemetry (billing already includes it in output_tokens). + yield { + type: "usage", + inputTokens: usage.input_tokens || 0, + outputTokens: usage.output_tokens || 0, + cacheWriteTokens: usage.cache_creation_input_tokens || undefined, + cacheReadTokens: usage.cache_read_input_tokens || undefined, + ...(typeof usage.output_tokens_details?.thinking_tokens === "number" + ? { reasoningTokens: usage.output_tokens_details!.thinking_tokens } + : {}), + } + + break } + case "message_delta": { + const deltaThinkingTokens = chunk.usage!.output_tokens_details?.thinking_tokens + + yield { + type: "usage", + inputTokens: 0, + outputTokens: chunk.usage!.output_tokens || 0, + ...(typeof deltaThinkingTokens === "number" + ? { reasoningTokens: deltaThinkingTokens } + : {}), + } - break - } - case "content_block_start": { - switch (chunk.content_block!.type) { - case "text": { - if (chunk.index! > 0) { - yield { type: "text", text: "\n" } + break + } + case "content_block_start": { + switch (chunk.content_block!.type) { + case "text": { + if (chunk.index! > 0) { + yield { type: "text", text: "\n" } + } + + yield { type: "text", text: chunk.content_block!.text } + break } + case "thinking": { + if (chunk.index! > 0) { + yield { type: "reasoning", text: "\n" } + } - yield { type: "text", text: chunk.content_block!.text } - break - } - case "thinking": { - if (chunk.index! > 0) { - yield { type: "reasoning", text: "\n" } + yield { type: "reasoning", text: (chunk.content_block as any).thinking } + break } - - yield { type: "reasoning", text: (chunk.content_block as any).thinking } - break - } - case "tool_use": { - // Emit initial tool call partial with id and name - yield { - type: "tool_call_partial", - index: chunk.index, - id: chunk.content_block!.id, - name: chunk.content_block!.name, - arguments: undefined, + case "tool_use": { + // Emit initial tool call partial with id and name + yield { + type: "tool_call_partial", + index: chunk.index, + id: chunk.content_block!.id, + name: chunk.content_block!.name, + arguments: undefined, + } + break } - break } - } - break - } - case "content_block_delta": { - switch (chunk.delta!.type) { - case "text_delta": { - yield { type: "text", text: chunk.delta!.text } - break - } - case "thinking_delta": { - yield { type: "reasoning", text: (chunk.delta as any).thinking } - break - } - case "input_json_delta": { - // Emit tool call partial chunks as arguments stream in - yield { - type: "tool_call_partial", - index: chunk.index, - id: undefined, - name: undefined, - arguments: (chunk.delta as any).partial_json, + break + } + case "content_block_delta": { + switch (chunk.delta!.type) { + case "text_delta": { + yield { type: "text", text: chunk.delta!.text } + break + } + case "thinking_delta": { + yield { type: "reasoning", text: (chunk.delta as any).thinking } + break + } + case "input_json_delta": { + // Emit tool call partial chunks as arguments stream in + yield { + type: "tool_call_partial", + index: chunk.index, + id: undefined, + name: undefined, + arguments: (chunk.delta as any).partial_json, + } + break } - break } - } - break - } - case "content_block_stop": { - // Block complete - no action needed for now. - // NativeToolCallParser handles tool call completion - // Note: Signature for multi-turn thinking would require using stream.finalMessage() - // after iteration completes, which requires restructuring the streaming approach. - break + break + } + case "content_block_stop": { + // Block complete - no action needed for now. + // NativeToolCallParser handles tool call completion + // Note: Signature for multi-turn thinking would require using stream.finalMessage() + // after iteration completes, which requires restructuring the streaming approach. + break + } } } + } finally { + // Release the listener once the stream is consumed, whether the + // request completed, failed, or the generator was closed early. + removeExternalAbortListener?.() } } @@ -280,7 +309,7 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple } } - async completePrompt(prompt: string, _options?: CompletePromptOptions) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { try { const { id, @@ -306,7 +335,19 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple stream: false, } as Anthropic.Messages.MessageCreateParamsNonStreaming - const response = await this.client.messages.create(params) + // Build request options with abortSignal and/or timeout handling + const requestOptions: Anthropic.RequestOptions = {} + if (options?.abortSignal) { + requestOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeout = options.timeoutMs + } + + const response = await this.client.messages.create( + params, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = response.content.find(({ type }) => type === "text") return content?.type === "text" ? content.text : "" diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index 9dddc75248..6ee3aba928 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -1,6 +1,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming" import { CacheControlEphemeral } from "@anthropic-ai/sdk/resources" +import { withFinallyCleanup } from "./stream-cleanup" import { type ModelInfo, @@ -20,6 +21,7 @@ import { filterNonAnthropicBlocks } from "../transform/anthropic-filter" import { getAnthropicProviderReasoning } from "../transform/reasoning" import { BaseProvider } from "./base-provider" +import { bridgeExternalAbortSignal } from "./utils/abort-signal" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { calculateApiCostAnthropic } from "../../shared/cost" import { @@ -111,6 +113,15 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa tool_choice: toolChoice, } + // Bridge the external abort signal from request metadata into a per-request + // controller so the SDK call is cancelled when the owning request is aborted + // (or when the signal is already aborted). Without an external signal the + // client-level timeout configured in the constructor remains the only + // cancellation mechanism, preserving the existing behavior. + const externalAbortSignal = metadata?.abortSignal + const { signal: abortSignal, dispose: removeExternalAbortListener } = + bridgeExternalAbortSignal(externalAbortSignal) + switch (modelId) { case "claude-sonnet-5": case "claude-sonnet-4-6": @@ -177,41 +188,19 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa stream: true, ...nativeToolParams, } + // prompt caching: https://x.com/alexalbert__/status/1823751995901272068 + // https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers + // https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393 + // Every model that reaches this branch supports prompt caching, so the + // beta applies to all of them. + betas.push("prompt-caching-2024-07-31") + const requestOptions: Anthropic.RequestOptions = { + headers: { "anthropic-beta": betas.join(",") }, + ...(abortSignal && { signal: abortSignal }), + } stream = await this.client.messages.create( requestParams as Anthropic.Messages.MessageCreateParamsStreaming, - (() => { - // prompt caching: https://x.com/alexalbert__/status/1823751995901272068 - // https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers - // https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393 - - // Then check for models that support prompt caching - switch (modelId) { - case "claude-sonnet-5": - case "claude-sonnet-4-6": - case "claude-sonnet-4-5": - case "claude-sonnet-4-20250514": - case "claude-opus-4-6": - case "claude-opus-4-7": - case "claude-opus-4-8": - case "claude-opus-5": - case "claude-opus-5-5": - case "claude-fable-5-1": - case "claude-fable-5": - case "claude-opus-4-5-20251101": - case "claude-opus-4-1-20250805": - case "claude-opus-4-20250514": - case "claude-3-7-sonnet-20250219": - case "claude-3-5-sonnet-20241022": - case "claude-3-5-haiku-20241022": - case "claude-3-opus-20240229": - case "claude-haiku-4-5-20251001": - case "claude-3-haiku-20240307": - betas.push("prompt-caching-2024-07-31") - return { headers: { "anthropic-beta": betas.join(",") } } - default: - return undefined - } - })(), + requestOptions, ) } catch (error) { TelemetryService.instance.captureException( @@ -222,6 +211,9 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa "createMessage", ), ) + // The request failed before streaming: detach the bridged listener so + // it cannot outlive this request. + removeExternalAbortListener?.() throw error } break @@ -240,6 +232,7 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa } stream = (await this.client.messages.create( requestParams as Anthropic.Messages.MessageCreateParamsStreaming, + abortSignal ? { signal: abortSignal } : undefined, )) as any } catch (error) { TelemetryService.instance.captureException( @@ -250,6 +243,9 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa "createMessage", ), ) + // The request failed before streaming: detach the bridged listener so + // it cannot outlive this request. + removeExternalAbortListener?.() throw error } break @@ -261,7 +257,7 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa let cacheWriteTokens = 0 let cacheReadTokens = 0 - for await (const chunk of stream) { + for await (const chunk of withFinallyCleanup(stream, removeExternalAbortListener)) { switch (chunk.type) { case "message_start": { // Tells us cache reads/writes/input/output. @@ -461,19 +457,31 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa } } - async completePrompt(prompt: string, _options?: CompletePromptOptions) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { const { id: model, temperature } = this.getModel() let message try { - message = await this.client.messages.create({ - model, - max_tokens: ANTHROPIC_DEFAULT_MAX_TOKENS, - thinking: undefined, - temperature, - messages: [{ role: "user", content: prompt }], - stream: false, - }) + // Build request options with both abortSignal and timeout handling + const requestOptions: Anthropic.RequestOptions = {} + if (options?.abortSignal) { + requestOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeout = options.timeoutMs + } + + message = await this.client.messages.create( + { + model, + max_tokens: ANTHROPIC_DEFAULT_MAX_TOKENS, + thinking: undefined, + temperature, + messages: [{ role: "user", content: prompt }], + stream: false, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) } catch (error) { TelemetryService.instance.captureException( new ApiProviderError( diff --git a/src/api/providers/minimax.ts b/src/api/providers/minimax.ts index b79e96568a..f10bca8a7a 100644 --- a/src/api/providers/minimax.ts +++ b/src/api/providers/minimax.ts @@ -1,5 +1,7 @@ import { Anthropic } from "@anthropic-ai/sdk" +import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming" import { CacheControlEphemeral } from "@anthropic-ai/sdk/resources" +import { withFinallyCleanup } from "./stream-cleanup" import OpenAI from "openai" import { type MinimaxModelId, minimaxDefaultModelId, minimaxModels } from "@roo-code/types" @@ -11,6 +13,7 @@ import { getModelParams } from "../transform/model-params" import { mergeEnvironmentDetailsForMiniMax } from "../transform/minimax-format" import { BaseProvider } from "./base-provider" +import { bridgeExternalAbortSignal } from "./utils/abort-signal" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { calculateApiCostAnthropic } from "../../shared/cost" import { convertOpenAIToolsToAnthropic } from "../../core/prompts/tools/native-tools/converters" @@ -84,6 +87,15 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand const cacheControl: CacheControlEphemeral = { type: "ephemeral" } const { id: modelId, info, maxTokens, temperature } = this.getModel() + // Bridge the external abort signal from request metadata into a per-request + // controller so the SDK call is cancelled when the owning request is aborted + // (or when the signal is already aborted). Without an external signal the + // client-level timeout configured in the constructor remains the only + // cancellation mechanism, preserving the existing behavior. + const externalAbortSignal = metadata?.abortSignal + const { signal: abortSignal, dispose: removeExternalAbortListener } = + bridgeExternalAbortSignal(externalAbortSignal) + // MiniMax M2 models support prompt caching const supportsPromptCache = info.supportsPromptCache ?? false @@ -112,14 +124,22 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand tool_choice: convertOpenAIToolChoice(metadata?.tool_choice), } - const stream = await this.client.messages.create(requestParams) + let stream: AnthropicStream + try { + stream = await this.client.messages.create(requestParams, abortSignal ? { signal: abortSignal } : undefined) + } catch (error) { + // Creation failed before streaming: detach the bridged listener so it + // cannot outlive this request. + removeExternalAbortListener?.() + throw error + } let inputTokens = 0 let outputTokens = 0 let cacheWriteTokens = 0 let cacheReadTokens = 0 - for await (const chunk of stream) { + for await (const chunk of withFinallyCleanup(stream, removeExternalAbortListener)) { switch (chunk.type) { case "message_start": { // Tells us cache reads/writes/input/output. @@ -288,16 +308,28 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand } } - async completePrompt(prompt: string, _options?: CompletePromptOptions) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { const { id: model, temperature } = this.getModel() - const message = await this.client.messages.create({ - model, - max_tokens: 16_384, - temperature: temperature ?? 1.0, - messages: [{ role: "user", content: prompt }], - stream: false, - }) + // Build request options with abortSignal and/or timeout handling + const requestOptions: Anthropic.RequestOptions = {} + if (options?.abortSignal) { + requestOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeout = options.timeoutMs + } + + const message = await this.client.messages.create( + { + model, + max_tokens: 16_384, + temperature: temperature ?? 1.0, + messages: [{ role: "user", content: prompt }], + stream: false, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) const content = message.content.find(({ type }) => type === "text") return content?.type === "text" ? content.text : "" diff --git a/src/api/providers/stream-cleanup.ts b/src/api/providers/stream-cleanup.ts new file mode 100644 index 0000000000..d20168c954 --- /dev/null +++ b/src/api/providers/stream-cleanup.ts @@ -0,0 +1,19 @@ +/** + * Forward chunks from an async iterable and run the cleanup callback once when + * iteration ends: normal completion, a failure from the underlying stream, or + * an early stop by the consumer. Providers use this to detach a bridged + * external abort listener as soon as the request's stream has finished, + * keeping the cleanup out of the (existing) stream loop body. + */ +export async function* withFinallyCleanup( + stream: AsyncIterable, + cleanup: (() => void) | undefined, +): AsyncGenerator { + try { + for await (const chunk of stream) { + yield chunk + } + } finally { + cleanup?.() + } +} diff --git a/src/api/providers/utils/__tests__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index f779d48e1f..1d0bbff9d4 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -1,4 +1,5 @@ import { + bridgeExternalAbortSignal, createAbortError, isRequestAborted, mergeAbortSignalAndTimeout, @@ -275,6 +276,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", () => { @@ -416,3 +467,49 @@ describe("abort-signal utilities", () => { }) }) }) + +describe("bridgeExternalAbortSignal", () => { + it("returns an undefined signal and a noop dispose without an external signal", () => { + const { signal, dispose } = bridgeExternalAbortSignal() + expect(signal).toBeUndefined() + expect(() => dispose()).not.toThrow() + }) + + it("returns an already-aborted signal when the external signal is aborted", () => { + const controller = new AbortController() + controller.abort() + const { signal, dispose } = bridgeExternalAbortSignal(controller.signal) + expect(signal?.aborted).toBe(true) + expect(() => dispose()).not.toThrow() + }) + + it("aborts the per-request signal when the external signal aborts", () => { + const controller = new AbortController() + const { signal } = bridgeExternalAbortSignal(controller.signal) + expect(signal?.aborted).toBe(false) + controller.abort() + expect(signal?.aborted).toBe(true) + }) + + it("detaches the retained listener on dispose", () => { + const controller = new AbortController() + const addSpy = vitest.spyOn(controller.signal, "addEventListener") + const removeSpy = vitest.spyOn(controller.signal, "removeEventListener") + + const { signal, dispose } = bridgeExternalAbortSignal(controller.signal) + + expect(addSpy).toHaveBeenCalledTimes(1) + const [event, listener, opts] = addSpy.mock.calls[0] + expect(event).toBe("abort") + expect(opts).toEqual({ once: true }) + + dispose() + expect(removeSpy).toHaveBeenCalledTimes(1) + expect(removeSpy).toHaveBeenCalledWith("abort", listener) + + // After dispose the bridge no longer forwards aborts: the per-request + // signal stays live, proving the exact retained callback was detached. + controller.abort() + expect(signal?.aborted).toBe(false) + }) +}) diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index cbb7a13aaf..3f19335158 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 @@ -36,6 +42,45 @@ export function mergeAbortSignals(primarySignal: AbortSignal, secondarySignal?: return AbortSignal.any([primarySignal, secondarySignal]) } +/** + * Bridge an optional caller-owned (external) abort signal to a per-request + * signal, with an explicit detach handle. + * + * The returned `signal` aborts together with the external signal (immediately + * if it is already aborted). Callers must invoke `dispose()` once the request + * ends - on completion, failure, and early generator exit alike: the listener + * is registered on the caller's long-lived signal, and an anonymous listener + * left attached would outlive the request and retain the per-request + * controller for every completed request. + * + * With no external signal, returns `{ signal: undefined, dispose: noop }` so + * the client-level timeout remains the only cancellation mechanism. + */ +export function bridgeExternalAbortSignal(externalSignal?: AbortSignal): { + signal?: AbortSignal + dispose: () => void +} { + if (!externalSignal) { + return { signal: undefined, dispose: () => {} } + } + + const controller = new AbortController() + if (externalSignal.aborted) { + controller.abort() + return { signal: controller.signal, dispose: () => {} } + } + + // Retain the callback so the listener can be detached once this request + // ends: an anonymous listener on a long-lived external signal would + // outlive the request and retain the per-request controller. + const onExternalAbort = () => controller.abort() + externalSignal.addEventListener("abort", onExternalAbort, { once: true }) + return { + signal: controller.signal, + dispose: () => externalSignal.removeEventListener("abort", onExternalAbort), + } +} + /** * Throw an AbortError if the given signal is already aborted. * @@ -115,7 +160,6 @@ export function rejectOnAbort(pending: Promise, signal: AbortSignal, provi } return new Promise((resolve, reject) => { const onAbort = () => reject(createAbortError(providerName)) - // Stryker disable next-line ObjectLiteral,BooleanLiteral: a signal fires its abort event exactly once and the settle handler removes this listener, so the once flag is unobservable signal.addEventListener("abort", onAbort, { once: true }) void pending.then( (value) => { diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index 7dc4fa6b1b..bcf401284c 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -13,6 +13,7 @@ import { getModelParams } from "../transform/model-params" import { DEFAULT_HEADERS, NOT_PROVIDED } from "./constants" import { BaseProvider } from "./base-provider" +import { bridgeExternalAbortSignal, createAbortError, isRequestAborted } from "./utils/abort-signal" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { isMcpTool } from "../../utils/mcp-name" @@ -95,6 +96,15 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler const input = convertToResponsesApiInput(messages) const responseTools = this.mapResponseTools(metadata?.tools) + // Bridge the external abort signal from request metadata into a per-request + // controller so the SDK call is cancelled when the owning request is aborted + // (or when the signal is already aborted). Without an external signal the + // client-level timeout configured in the constructor remains the only + // cancellation mechanism, preserving the existing behavior. + const externalAbortSignal = metadata?.abortSignal + const { signal: abortSignal, dispose: removeExternalAbortListener } = + bridgeExternalAbortSignal(externalAbortSignal) + // Build request options const requestBody: Record = { model: model.id, @@ -127,36 +137,80 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler requestBody.reasoning = { effort: model.reasoning.reasoning_effort } } - let stream: AsyncIterable try { - stream = (await this.client.responses.create({ - ...requestBody, - stream: true, - } as any)) as unknown as AsyncIterable - } catch (error) { - const errorMessage = error instanceof Error ? error.message : String(error) - const apiError = new ApiProviderError(errorMessage, this.providerName, model.id, "createMessage") - TelemetryService.instance.captureException(apiError) - throw handleOpenAIError(error, this.providerName) + let stream: AsyncIterable + try { + stream = (await this.client.responses.create( + { ...requestBody, stream: true } as any, + abortSignal ? { signal: abortSignal } : undefined, + )) as unknown as AsyncIterable + } catch (error) { + // Normalize cancellations to the contract error: the OpenAI SDK's + // APIUserAbortError carries name "Error" and message "Request was + // aborted.", so rethrowing it raw would leak a non-conforming shape to + // the caller's AbortError guard. + if (isRequestAborted(error, abortSignal)) { + throw createAbortError(this.providerName) + } + const errorMessage = error instanceof Error ? error.message : String(error) + const apiError = new ApiProviderError(errorMessage, this.providerName, model.id, "createMessage") + TelemetryService.instance.captureException(apiError) + throw handleOpenAIError(error, this.providerName) + } + + const normalizeUsage = createUsageNormalizer() + try { + yield* processResponsesApiStream(stream, normalizeUsage) + } catch (error) { + // An abort landing after responses.create resolved surfaces here as + // the SDK's APIUserAbortError thrown from stream iteration (name + // "Error", message "Request was aborted."), outside the create + // catch above. Normalize it to the contract AbortError too, so + // callers never classify a cancellation as a provider failure. + if (isRequestAborted(error, abortSignal)) { + throw createAbortError(this.providerName) + } + throw error + } + } finally { + // Release the listener once the stream is consumed, whether the + // request completed, failed, or the generator was closed early. + removeExternalAbortListener?.() } - - const normalizeUsage = createUsageNormalizer() - yield* processResponsesApiStream(stream, normalizeUsage) } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const model = this.getModel() try { - const response = await this.client.responses.create({ - model: model.id, - input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }], - store: false, - }) + // Build request options with abortSignal and/or timeout handling + const requestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal) { + requestOptions.signal = options.abortSignal + } + if (options?.timeoutMs !== undefined) { + requestOptions.timeout = options.timeoutMs + } + + const response = await this.client.responses.create( + { + model: model.id, + input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }], + store: false, + }, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) // output_text is a convenience field on the Responses API response return response.output_text || "" } catch (error) { + // Normalize cancellations to the contract error: the OpenAI SDK's + // APIUserAbortError carries name "Error" and message "Request was + // aborted.", so rethrowing it raw would leak a non-conforming shape to + // the caller's AbortError guard. + if (isRequestAborted(error, options?.abortSignal)) { + throw createAbortError(this.providerName) + } const errorMessage = error instanceof Error ? error.message : String(error) const apiError = new ApiProviderError(errorMessage, this.providerName, model.id, "completePrompt") TelemetryService.instance.captureException(apiError)