diff --git a/packages/config-eslint/api-provider-hygiene.js b/packages/config-eslint/api-provider-hygiene.js index d9a560fd33..99dac3a62a 100644 --- a/packages/config-eslint/api-provider-hygiene.js +++ b/packages/config-eslint/api-provider-hygiene.js @@ -3,12 +3,17 @@ // Paths are relative to the `src` package, where this config is loaded. const sharedAbortHelpers = - "mergeAbortSignalAndTimeout|mergeAbortSignals|throwIfAborted|isRequestAborted|createAbortError|rejectOnAbort|resolveModelWithAbort" + "mergeAbortSignalAndTimeout|mergeAbortSignals|throwIfAborted|isRequestAborted|createAbortError|rejectOnAbort|resolveModelWithAbort|bridgeExternalAbortSignal|hasUsableTimeout|toSdkTimeoutMs" // Names that usually hold a secret or an endpoint that can embed credentials. // `cacheKey` and `contextTokens` must not match. const sensitiveName = "^(url|apiKey|secret|password)$|(Url|URL|ApiKey|Secret|Token)$" +// A per-request `timeoutMs` from CompletePromptOptions. `this.timeoutMs` is the client-level +// timeout from getApiRequestTimeout, which is already valid, so it must not match. +const rawRequestTimeout = + ":matches(Identifier[name='timeoutMs'], MemberExpression[property.name='timeoutMs']:not([object.type='ThisExpression']), ChainExpression[expression.property.name='timeoutMs'])" + export const apiProviderHygieneSelectors = [ { selector: `:matches(FunctionDeclaration, VariableDeclarator)[id.name=/^(${sharedAbortHelpers})$/]`, @@ -20,6 +25,16 @@ export const apiProviderHygieneSelectors = [ message: "Use mergeAbortSignalAndTimeout for a variable delay. AbortSignal.timeout throws on Infinity and fractions, and clamps large values to 1 ms.", }, + { + selector: `:matches(Property[key.name='timeout'], AssignmentExpression[left.property.name='timeout']) > ${rawRequestTimeout}`, + message: + "Do not forward a request timeoutMs as is. A SDK treats 0 as an immediate abort and throws on negative, non-finite or fractional values. Use hasUsableTimeout and toSdkTimeoutMs.", + }, + { + selector: "Literal[value=2147483647]", + message: + "Do not write the maximum timer delay by hand. Use MAX_TIMEOUT_MS or toSdkTimeoutMs from providers/utils/abort-signal.", + }, { selector: `CallExpression[callee.object.name='console'] TemplateLiteral > :matches(Identifier[name=/${sensitiveName}/], MemberExpression[property.name=/${sensitiveName}/])`, message: @@ -35,3 +50,12 @@ export const apiProviderHygieneConfig = { "no-restricted-syntax": ["error", ...apiProviderHygieneSelectors], }, } + +// Spec files under src/api. A loop of expects in one test stops at the first failing case and hides the rest. +export const apiProviderTestSelectors = [ + { + selector: "ForOfStatement[right.type='ArrayExpression']:has(CallExpression[callee.name='expect'])", + message: + "Use it.each for repeated cases. A loop of expects in one test stops at the first failing case and does not name it.", + }, +] diff --git a/packages/config-eslint/api-provider-hygiene.test.js b/packages/config-eslint/api-provider-hygiene.test.js index ae981a0d05..616ce32730 100644 --- a/packages/config-eslint/api-provider-hygiene.test.js +++ b/packages/config-eslint/api-provider-hygiene.test.js @@ -4,7 +4,7 @@ import { describe, it } from "node:test" import { Linter } from "eslint" import typescriptEslint from "typescript-eslint" -import { apiProviderHygieneConfig } from "./api-provider-hygiene.js" +import { apiProviderHygieneConfig, apiProviderTestSelectors } from "./api-provider-hygiene.js" const linter = new Linter({ configType: "flat" }) const config = [{ files: ["**/*.ts"], languageOptions: { parser: typescriptEslint.parser } }, apiProviderHygieneConfig] @@ -88,3 +88,89 @@ describe("api provider hygiene: sensitive values in logs", () => { }) } }) + +describe("api provider hygiene: request timeout forwarding", () => { + for (const code of [ + "create(body, { timeout: options.timeoutMs })", + "create(body, { timeout: options?.timeoutMs })", + "create(body, { signal, timeout: timeoutMs })", + "requestOptions.timeout = options.timeoutMs", + "requestOptions.timeout = timeoutMs", + ]) { + it(`flags: ${code}`, () => { + assert.equal(lint(code).length, 1) + }) + } + + for (const code of [ + "new OpenAI({ timeout: this.timeoutMs })", + "requestOptions.timeout = toSdkTimeoutMs(timeoutMs)", + "create(body, { timeout: toSdkTimeoutMs(options.timeoutMs) })", + "if (hasUsableTimeout(timeoutMs)) requestOptions.timeout = toSdkTimeoutMs(timeoutMs)", + "const timeoutMs = options?.timeoutMs", + "mergeAbortSignalAndTimeout(signal, options?.timeoutMs)", + ]) { + it(`allows: ${code}`, () => { + assert.deepEqual(lint(code), []) + }) + } +}) + +describe("api provider hygiene: timeout helpers", () => { + for (const code of [ + "const MAX = 2147483647", + "Math.min(Math.ceil(timeoutMs), 2_147_483_647)", + "function hasUsableTimeout(t?: number) { return true }", + "const toSdkTimeoutMs = (t: number) => t", + "function bridgeExternalAbortSignal() {}", + ]) { + it(`flags: ${code}`, () => { + assert.equal(lint(code).length, 1) + }) + } + + it("allows importing the timeout helpers", () => { + assert.deepEqual( + lint('import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal"\ntoSdkTimeoutMs(1)'), + [], + ) + }) + + it("allows the literal in the helper module", () => { + assert.deepEqual(lint("export const MAX_TIMEOUT_MS = 2_147_483_647", "api/providers/utils/abort-signal.ts"), []) + }) +}) + +describe("api provider hygiene: spec loops", () => { + const testLinter = new Linter({ configType: "flat" }) + const testConfig = [ + { files: ["**/*.ts"], languageOptions: { parser: typescriptEslint.parser } }, + { + files: ["api/**/__tests__/**/*.spec.ts"], + rules: { "no-restricted-syntax": ["error", ...apiProviderTestSelectors] }, + }, + ] + const lintSpec = (code) => + testLinter + .verify(code, testConfig, { filename: "api/providers/__tests__/example.spec.ts" }) + .map((m) => m.message) + + for (const code of [ + "it('x', () => { for (const v of [0, -5]) { expect(f(v)).toBe(1) } })", + "it('x', async () => { for (const v of [0, -5]) { await run(v); expect(f(v)).toBe(1) } })", + ]) { + it(`flags a loop of expects over a literal array: ${code}`, () => { + assert.equal(lintSpec(code).length, 1) + }) + } + + for (const code of [ + "it.each([0, -5])('x %s', (v) => { expect(f(v)).toBe(1) })", + "it('x', () => { for (const v of [0, -5]) { run(v) } })", + "it('x', () => { for (const v of cases) { expect(f(v)).toBe(1) } })", + ]) { + it(`allows: ${code}`, () => { + assert.deepEqual(lintSpec(code), []) + }) + } +}) diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts index 339fa36a07..31f26f00ce 100644 --- a/src/api/providers/__tests__/anthropic-vertex.spec.ts +++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts @@ -7,7 +7,9 @@ import { GoogleAuth } from "google-auth-library" import { VERTEX_1M_CONTEXT_MODEL_IDS } from "@roo-code/types" import { AnthropicVertexHandler } from "../anthropic-vertex" +import { MAX_TIMEOUT_MS } from "../utils/abort-signal" 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 +748,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 +952,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 +1018,104 @@ 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.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + // Only a call through toSdkTimeoutMs rounds this up. + { label: "a fractional", input: 12.5, forwarded: 13 }, + // Only a call through toSdkTimeoutMs caps at the Node timer maximum. + { label: "an over-cap", input: MAX_TIMEOUT_MS + 1, forwarded: MAX_TIMEOUT_MS }, + ])("completePrompt should forward $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + const mockCreate = vitest + .spyOn(handler["client"].messages, "create") + .mockResolvedValue({ content: [{ type: "text", text: "response" }] } as never) + + await handler.completePrompt("test prompt", { timeoutMs: input }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + expect.objectContaining({ timeout: forwarded }), + ) + }) + + // The SDK turns timeout: 0 into an immediate abort and rejects + // negative/non-finite delays, so "no explicit timeout" means no option. + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "completePrompt should not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + 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 }) + expect(mockCreate).toHaveBeenCalledWith(expect.any(Object), undefined) + }, + ) + 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..b80e9887b7 100644 --- a/src/api/providers/__tests__/anthropic.spec.ts +++ b/src/api/providers/__tests__/anthropic.spec.ts @@ -1,8 +1,10 @@ // npx vitest run src/api/providers/__tests__/anthropic.spec.ts import { AnthropicHandler } from "../anthropic" +import { MAX_TIMEOUT_MS } from "../utils/abort-signal" 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 +769,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 +1058,115 @@ 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) + }) + + // The Anthropic SDK turns timeout: 0 into an immediate abort and throws + // on negative/non-finite delays, so "no explicit timeout" must mean the + // option is absent rather than 0. + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "should not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockCreate.mockClear() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { timeoutMs }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: mockOptions.apiModelId }), + undefined, + ) + }, + ) + + it.each([ + // Only a call through toSdkTimeoutMs rounds this up. + { label: "a fractional", input: 12.5, forwarded: 13 }, + // Only a call through toSdkTimeoutMs caps at the Node timer maximum. + { label: "an over-cap", input: MAX_TIMEOUT_MS + 1, forwarded: MAX_TIMEOUT_MS }, + ])("should forward $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { timeoutMs: input }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: mockOptions.apiModelId }), { + timeout: forwarded, + }) + }) }) describe("getModel", () => { diff --git a/src/api/providers/__tests__/bedrock.spec.ts b/src/api/providers/__tests__/bedrock.spec.ts index fcc98b87b1..0d8e5a784a 100644 --- a/src/api/providers/__tests__/bedrock.spec.ts +++ b/src/api/providers/__tests__/bedrock.spec.ts @@ -138,37 +138,49 @@ describe("AwsBedrockHandler", () => { }, ) - it.each([ - ["anthropic.claude-sonnet-4-5-20250929-v1:0", 64_000], - ["anthropic.claude-sonnet-4-6", 64_000], - ["anthropic.claude-sonnet-5", 128_000], - ["anthropic.claude-opus-4-7", 128_000], - ["anthropic.claude-opus-4-8", 128_000], - ["anthropic.claude-opus-5", 128_000], - ["anthropic.claude-opus-5-5", 128_000], - ["anthropic.claude-haiku-5-5", 128_000], - ] as const)("exposes %s output ceiling without raising request defaults", async (apiModelId, ceiling) => { - for (const enableReasoningEffort of [true, false]) { - for (const modelMaxTokens of [undefined, ceiling]) { - const provider = new AwsBedrockHandler({ apiModelId, enableReasoningEffort, modelMaxTokens }) - const modelInfo = provider.getModel().info - expect(modelInfo.maxTokens).toBe(ceiling) - provider["client"].send = vi.fn().mockResolvedValue({ stream: asyncStreamFrom([]) }) - await collectStream(provider.createMessage("system", [{ role: "user", content: "hello" }])) - const expected = - modelMaxTokens && (enableReasoningEffort || modelInfo.supportsMaxTokens) - ? ceiling - : enableReasoningEffort - ? 16_384 - : 8192 - expect(mockConverseStreamCommand).toHaveBeenLastCalledWith( - expect.objectContaining({ - inferenceConfig: expect.objectContaining({ maxTokens: expected }), - }), - ) - } - } - }) + it.each( + ( + [ + ["anthropic.claude-sonnet-4-5-20250929-v1:0", 64_000], + ["anthropic.claude-sonnet-4-6", 64_000], + ["anthropic.claude-sonnet-5", 128_000], + ["anthropic.claude-opus-4-7", 128_000], + ["anthropic.claude-opus-4-8", 128_000], + ["anthropic.claude-opus-5", 128_000], + ["anthropic.claude-opus-5-5", 128_000], + ["anthropic.claude-haiku-5-5", 128_000], + ] as const + ).flatMap(([apiModelId, ceiling]) => + [true, false].flatMap((enableReasoningEffort) => + [undefined, ceiling].map((modelMaxTokens) => ({ + apiModelId, + ceiling, + enableReasoningEffort, + modelMaxTokens, + })), + ), + ), + )( + "exposes $apiModelId output ceiling without raising request defaults (reasoning effort: $enableReasoningEffort, modelMaxTokens: $modelMaxTokens)", + async ({ apiModelId, ceiling, enableReasoningEffort, modelMaxTokens }) => { + const provider = new AwsBedrockHandler({ apiModelId, enableReasoningEffort, modelMaxTokens }) + const modelInfo = provider.getModel().info + expect(modelInfo.maxTokens).toBe(ceiling) + provider["client"].send = vi.fn().mockResolvedValue({ stream: asyncStreamFrom([]) }) + await collectStream(provider.createMessage("system", [{ role: "user", content: "hello" }])) + const expected = + modelMaxTokens && (enableReasoningEffort || modelInfo.supportsMaxTokens) + ? ceiling + : enableReasoningEffort + ? 16_384 + : 8192 + expect(mockConverseStreamCommand).toHaveBeenLastCalledWith( + expect.objectContaining({ + inferenceConfig: expect.objectContaining({ maxTokens: expected }), + }), + ) + }, + ) describe("getModel", () => { it("should return the correct model info for a standard model", () => { diff --git a/src/api/providers/__tests__/friendli.spec.ts b/src/api/providers/__tests__/friendli.spec.ts index 6ac7d9c1fe..acf03687f3 100644 --- a/src/api/providers/__tests__/friendli.spec.ts +++ b/src/api/providers/__tests__/friendli.spec.ts @@ -181,6 +181,32 @@ describe("FriendliHandler", () => { expect(result).toBe(expectedResponse) }) + it("completePrompt should forward the abort signal", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + const { signal } = new AbortController() + await handler.completePrompt("test prompt", { abortSignal: signal }) + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal }) + }) + + it.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + { label: "a fractional", input: 12.5, forwarded: 13 }, + { label: "an over-cap", input: 2_147_483_648, forwarded: 2_147_483_647 }, + ])("completePrompt should forward $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + await handler.completePrompt("test prompt", { timeoutMs: input }) + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { timeout: forwarded }) + }) + + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "completePrompt should not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + await handler.completePrompt("test prompt", { timeoutMs }) + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), undefined) + }, + ) + it("should handle errors in completePrompt", async () => { const errorMessage = "Friendli API error" mockCreate.mockRejectedValueOnce(new Error(errorMessage)) diff --git a/src/api/providers/__tests__/io-intelligence.spec.ts b/src/api/providers/__tests__/io-intelligence.spec.ts index 11166a1533..28ce7010b7 100644 --- a/src/api/providers/__tests__/io-intelligence.spec.ts +++ b/src/api/providers/__tests__/io-intelligence.spec.ts @@ -364,6 +364,27 @@ describe("IOIntelligenceHandler", () => { expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ stream: false }), { signal, timeout: 5_000 }) }) + it.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + { label: "a fractional", input: 12.5, forwarded: 13 }, + { label: "an over-cap", input: 2_147_483_648, forwarded: 2_147_483_647 }, + ])("completePrompt forwards $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + mockCreate.mockResolvedValue({ choices: [{ message: { role: "assistant", content: "ok" } }] }) + const handler = new IOIntelligenceHandler({ ioIntelligenceModelId: "meta-llama/Llama-3.3-70B-Instruct" }) + await handler.completePrompt("ping", { timeoutMs: input }) + expect(mockCreate.mock.calls[0][1]).toEqual({ timeout: forwarded }) + }) + + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "completePrompt does not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockCreate.mockResolvedValue({ choices: [{ message: { role: "assistant", content: "ok" } }] }) + const handler = new IOIntelligenceHandler({ ioIntelligenceModelId: "meta-llama/Llama-3.3-70B-Instruct" }) + await handler.completePrompt("ping", { timeoutMs }) + expect(mockCreate.mock.calls[0][1]).toBeUndefined() + }, + ) + it("does not send the OpenAI Compatible custom headers to the fixed endpoint", () => { new IOIntelligenceHandler({ ioIntelligenceApiKey: "test-key", diff --git a/src/api/providers/__tests__/minimax.spec.ts b/src/api/providers/__tests__/minimax.spec.ts index 063932b303..5d87ef3de8 100644 --- a/src/api/providers/__tests__/minimax.spec.ts +++ b/src/api/providers/__tests__/minimax.spec.ts @@ -13,7 +13,9 @@ 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 { MAX_TIMEOUT_MS } from "../utils/abort-signal" +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 +241,96 @@ 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.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + // Only a call through toSdkTimeoutMs rounds this up. + { label: "a fractional", input: 12.5, forwarded: 13 }, + // Only a call through toSdkTimeoutMs caps at the Node timer maximum. + { label: "an over-cap", input: MAX_TIMEOUT_MS + 1, forwarded: MAX_TIMEOUT_MS }, + ])("should forward $label timeoutMs as $forwarded when no signal provided", async ({ input, forwarded }) => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "response" }], + }) + await handler.completePrompt("test prompt", { timeoutMs: input }) + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeout: forwarded, + }) + }) + + // "No explicit timeout" must leave the option absent: timeout: 0 aborts + // immediately and negative/non-finite delays throw inside the SDK. + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "should not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockCreate.mockClear() + mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] }) + await handler.completePrompt("test prompt", { timeoutMs }) + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, + ) + }, + ) + + 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 +398,7 @@ describe("MiniMaxHandler", () => { messages: expect.any(Array), stream: true, }), + undefined, ) }) @@ -319,6 +412,7 @@ describe("MiniMaxHandler", () => { expect.objectContaining({ temperature: 1, }), + undefined, ) }) @@ -375,6 +469,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__/nanogpt.spec.ts b/src/api/providers/__tests__/nanogpt.spec.ts index 0173e2c2a9..0a13f779c9 100644 --- a/src/api/providers/__tests__/nanogpt.spec.ts +++ b/src/api/providers/__tests__/nanogpt.spec.ts @@ -773,6 +773,42 @@ describe("NanoGptHandler", () => { expect(await handler.completePrompt("prompt")).toBe("") }) + it("forwards no request options without a signal or timeout", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await new NanoGptHandler({ nanoGptModelId: "model:thinking" }).completePrompt("prompt") + expect(mockCreate.mock.calls[0][1]).toBeUndefined() + }) + + it("forwards the abort signal", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + const { signal } = new AbortController() + await new NanoGptHandler({ nanoGptModelId: "model:thinking" }).completePrompt("prompt", { + abortSignal: signal, + }) + expect(mockCreate.mock.calls[0][1]).toEqual({ signal }) + }) + + it.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + { label: "a fractional", input: 12.5, forwarded: 13 }, + { label: "an over-cap", input: 2_147_483_648, forwarded: 2_147_483_647 }, + ])("forwards $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await new NanoGptHandler({ nanoGptModelId: "model:thinking" }).completePrompt("prompt", { + timeoutMs: input, + }) + expect(mockCreate.mock.calls[0][1]).toEqual({ timeout: forwarded }) + }) + + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "does not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] }) + await new NanoGptHandler({ nanoGptModelId: "model:thinking" }).completePrompt("prompt", { timeoutMs }) + expect(mockCreate.mock.calls[0][1]).toBeUndefined() + }, + ) + it("preserves completion error metadata without leaking the API key", async () => { const errorDetails = [{ "@type": "type.googleapis.com/google.rpc.RetryInfo", retryDelay: "20s" }] mockCreate.mockRejectedValue( 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__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 06681fe374..c6926e3f3d 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -484,65 +484,66 @@ describe("VsCodeLmHandler", () => { }) describe("offered tool name recovery", () => { - it.each([ - { - label: "encoded surrogate", - rawName: "read\uD800file", - declaredName: "read_uD800file", - leakedName: "read_uD800file", - expectedName: "read\uD800file", - }, - { - label: "encoded marker", - rawName: "get_uuid", - declaredName: "get_uuuid", - leakedName: "get_uuuid", - expectedName: "get_uuid", - }, - { - label: "raw marker", - rawName: "get_uuid", - declaredName: "get_uuuid", - leakedName: "get_uuid", - expectedName: "get_uuid", - }, - ])("recovers $label with its original schema and split invariance", async (fixture) => { + it.each( + [ + { + label: "encoded surrogate", + rawName: "read\uD800file", + declaredName: "read_uD800file", + leakedName: "read_uD800file", + expectedName: "read\uD800file", + }, + { + label: "encoded marker", + rawName: "get_uuid", + declaredName: "get_uuuid", + leakedName: "get_uuuid", + expectedName: "get_uuid", + }, + { + label: "raw marker", + rawName: "get_uuid", + declaredName: "get_uuuid", + leakedName: "get_uuid", + expectedName: "get_uuid", + }, + ].flatMap((fixture) => [false, true].map((split) => ({ ...fixture, split }))), + )("recovers $label with its original schema (split stream: $split)", async (fixture) => { const text = `` + '7' - const split = text.indexOf(fixture.leakedName) + 5 - for (const parts of [[text], [text.slice(0, split), text.slice(split)]]) { - streamTextParts(parts) - const chunks = await collectStream( - handler.createMessage("system", [{ role: "user", content: "hi" }], { - taskId: "test-task", - tools: [ - { - type: "function", - function: { - name: fixture.rawName, - parameters: { type: "object", properties: { count: { type: "number" } } }, - }, + const splitAt = text.indexOf(fixture.leakedName) + 5 + const parts = fixture.split ? [text.slice(0, splitAt), text.slice(splitAt)] : [text] + streamTextParts(parts) + const chunks = await collectStream( + handler.createMessage("system", [{ role: "user", content: "hi" }], { + taskId: "test-task", + tools: [ + { + type: "function", + function: { + name: fixture.rawName, + parameters: { type: "object", properties: { count: { type: "number" } } }, }, - ], - }), - ) - expect(mockLanguageModelChat.sendRequest).toHaveBeenLastCalledWith( - expect.any(Array), - expect.objectContaining({ - tools: [expect.objectContaining({ name: fixture.declaredName })], - }), - expect.anything(), - ) - expect(chunks.filter((chunk) => chunk.type !== "usage")).toEqual([ - { - type: "tool_call", - id: expect.stringMatching(/^vscodelm-salvaged-\d+-0$/), - name: fixture.expectedName, - arguments: '{"count":7}', - }, - ]) - } + }, + ], + }), + ) + expect(mockLanguageModelChat.sendRequest).toHaveBeenLastCalledWith( + expect.any(Array), + expect.objectContaining({ + tools: [expect.objectContaining({ name: fixture.declaredName })], + }), + expect.anything(), + ) + expect(chunks.filter((chunk) => chunk.type !== "usage")).toEqual([ + { + type: "tool_call", + id: expect.stringMatching(/^vscodelm-salvaged-\d+-0$/), + name: fixture.expectedName, + arguments: '{"count":7}', + }, + ]) }) it.each([false, true])( @@ -4172,12 +4173,8 @@ describe("context-window tool_result truncation", () => { expect(middleOutTruncate(text, 200)).toBe(text) }) - it("stays within a limit too small to hold the truncation marker", () => { - const text = "A".repeat(500) - for (const limit of [1, 5, 50]) { - const result = middleOutTruncate(text, limit) - expect(result).toBe("A".repeat(limit)) - } + it.each([1, 5, 50])("stays within a limit too small to hold the truncation marker (%s)", (limit) => { + expect(middleOutTruncate("A".repeat(500), limit)).toBe("A".repeat(limit)) }) it("does not end a marker-less truncation on a lone high surrogate", () => { diff --git a/src/api/providers/__tests__/xai.spec.ts b/src/api/providers/__tests__/xai.spec.ts index d0524ff0e4..7d66f6985d 100644 --- a/src/api/providers/__tests__/xai.spec.ts +++ b/src/api/providers/__tests__/xai.spec.ts @@ -14,15 +14,23 @@ 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 { MAX_TIMEOUT_MS } from "../utils/abort-signal" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" describe("XAIHandler", () => { @@ -97,6 +105,7 @@ describe("XAIHandler", () => { store: false, include: ["reasoning.encrypted_content"], }), + undefined, ) }) @@ -226,9 +235,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 +312,115 @@ 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.each([ + { label: "a whole-number", input: 3000, forwarded: 3000 }, + // Only a call through toSdkTimeoutMs rounds this up. + { label: "a fractional", input: 12.5, forwarded: 13 }, + // Only a call through toSdkTimeoutMs caps at the Node timer maximum. + { label: "an over-cap", input: MAX_TIMEOUT_MS + 1, forwarded: MAX_TIMEOUT_MS }, + ])( + "completePrompt should forward $label timeoutMs as $forwarded when no signal provided", + async ({ input, forwarded }) => { + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + + await handler.completePrompt("test prompt", { timeoutMs: input }) + expect(mockResponsesCreate).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), { + timeout: forwarded, + }) + }, + ) + + // "No explicit timeout" must leave the option absent: timeout: 0 aborts + // immediately and negative/non-finite delays throw inside the SDK. + it.each([0, -5, Number.NaN, Number.POSITIVE_INFINITY])( + "completePrompt should not forward an unusable timeoutMs (%s)", + async (timeoutMs) => { + mockResponsesCreate.mockClear() + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + await handler.completePrompt("test prompt", { timeoutMs }) + expect(mockResponsesCreate).toHaveBeenCalledWith( + expect.objectContaining({ model: expect.any(String) }), + undefined, + ) + }, + ) + it("should include reasoning effort for mini models in Responses API format", async () => { const miniModelHandler = new XAIHandler({ apiModelId: "grok-3-mini", @@ -263,6 +438,7 @@ describe("XAIHandler", () => { effort: "high", }), }), + undefined, ) }) @@ -283,6 +459,7 @@ describe("XAIHandler", () => { effort: "high", }), }), + undefined, ) }) @@ -304,6 +481,7 @@ describe("XAIHandler", () => { effort: "low", }), }), + undefined, ) }) @@ -329,4 +507,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..87da090346 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, hasUsableTimeout, toSdkTimeoutMs } 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,23 @@ 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 + } + const timeoutMs = options?.timeoutMs + // Only a usable finite positive timeout may be forwarded: the SDKs treat + // timeout: 0 as an immediate abort and throw on negative/non-finite values. + // Fractional delays are rounded up, matching mergeAbortSignalAndTimeout. + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(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..f00b49cd08 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, hasUsableTimeout, toSdkTimeoutMs } 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,35 @@ 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 + } + const timeoutMs = options?.timeoutMs + // Only a usable finite positive timeout may be forwarded: the SDKs treat + // timeout: 0 as an immediate abort and throw on negative/non-finite values. + // Fractional delays are rounded up, matching mergeAbortSignalAndTimeout. + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(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/friendli.ts b/src/api/providers/friendli.ts index fc4f66e594..86d57a9d2d 100644 --- a/src/api/providers/friendli.ts +++ b/src/api/providers/friendli.ts @@ -11,6 +11,7 @@ import { getModelParams } from "../transform/model-params" import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider" import { handleOpenAIError } from "./utils/error-handler" +import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal" import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** @@ -193,13 +194,19 @@ export class FriendliHandler extends BaseOpenAiCompatibleProvider 0 ? requestOptions : undefined, )) as OpenAI.Chat.Completions.ChatCompletion return response.choices?.[0]?.message.content || "" } catch (error) { diff --git a/src/api/providers/io-intelligence.ts b/src/api/providers/io-intelligence.ts index 0c133d94fa..f5cd7470a0 100644 --- a/src/api/providers/io-intelligence.ts +++ b/src/api/providers/io-intelligence.ts @@ -23,6 +23,7 @@ import { RouterProvider } from "./router-provider" import { getIOIntelligenceModels } from "./fetchers/io-intelligence" import { handleProviderError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal" /** * IO Intelligence (io.net) provider. @@ -201,8 +202,9 @@ export class IOIntelligenceHandler extends RouterProvider implements SingleCompl if (options?.abortSignal !== undefined) { completionOptions.signal = options.abortSignal } - if (options?.timeoutMs !== undefined) { - completionOptions.timeout = options.timeoutMs + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + completionOptions.timeout = toSdkTimeoutMs(timeoutMs) } const hasCompletionOptions = completionOptions.signal !== undefined || completionOptions.timeout !== undefined diff --git a/src/api/providers/minimax.ts b/src/api/providers/minimax.ts index b79e96568a..52c42a4138 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, hasUsableTimeout, toSdkTimeoutMs } 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,32 @@ 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 + } + const timeoutMs = options?.timeoutMs + // Only a usable finite positive timeout may be forwarded: the SDKs treat + // timeout: 0 as an immediate abort and throw on negative/non-finite values. + // Fractional delays are rounded up, matching mergeAbortSignalAndTimeout. + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(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/nanogpt.ts b/src/api/providers/nanogpt.ts index 2dd2236ad6..58cb29ec6e 100644 --- a/src/api/providers/nanogpt.ts +++ b/src/api/providers/nanogpt.ts @@ -27,6 +27,7 @@ import { RouterProvider } from "./router-provider" import { getNanoGptModels } from "./fetchers/nanogpt" import { handleProviderError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal" type NanoGptUsage = OpenAI.CompletionUsage & { cache_read_input_tokens?: number @@ -228,10 +229,20 @@ export class NanoGptHandler extends RouterProvider implements SingleCompletionHa } try { - const response = await this.client.chat.completions.create(body, { - signal: options?.abortSignal, - timeout: options?.timeoutMs, - }) + // The OpenAI SDK rejects a present-but-undefined or unusable `timeout`, so + // only forward a real value, rounded and capped like mergeAbortSignalAndTimeout. + const requestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + requestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + const response = await this.client.chat.completions.create( + body, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) return response.choices[0]?.message.content ?? "" } catch (error) { throw this.createSafeError("completion", error) diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index 7e17ad17ec..86229b4918 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -6,7 +6,13 @@ import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" import type { ApiHandlerOptions, GetModelsOptions } from "../../shared/api" import { getOllamaModels, isSecureOllamaEndpoint } from "./fetchers/ollama" -import { createAbortError, mergeAbortSignalAndTimeout, rejectOnAbort } from "./utils/abort-signal" +import { + createAbortError, + hasUsableTimeout, + mergeAbortSignalAndTimeout, + rejectOnAbort, + toSdkTimeoutMs, +} from "./utils/abort-signal" import { getLoggableOrigin } from "./utils/endpoint-origin" import { TagMatcher } from "../../utils/tag-matcher" import { isProxyEnabled } from "../../utils/networkProxy" @@ -688,19 +694,13 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio // Handle timeoutMs if provided (the client exists, so the timer can // abort both the per-request transport and the client) const timeoutMs = options?.timeoutMs - if (timeoutMs !== undefined && Number.isFinite(timeoutMs) && timeoutMs > 0) { - // Node clamps delays above 2^31-1 ms to 1 ms (and coerces - // Infinity the same way), which would abort almost immediately, - // and truncates fractional delays; round up and cap so the - // semantics match mergeAbortSignalAndTimeout, and treat - // non-finite values as "no timeout". - timeoutId = setTimeout( - () => { - requestController.abort() - client.abort() - }, - Math.min(Math.ceil(timeoutMs), 2_147_483_647), - ) + if (hasUsableTimeout(timeoutMs)) { + // toSdkTimeoutMs rounds up and caps the delay (Node clamps larger + // values to 1 ms), matching mergeAbortSignalAndTimeout. + timeoutId = setTimeout(() => { + requestController.abort() + client.abort() + }, toSdkTimeoutMs(timeoutMs)) } // Propagate abortSignal into the per-request client via client.abort() 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 3cd64e8fd1..fc65239db2 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -1,11 +1,15 @@ import { + bridgeExternalAbortSignal, createAbortError, isRequestAborted, mergeAbortSignalAndTimeout, + hasUsableTimeout, mergeAbortSignals, rejectOnAbort, resolveModelWithAbort, throwIfAborted, + MAX_TIMEOUT_MS, + toSdkTimeoutMs, } from "../abort-signal" import { withSettleGuard } from "../../../../test-utils/settle-guard" @@ -466,3 +470,78 @@ 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) + }) +}) + +describe("hasUsableTimeout", () => { + it("accepts only finite positive values", () => { + expect(hasUsableTimeout(1)).toBe(true) + expect(hasUsableTimeout(50)).toBe(true) + expect(hasUsableTimeout(12.5)).toBe(true) + }) + + it("rejects values the SDK timeout option cannot take", () => { + expect(hasUsableTimeout(undefined)).toBe(false) + expect(hasUsableTimeout(0)).toBe(false) + expect(hasUsableTimeout(-5)).toBe(false) + expect(hasUsableTimeout(Number.NaN)).toBe(false) + expect(hasUsableTimeout(Number.POSITIVE_INFINITY)).toBe(false) + }) +}) + +describe("toSdkTimeoutMs", () => { + it("rounds fractional delays up", () => { + expect(toSdkTimeoutMs(1500.4)).toBe(1501) + expect(toSdkTimeoutMs(1)).toBe(1) + }) + + it("caps delays above the maximum Node timer delay", () => { + expect(toSdkTimeoutMs(MAX_TIMEOUT_MS)).toBe(MAX_TIMEOUT_MS) + expect(toSdkTimeoutMs(MAX_TIMEOUT_MS + 1)).toBe(MAX_TIMEOUT_MS) + expect(toSdkTimeoutMs(Number.MAX_SAFE_INTEGER)).toBe(MAX_TIMEOUT_MS) + }) +}) diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index e1ab5dc7c0..c26bce4f6f 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -1,3 +1,28 @@ +/** + * Whether a timeout value is usable as an SDK `timeout` option: a finite + * number greater than zero. + * + * The value `0` means "no explicit timeout" across this series (the SDKs turn + * `timeout: 0` into an immediate abort), and negative, `NaN`, `Infinity` and + * fractional values make the SDK's own validation throw — so none of them may + * be forwarded. + */ +export function hasUsableTimeout(timeoutMs?: number): timeoutMs is number { + return typeof timeoutMs === "number" && Number.isFinite(timeoutMs) && timeoutMs > 0 +} + +/** Maximum delay a Node timer (and therefore an SDK timeout option) accepts. */ +export const MAX_TIMEOUT_MS = 2_147_483_647 + +/** + * Convert a validated timeout to the value forwarded to an SDK or Node timer: + * round fractions up (uint32 validation rejects them) and cap at the maximum + * delay (Node clamps larger values to 1 ms otherwise). + */ +export function toSdkTimeoutMs(timeoutMs: number): number { + return Math.min(Math.max(1, Math.ceil(timeoutMs)), MAX_TIMEOUT_MS) +} + /** * Merge an optional external abort signal with an optional timeout. * @@ -12,13 +37,13 @@ export function mergeAbortSignalAndTimeout(externalSignal?: AbortSignal, timeout // 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 + const hasTimeout = hasUsableTimeout(timeoutMs) if (!hasTimeout) { return externalSignal } - const timeoutSignal = AbortSignal.timeout(Math.min(Math.max(1, Math.ceil(timeoutMs)), 2_147_483_647)) + const timeoutSignal = AbortSignal.timeout(toSdkTimeoutMs(timeoutMs)) if (!externalSignal) { return timeoutSignal @@ -42,6 +67,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. * @@ -110,7 +174,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..f09f3df252 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -13,6 +13,13 @@ import { getModelParams } from "../transform/model-params" import { DEFAULT_HEADERS, NOT_PROVIDED } from "./constants" import { BaseProvider } from "./base-provider" +import { + bridgeExternalAbortSignal, + createAbortError, + hasUsableTimeout, + toSdkTimeoutMs, + 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 +102,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 +143,84 @@ 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 + } + const timeoutMs = options?.timeoutMs + // Only a usable finite positive timeout may be forwarded: the SDKs treat + // timeout: 0 as an immediate abort and throw on negative/non-finite values. + // Fractional delays are rounded up, matching mergeAbortSignalAndTimeout. + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(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) diff --git a/src/api/transform/__tests__/vscode-lm-format.spec.ts b/src/api/transform/__tests__/vscode-lm-format.spec.ts index f3dc99a4a7..e66e3f869f 100644 --- a/src/api/transform/__tests__/vscode-lm-format.spec.ts +++ b/src/api/transform/__tests__/vscode-lm-format.spec.ts @@ -936,20 +936,22 @@ describe("convertToVsCodeLmMessages surrogate-safe identifiers", () => { expect(codeUnits(sanitizeToolNameSurrogates(`read-${VALID_PAIR}`))).toEqual(codeUnits(`read-${VALID_PAIR}`)) }) - it("leaves a valid tool name untouched so it still matches its registry entry", () => { - for (const name of ["get_user", "read_file", "apply_diff", "update_todo_list"]) { + it.each(["get_user", "read_file", "apply_diff", "update_todo_list"])( + "leaves the valid tool name %s untouched so it still matches its registry entry", + (name) => { expect(sanitizeToolNameSurrogates(name)).toBe(name) expect(decodeToolNameSurrogates(sanitizeToolNameSurrogates(name))).toBe(name) - } - }) + }, + ) - it("round-trips encoded tool names back to the original registry name", () => { - for (const name of [`read${LONE_HIGH}file`, `read${LONE_LOW}file`, "a_uu_b", "x_uD800y", "get_user"]) { + it.each([`read${LONE_HIGH}file`, `read${LONE_LOW}file`, "a_uu_b", "x_uD800y", "get_user"])( + "round-trips the encoded tool name %j back to the original registry name", + (name) => { const declared = sanitizeToolNameSurrogates(name) expect(declared).toMatch(/^[\w-]+$/) expect(codeUnits(decodeToolNameSurrogates(declared))).toEqual(codeUnits(name)) - } - }) + }, + ) it("does not decode a non-marker u+hex sequence", () => { // "queue1234" has "u" followed by hex but no preceding "_", so it must survive untouched. diff --git a/src/eslint.config.mjs b/src/eslint.config.mjs index f05cbeb284..a5286ca0e6 100644 --- a/src/eslint.config.mjs +++ b/src/eslint.config.mjs @@ -1,8 +1,28 @@ import { config } from "@roo-code/config-eslint/base" -import { apiProviderHygieneConfig } from "@roo-code/config-eslint/api-provider-hygiene" +import { apiProviderHygieneConfig, apiProviderTestSelectors } from "@roo-code/config-eslint/api-provider-hygiene" import { createProviderIdentifierConfig } from "@roo-code/config-eslint/provider-identifiers" import { providerIdentifiers, retiredProviderIdentifiers } from "@roo-code/types/provider-identifiers" +// Tests must not read files 3+ levels above the test file (see the test override below). +const testPathTraversalSelectors = [ + { + // Catches path.resolve(..., "../../..") — the three-level literal ends here. + selector: "CallExpression > Literal[value=/\\.\\.\\W\\.\\.\\W\\.\\.$/]", + message: + "Do not read files 3+ levels above the test file. Tests reading root-level config must run in a dedicated non-Turbo-cached CI step.", + }, + { + // Catches readFile("../../../.coderabbit.yaml") — three-level traversal + // embedded inside a longer path string. Scoped to FS-reading calls so + // vi.mock("../../../module") (module paths, not file reads) is not flagged. + // Matches both bare readFileSync(...) and fs.readFileSync(...). + selector: + ":matches(CallExpression[callee.name=/^(readFile|readFileSync|open|openSync|stat|statSync|access|accessSync|watch|watchFile|readdir|readdirSync)/], CallExpression[callee.property.name=/^(readFile|readFileSync|open|openSync|stat|statSync|access|accessSync|watch|watchFile|readdir|readdirSync)/]) > Literal[value=/\\.\\.\\W\\.\\.\\W\\.\\./]", + message: + "Do not read files 3+ levels above the test file. Tests reading root-level config must run in a dedicated non-Turbo-cached CI step.", + }, +] + /** @type {import("eslint").Linter.Config} */ export default [ ...config, @@ -77,25 +97,14 @@ export default [ "**/services/__tests__/workflow-self-listen.test.ts", ], rules: { - "no-restricted-syntax": [ - "error", - { - // Catches path.resolve(..., "../../..") — the three-level literal ends here. - selector: "CallExpression > Literal[value=/\\.\\.\\W\\.\\.\\W\\.\\.$/]", - message: - "Do not read files 3+ levels above the test file. Tests reading root-level config must run in a dedicated non-Turbo-cached CI step.", - }, - { - // Catches readFile("../../../.coderabbit.yaml") — three-level traversal - // embedded inside a longer path string. Scoped to FS-reading calls so - // vi.mock("../../../module") (module paths, not file reads) is not flagged. - // Matches both bare readFileSync(...) and fs.readFileSync(...). - selector: - ":matches(CallExpression[callee.name=/^(readFile|readFileSync|open|openSync|stat|statSync|access|accessSync|watch|watchFile|readdir|readdirSync)/], CallExpression[callee.property.name=/^(readFile|readFileSync|open|openSync|stat|statSync|access|accessSync|watch|watchFile|readdir|readdirSync)/]) > Literal[value=/\\.\\.\\W\\.\\.\\W\\.\\./]", - message: - "Do not read files 3+ levels above the test file. Tests reading root-level config must run in a dedicated non-Turbo-cached CI step.", - }, - ], + "no-restricted-syntax": ["error", ...testPathTraversalSelectors], + }, + }, + { + // A later block replaces no-restricted-syntax for the same files, so repeat the traversal selectors. + files: ["api/**/__tests__/**/*.spec.ts"], + rules: { + "no-restricted-syntax": ["error", ...testPathTraversalSelectors, ...apiProviderTestSelectors], }, }, ]