diff --git a/package.json b/package.json index 1a44a12680..ee83c61429 100644 --- a/package.json +++ b/package.json @@ -44,7 +44,8 @@ "rimraf": "6.0.1", "tsx": "4.22.4", "turbo": "2.10.0", - "typescript": "5.9.3" + "typescript": "5.9.3", + "vitest": "4.1.9" }, "lint-staged": { "*.{js,jsx,ts,tsx,json,css,md}": [ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 183a5e02d1..a8972fd426 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -74,6 +74,9 @@ importers: typescript: specifier: 5.9.3 version: 5.9.3 + vitest: + specifier: 4.1.9 + version: 4.1.9(@opentelemetry/api@1.9.0)(@types/node@22.20.1)(@vitest/coverage-v8@4.1.9)(@vitest/ui@4.1.9)(jsdom@26.1.0)(vite@8.1.0(@types/node@22.20.1)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.22.4)(yaml@2.9.0)) apps/cli: dependencies: @@ -5037,6 +5040,7 @@ packages: eslint@9.39.4: resolution: {integrity: sha512-XoMjdBOwe/esVgEvLmNsD3IRHkm7fbKIUGvrleloJXUZgDHig2IPWNniv+GwjyJXzuNqVjlr5+4yVUZjycJwfQ==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} + deprecated: This version is no longer supported. Please see https://eslint.org/version-support for other options. hasBin: true peerDependencies: jiti: '*' diff --git a/src/api/providers/__tests__/lm-studio-timeout.spec.ts b/src/api/providers/__tests__/lm-studio-timeout.spec.ts index f661d9092e..2f84457cb6 100644 --- a/src/api/providers/__tests__/lm-studio-timeout.spec.ts +++ b/src/api/providers/__tests__/lm-studio-timeout.spec.ts @@ -11,21 +11,33 @@ vitest.mock("../utils/timeout-config", () => ({ import { getApiRequestTimeout } from "../utils/timeout-config" import { clearAllMocks } from "../../../test-utils/reset" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" -// Mock OpenAI +interface MockOpenAiClient { + chat: { + completions: { + create: ReturnType + } + } +} + +// Mock OpenAI (records each created client so tests can drive its create call) const mockOpenAIConstructor = vitest.fn() +const createdClients: MockOpenAiClient[] = [] vitest.mock("openai", () => { return { __esModule: true, default: vitest.fn().mockImplementation(function (config) { - mockOpenAIConstructor(config) - return { + const client: MockOpenAiClient = { chat: { completions: { create: vitest.fn(), }, }, } + createdClients.push(client) + mockOpenAIConstructor(config) + return client }), } }) @@ -36,7 +48,7 @@ describe("LmStudioHandler timeout configuration", () => { }) it("should use default timeout of 600 seconds when no configuration is set", () => { - ;(getApiRequestTimeout as any).mockReturnValue(600000) + vitest.mocked(getApiRequestTimeout).mockReturnValue(600000) const options: ApiHandlerOptions = { apiModelId: "llama2", @@ -57,7 +69,7 @@ describe("LmStudioHandler timeout configuration", () => { }) it("should use custom timeout when configuration is set", () => { - ;(getApiRequestTimeout as any).mockReturnValue(1200000) // 20 minutes + vitest.mocked(getApiRequestTimeout).mockReturnValue(1200000) // 20 minutes const options: ApiHandlerOptions = { apiModelId: "llama2", @@ -75,7 +87,7 @@ describe("LmStudioHandler timeout configuration", () => { }) it("should handle zero timeout (no timeout)", () => { - ;(getApiRequestTimeout as any).mockReturnValue(0) + vitest.mocked(getApiRequestTimeout).mockReturnValue(0) const options: ApiHandlerOptions = { apiModelId: "llama2", @@ -91,3 +103,278 @@ describe("LmStudioHandler timeout configuration", () => { ) }) }) + +describe("LmStudioHandler abort signal wiring", () => { + let options: ApiHandlerOptions + + // Mirror the OpenAI SDK's APIUserAbortError shape: name "Error", message + // "Request was aborted." It does not satisfy the Task.ts abort contract + // (message must end in "aborted"), so the provider must normalize it. + const sdkAbortError = (): Error => { + const err = new Error("Request was aborted.") + err.name = "Error" + return err + } + + const waitForCreateCall = async (create: { mock: { calls: unknown[][] } }, timeoutMs = 5000): Promise => { + const start = Date.now() + while (create.mock.calls.length === 0) { + if (Date.now() - start > timeoutMs) { + throw new Error("timed out waiting for the SDK create call") + } + await new Promise((resolve) => setTimeout(resolve, 5)) + } + } + + const waitForSignalAbort = (signal: AbortSignal | undefined): Promise => { + return new Promise((resolve, reject) => { + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + if (signal.aborted) { + resolve() + return + } + signal.addEventListener("abort", () => resolve(), { once: true }) + }) + } + + const lastCreate = (): MockOpenAiClient["chat"]["completions"]["create"] => { + const client = createdClients[createdClients.length - 1] + if (!client) { + throw new Error("no OpenAI client was created") + } + return client.chat.completions.create + } + + beforeEach(() => { + clearAllMocks() + vitest.mocked(getApiRequestTimeout).mockReturnValue(600000) + options = { + apiModelId: "llama2", + lmStudioModelId: "llama2", + lmStudioBaseUrl: "http://localhost:1234", + } + }) + + describe("createMessage", () => { + it("should pass a request-local AbortSignal to the SDK and bridge the external signal", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([])) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + await stream.next() + + const opts = create.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // request-local, not the external signal + expect(opts.signal.aborted).toBe(false) + + external.abort() + expect(opts.signal.aborted).toBe(true) // the external abort is bridged to the SDK signal + + await stream.next() // drain the generator + }) + + it("should fast-fail with a normalized AbortError when the signal is pre-aborted", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + const external = new AbortController() + external.abort() + + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await stream.next() + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + expect(create).not.toHaveBeenCalled() + }) + + it("should abort the in-flight SDK request when the external signal fires", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + // Simulate the OpenAI SDK: reject with its abort error when the signal aborts. + create.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + if (!opts?.signal) { + reject(new Error("SDK create was called without a signal")) + return + } + opts.signal.addEventListener("abort", () => reject(sdkAbortError()), { once: true }) + }) + }) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await waitForCreateCall(create) + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should normalize an abort error thrown mid-stream", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + const external = new AbortController() + // Simulate the OpenAI SDK stream: yield once, then reject with its + // abort error once the request-local signal is aborted. + create.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + await waitForSignalAbort(opts?.signal) + throw sdkAbortError() + })() + }) + + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const chunks: { type: string; text?: string }[] = [] + let caught: unknown + try { + for await (const chunk of stream) { + chunks.push(chunk) + if (chunk.type === "text") { + external.abort() + } + } + } catch (error) { + caught = error + } + + expect(chunks).toContainEqual({ type: "text", text: "partial" }) + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should stream reasoning chunks from a reasoning_content delta", async () => { + // Changed-line coverage regression: reasoning models served by LM Studio + // stream thinking via delta.reasoning_content, and createMessage must yield + // a reasoning chunk from that dedicated field. + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue( + asyncStreamFrom([ + { choices: [{ delta: { reasoning_content: "thinking..." }, index: 0 }] }, + { choices: [{ delta: { content: "answer" }, index: 0 }] }, + { + choices: [{ delta: {}, index: 0 }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + ]), + ) + + const chunks = await collectStream(handler.createMessage("system", [])) + + expect(chunks).toContainEqual({ type: "reasoning", text: "thinking..." }) + expect(chunks).toContainEqual({ type: "text", text: "answer" }) + }) + }) + + describe("completePrompt", () => { + it("should pass the external signal through, and nothing without a signal or with a zero timeout", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const external = new AbortController() + + expect(await handler.completePrompt("hi")).toBe("ok") + expect(create.mock.calls[0][1]).toBeUndefined() // no signal, no timeout: nothing reaches the SDK + + expect(await handler.completePrompt("hi", { abortSignal: external.signal })).toBe("ok") + // no timeout: the merged signal is the external signal itself + expect(create.mock.calls[1][1]?.signal).toBe(external.signal) + + // timeoutMs <= 0 means "no explicit timeout": nothing may reach the SDK + expect(await handler.completePrompt("hi", { timeoutMs: 0 })).toBe("ok") + expect(create.mock.calls[2][1]).toBeUndefined() + }) + + it("should merge the external signal with a positive timeoutMs", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = opts?.signal + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + signal.addEventListener("abort", () => reject(sdkAbortError()), { once: true }) + }) + }) + const external = new AbortController() + + const pending = handler.completePrompt("hi", { abortSignal: external.signal, timeoutMs: 60_000 }) + const opts = create.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // merged via AbortSignal.any + expect(opts.signal.aborted).toBe(false) + + external.abort() + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should normalize SDK abort errors instead of wrapping them", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockRejectedValue(sdkAbortError()) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should keep wrapping non-abort errors in the LM Studio debug message", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockRejectedValue(new Error("boom")) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).message).toContain("Please check the LM Studio developer logs") + }) + }) +}) diff --git a/src/api/providers/__tests__/lm-studio.spec.ts b/src/api/providers/__tests__/lm-studio.spec.ts new file mode 100644 index 0000000000..9f789d8d07 --- /dev/null +++ b/src/api/providers/__tests__/lm-studio.spec.ts @@ -0,0 +1,714 @@ +// npx vitest run api/providers/__tests__/lm-studio.spec.ts + +import { LmStudioHandler } from "../lm-studio" +import { ApiHandlerOptions } from "../../../shared/api" + +// Mock the timeout config utility +vitest.mock("../utils/timeout-config", () => ({ + getApiRequestTimeout: vitest.fn(), +})) + +import { getApiRequestTimeout } from "../utils/timeout-config" + +import { clearAllMocks } from "../../../test-utils/reset" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" + +interface MockOpenAiClient { + chat: { + completions: { + create: ReturnType + } + } +} + +// Mock OpenAI (records each created client so tests can drive its create call) +const mockOpenAIConstructor = vitest.fn() +const createdClients: MockOpenAiClient[] = [] +vitest.mock("openai", () => { + return { + __esModule: true, + default: vitest.fn().mockImplementation(function (config) { + const client: MockOpenAiClient = { + chat: { + completions: { + create: vitest.fn(), + }, + }, + } + createdClients.push(client) + mockOpenAIConstructor(config) + return client + }), + } +}) + +describe("LmStudioHandler abort wiring", () => { + let options: ApiHandlerOptions + + // Mirror the OpenAI SDK's APIUserAbortError shape: name "Error", message + // "Request was aborted." It does not satisfy the Task.ts abort contract + // (message must end in "aborted"), so the provider must normalize it. + const sdkAbortError = (): Error => { + const err = new Error("Request was aborted.") + err.name = "Error" + return err + } + + const waitForCreateCall = async (create: { mock: { calls: unknown[][] } }, timeoutMs = 5000): Promise => { + const start = Date.now() + while (create.mock.calls.length === 0) { + if (Date.now() - start > timeoutMs) { + throw new Error("timed out waiting for the SDK create call") + } + await new Promise((resolve) => setTimeout(resolve, 5)) + } + } + + const waitForSignalAbort = (signal: AbortSignal | undefined): Promise => { + return new Promise((resolve, reject) => { + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + if (signal.aborted) { + resolve() + return + } + signal.addEventListener("abort", () => resolve(), { once: true }) + }) + } + + const lastCreate = (): MockOpenAiClient["chat"]["completions"]["create"] => { + const client = createdClients[createdClients.length - 1] + if (!client) { + throw new Error("no OpenAI client was created") + } + return client.chat.completions.create + } + + beforeEach(() => { + clearAllMocks() + vitest.mocked(getApiRequestTimeout).mockReturnValue(600000) + options = { + apiModelId: "llama2", + lmStudioModelId: "llama2", + lmStudioBaseUrl: "http://localhost:1234", + } + }) + + describe("createMessage", () => { + it("passes a request-local AbortSignal to the SDK and bridges the external signal", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([])) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + await stream.next() + + const opts = create.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // request-local, not the external signal + expect(opts.signal.aborted).toBe(false) + + external.abort() + expect(opts.signal.aborted).toBe(true) // the external abort is bridged to the SDK signal + + await stream.next() // drain the generator + }) + + it("fast-fails with the abort contract error for a pre-aborted signal", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + const external = new AbortController() + external.abort() + + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await stream.next() + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("This operation was aborted") + expect(create).not.toHaveBeenCalled() + }) + + it("aborts the in-flight SDK request when the external signal fires", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + // Simulate the OpenAI SDK: reject with its abort error when the + // request-local signal aborts. A fallback resolution keeps the test + // fast if the signal never aborts (e.g. a bridging regression). + create.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + return new Promise((resolve, reject) => { + if (!opts?.signal) { + reject(new Error("SDK create was called without a signal")) + return + } + opts.signal.addEventListener("abort", () => reject(sdkAbortError()), { once: true }) + setTimeout(() => resolve(asyncStreamFrom([])), 300) + }) + }) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await waitForCreateCall(create) + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + expect(create).toHaveBeenCalledTimes(1) + }) + + it("normalizes an abort-shaped create rejection without an external signal", async () => { + // Without an external signal the outer catch cannot normalize via the + // signal, so the inner catch's own abort decision alone determines + // whether the SDK abort error is normalized to the abort contract. + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockRejectedValue(sdkAbortError()) + + let caught: unknown + try { + await collectStream(handler.createMessage("system", [])) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("fast-fails when the external signal aborts while the input token count is pending", async () => { + const handler = new LmStudioHandler(options) + let resolveCountTokens!: (tokens: number) => void + // The first (input) count stays pending so the abort can land while + // it is; any later count (the output count) resolves, so a mutant + // that slips past the fast-fail guard fails fast instead of hanging + // the mutation-test run at the pending output count. + vitest + .spyOn(handler, "countTokens") + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveCountTokens = resolve + }), + ) + .mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([])) // must never be reached + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await new Promise((resolve) => setTimeout(resolve, 10)) // let the generator reach the pending count + external.abort() + resolveCountTokens(1) + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + expect(create).not.toHaveBeenCalled() + }) + + it("settles with the abort contract when the external signal aborts and the input count never settles", async () => { + // The input count is left pending forever: without the abort race + // the generator would hang on it after a Stop. + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockImplementation(() => new Promise(() => {})) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([])) // must never be reached + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await new Promise((resolve) => setTimeout(resolve, 10)) // let the generator reach the pending count + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + expect(create).not.toHaveBeenCalled() + }) + + it("settles with the abort contract when the external signal aborts while the output count is pending", async () => { + const handler = new LmStudioHandler(options) + // The input count resolves; the output count (second call) never settles. + vitest + .spyOn(handler, "countTokens") + .mockResolvedValueOnce(1) + .mockImplementation(() => new Promise(() => {})) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const first = await stream.next() + expect(first.value?.type).toBe("text") + const pending = stream.next() // runs past the stream into the pending output count + await new Promise((resolve) => setTimeout(resolve, 10)) + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("does not report usage when a stop lands between the output count settling and the usage yield", async () => { + // The stop lands in the microtask gap: after the output count + // settles (and the race detaches its listener) but before the + // generator resumes, so only the pre-yield aborted check stops the + // usage chunk being reported for an aborted response. + const handler = new LmStudioHandler(options) + let resolveOutputCount!: (tokens: number) => void + vitest + .spyOn(handler, "countTokens") + .mockResolvedValueOnce(1) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveOutputCount = resolve + }), + ) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const first = await stream.next() + expect(first.value?.type).toBe("text") + const pending = stream.next() // runs past the stream toward the output count + await vitest.waitFor(() => { + expect(typeof resolveOutputCount).toBe("function") + }) // wait until the generator reaches the (deferred) output count + resolveOutputCount(3) + await Promise.resolve() // let the race settle and detach its listener + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("succeeds without metadata and cancels the request-local signal on completion", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const chunks = await collectStream(handler.createMessage("system", [])) + + expect(chunks).toContainEqual({ type: "text", text: "hi" }) + const opts = create.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal.aborted).toBe(true) // the finally block cancels the request-local request + }) + + it("normalizes an abort error thrown mid-stream", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + const external = new AbortController() + // Simulate the OpenAI SDK stream: yield once, then reject with its + // abort error once the request-local signal is aborted. + create.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + await waitForSignalAbort(opts?.signal) + throw sdkAbortError() + })() + }) + + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + for await (const chunk of stream) { + if (chunk.type === "text") { + external.abort() + } + } + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("wraps a non-abort create rejection in the generic debug message", async () => { + // The inner catch's handleOpenAIError throw is re-wrapped by the outer + // catch into the generic debug message (it is not abort-shaped). + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockRejectedValue(new Error("boom")) + + let caught: unknown + try { + await collectStream(handler.createMessage("system", [])) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("Error") + expect((caught as Error).message).toBe( + "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", + ) + }) + + it("wraps a non-abort stream error in the generic debug message", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + const apiError = new Error("stream exploded") + create.mockImplementation(() => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + throw apiError + })() + }) + + let caught: unknown + try { + await collectStream(handler.createMessage("system", [])) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).not.toBe("AbortError") + expect((caught as Error).message).toBe( + "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", + ) + }) + + it("registers and removes the external abort listener", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const external = new AbortController() + const addSpy = vitest.spyOn(external.signal, "addEventListener") + const removeSpy = vitest.spyOn(external.signal, "removeEventListener") + + await collectStream(handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal })) + + expect(addSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + + it("counts the input tokens for the system prompt plus every message content block", async () => { + const handler = new LmStudioHandler(options) + const countSpy = vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([])) + + const chunks = await collectStream(handler.createMessage("system", [{ role: "user", content: "hello" }])) + + // The count must see the system prompt block plus each converted + // message block - a reduced payload would undercount the context. + expect(countSpy).toHaveBeenCalledTimes(2) + expect(countSpy.mock.calls[0][0]).toEqual([ + { type: "text", text: "system" }, + { type: "text", text: "hello" }, + ]) + expect(chunks).toEqual([{ type: "usage", inputTokens: 1, outputTokens: 1 }]) + }) + + it("counts the output tokens for the exact concatenation of reasoning and visible text", async () => { + const handler = new LmStudioHandler(options) + const countSpy = vitest.spyOn(handler, "countTokens").mockResolvedValue(1) + const create = lastCreate() + // Reasoning models stream their thinking in the delta's + // `reasoning_content` field alongside the visible `content`. + create.mockResolvedValue( + asyncStreamFrom([{ choices: [{ delta: { reasoning_content: "thinking", content: "answer" } }] }]), + ) + + const chunks = await collectStream(handler.createMessage("system", [])) + + expect(chunks).toEqual([ + { type: "reasoning", text: "thinking" }, + { type: "text", text: "answer" }, + { type: "usage", inputTokens: 1, outputTokens: 1 }, + ]) + // The output count must see exactly reasoning + visible text: + // reasoning tokens are billed as output. + expect(countSpy.mock.calls[1][0]).toEqual([{ type: "text", text: "thinkinganswer" }]) + }) + + it("falls back to zero input tokens and logs when the input count fails without an abort", async () => { + const handler = new LmStudioHandler(options) + const countSpy = vitest + .spyOn(handler, "countTokens") + .mockRejectedValueOnce(new Error("count failed")) + .mockResolvedValue(1) + const errorSpy = vitest.spyOn(console, "error").mockImplementation(() => {}) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const chunks = await collectStream(handler.createMessage("system", [])) + + // A count failure is not a request failure: the response still + // streams, with the failed count falling back to zero tokens. + expect(errorSpy).toHaveBeenCalledWith("[LmStudio] Failed to count input tokens:", expect.any(Error)) + expect(countSpy).toHaveBeenCalledTimes(2) + expect(chunks).toEqual([ + { type: "text", text: "hi" }, + { type: "usage", inputTokens: 0, outputTokens: 1 }, + ]) + errorSpy.mockRestore() + }) + + it("lets an abort-shaped input count rejection propagate to the abort contract", async () => { + const handler = new LmStudioHandler(options) + // Only the input count rejects: a permanent rejection would also hit + // the output count below, whose intact abort check would surface the + // same contract error and mask a dead input-side check. + vitest.spyOn(handler, "countTokens").mockRejectedValueOnce(sdkAbortError()).mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + let caught: unknown + try { + await collectStream(handler.createMessage("system", [])) + } catch (error) { + caught = error + } + + // An abort-shaped count failure is the caller's Stop, not a count + // failure: it must surface as the abort contract error instead of + // falling back to zero tokens and keeping the stream alive. + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + expect(create).not.toHaveBeenCalled() // no request goes out after the Stop + }) + + it("falls back to zero output tokens and logs when the output count fails without an abort", async () => { + const handler = new LmStudioHandler(options) + const countSpy = vitest + .spyOn(handler, "countTokens") + .mockResolvedValueOnce(1) + .mockRejectedValueOnce(new Error("count failed")) + const errorSpy = vitest.spyOn(console, "error").mockImplementation(() => {}) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const chunks = await collectStream(handler.createMessage("system", [])) + + // Same fallback as the input count above: an error is logged, + // zero tokens are reported, and the stream still completes. + expect(errorSpy).toHaveBeenCalledWith("[LmStudio] Failed to count output tokens:", expect.any(Error)) + expect(chunks).toEqual([ + { type: "text", text: "hi" }, + { type: "usage", inputTokens: 1, outputTokens: 0 }, + ]) + errorSpy.mockRestore() + }) + + it("lets an abort-shaped output count rejection propagate to the abort contract", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "countTokens").mockResolvedValueOnce(1).mockRejectedValueOnce(sdkAbortError()) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + let caught: unknown + try { + await collectStream(handler.createMessage("system", [])) + } catch (error) { + caught = error + } + + // Same contract as the input count above: an abort is not a count + // failure, so the usage chunk must not be reported for an + // aborted response. + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("does not issue the SDK request once the signal aborts while the input count settles", async () => { + // The stop lands in the microtask gap: after the input count + // settles (and the race detaches its listener) but before the + // generator resumes, so only the post-count fast-fail guard keeps + // the request from going out once the request-local signal has + // aborted. + const handler = new LmStudioHandler(options) + let resolveInputCount!: (tokens: number) => void + vitest + .spyOn(handler, "countTokens") + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveInputCount = resolve + }), + ) + .mockResolvedValue(1) + const create = lastCreate() + create.mockResolvedValue(asyncStreamFrom([{ choices: [{ delta: { content: "hi" } }] }])) + + const external = new AbortController() + const stream = handler.createMessage("system", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await new Promise((resolve) => setTimeout(resolve, 10)) // let the generator reach the pending count + resolveInputCount(1) + await Promise.resolve() // let the race settle and detach its listener + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + expect(create).not.toHaveBeenCalled() // the request must not be issued after the Stop + }) + }) + + describe("completePrompt", () => { + it("passes the merged signal through and nothing without options", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const external = new AbortController() + + expect(await handler.completePrompt("hi")).toBe("ok") + expect(create.mock.calls[0][1]).toBeUndefined() // no signal, no timeout: nothing reaches the SDK + + expect(await handler.completePrompt("hi", { abortSignal: external.signal })).toBe("ok") + // no timeout: the merged signal is the external signal itself + expect(create.mock.calls[1][1]?.signal).toBe(external.signal) + }) + + it("normalizes a create rejection that looks like an abort", async () => { + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockRejectedValue(sdkAbortError()) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("wraps a non-abort create rejection in the generic debug message", async () => { + // The inner catch's handleOpenAIError throw is re-wrapped by the outer + // catch into the generic debug message (it is not abort-shaped). + const handler = new LmStudioHandler(options) + const create = lastCreate() + create.mockRejectedValue(new Error("model not found")) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("Error") + expect((caught as Error).message).toBe( + "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", + ) + }) + + it("normalizes an abort error thrown while building the request", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "getModel").mockImplementation(() => { + const error = new Error("aborted") + error.name = "AbortError" + throw error + }) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("wraps a non-abort request-building error in the generic debug message", async () => { + const handler = new LmStudioHandler(options) + vitest.spyOn(handler, "getModel").mockImplementation(() => { + throw new Error("model not found") + }) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).not.toBe("AbortError") + expect((caught as Error).message).toBe( + "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", + ) + }) + }) +}) diff --git a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts index c6a63902a1..7bd28cf8ef 100644 --- a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts +++ b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts @@ -81,6 +81,7 @@ describe("LmStudioHandler Native Tools", () => { }), ]), }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) // parallel_tool_calls should be true by default when not explicitly set const callArgs = mockCreate.mock.calls[0][0] @@ -103,6 +104,7 @@ describe("LmStudioHandler Native Tools", () => { expect.objectContaining({ tool_choice: "auto", }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -204,6 +206,7 @@ describe("LmStudioHandler Native Tools", () => { expect.objectContaining({ parallel_tool_calls: true, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) diff --git a/src/api/providers/__tests__/lmstudio.spec.ts b/src/api/providers/__tests__/lmstudio.spec.ts index 7ab674a0a9..6aa354b8c4 100644 --- a/src/api/providers/__tests__/lmstudio.spec.ts +++ b/src/api/providers/__tests__/lmstudio.spec.ts @@ -204,18 +204,58 @@ describe("LmStudioHandler", () => { "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", ) }) + + it("should not issue the request when the caller aborts while input token counting is pending", async () => { + const controller = new AbortController() + let releaseCount!: () => void + const countGate = new Promise((resolve) => { + releaseCount = resolve + }) + const countSpy = vi.spyOn(handler, "countTokens").mockImplementation(async () => { + await countGate + return 10 + }) + + const stream = handler.createMessage(systemPrompt, messages, { + taskId: "test-task-id", + abortSignal: controller.signal, + }) + const pending = collectStream(stream).catch((error: unknown) => error) + + // Let the generator reach the token count, then abort while it is pending. + const start = Date.now() + while (countSpy.mock.calls.length === 0) { + if (Date.now() - start > 5000) { + throw new Error("timed out waiting for the token count") + } + await new Promise((resolve) => setTimeout(resolve, 5)) + } + controller.abort() + releaseCount() + + const caught = (await pending) as Error + countSpy.mockRestore() + + expect(caught).toBeInstanceOf(Error) + expect(caught.name).toBe("AbortError") + expect(caught.message).toBe("The LM Studio request was aborted") + expect(mockCreate).not.toHaveBeenCalled() + }) }) 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.lmStudioModelId, - messages: [{ role: "user", content: "Test prompt" }], - temperature: 0, - stream: false, - }) + expect(mockCreate).toHaveBeenCalledWith( + { + model: mockOptions.lmStudioModelId, + messages: [{ role: "user", content: "Test prompt" }], + temperature: 0, + stream: false, + }, + undefined, // no abort signal or timeout: no request options reach the SDK + ) }) it("should handle API errors", async () => { diff --git a/src/api/providers/__tests__/qwen-code-native-tools.spec.ts b/src/api/providers/__tests__/qwen-code-native-tools.spec.ts index 54df551d4e..d6c01330bf 100644 --- a/src/api/providers/__tests__/qwen-code-native-tools.spec.ts +++ b/src/api/providers/__tests__/qwen-code-native-tools.spec.ts @@ -101,6 +101,7 @@ describe("QwenCodeHandler Native Tools", () => { ]), parallel_tool_calls: true, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -120,6 +121,7 @@ describe("QwenCodeHandler Native Tools", () => { expect.objectContaining({ tool_choice: "auto", }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -220,6 +222,7 @@ describe("QwenCodeHandler Native Tools", () => { expect.objectContaining({ parallel_tool_calls: true, }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), ) }) @@ -418,4 +421,460 @@ describe("QwenCodeHandler Native Tools", () => { expect(endChunks).toHaveLength(1) }) }) + + describe("abort signal wiring", () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + // Mirror the OpenAI SDK's APIUserAbortError shape: name "Error", message + // "Request was aborted." It does not satisfy the Task.ts abort contract + // (message must end in "aborted"), so the provider must normalize it. + const sdkAbortError = (): Error => { + const err = new Error("Request was aborted.") + err.name = "Error" + return err + } + + const unauthorizedError = (): Error & { status: number } => + Object.assign(new Error("unauthorized"), { status: 401 }) + + const tokenResponse = (): { ok: boolean; json: () => Promise> } => ({ + ok: true, + json: async () => ({ + access_token: "new-access-token", + refresh_token: "new-refresh-token", + token_type: "Bearer", + expires_in: 3600, + }), + }) + + const waitForCreateCall = async (create: { mock: { calls: unknown[][] } }, timeoutMs = 5000): Promise => { + const start = Date.now() + while (create.mock.calls.length === 0) { + if (Date.now() - start > timeoutMs) { + throw new Error("timed out waiting for the SDK create call") + } + await new Promise((resolve) => setTimeout(resolve, 5)) + } + } + + const waitForSignalAbort = (signal: AbortSignal | undefined): Promise => { + return new Promise((resolve, reject) => { + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + if (signal.aborted) { + resolve() + return + } + signal.addEventListener("abort", () => resolve(), { once: true }) + }) + } + + describe("createMessage", () => { + it("should pass a request-local AbortSignal to the SDK and bridge the external signal", async () => { + const external = new AbortController() + let sdkSignal: AbortSignal | undefined + // A live stream: yield one chunk, then stay open until the SDK signal aborts. + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + sdkSignal = opts?.signal + return (async function* () { + yield { choices: [{ delta: { content: "x" } }] } + await waitForSignalAbort(sdkSignal) + })() + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const first = await stream.next() + expect(first.value?.type).toBe("text") + + const opts = mockCreate.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // request-local, not the external signal + expect(opts.signal.aborted).toBe(false) + + external.abort() + expect(opts.signal.aborted).toBe(true) // the external abort is bridged to the SDK signal + + await stream.next() // resume; the live stream ends once the signal aborts + }) + + it("should fast-fail with a normalized AbortError for a pre-aborted signal", async () => { + const external = new AbortController() + external.abort() + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await stream.next() + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should abort the in-flight SDK request when the external signal fires", async () => { + const external = new AbortController() + // Simulate the OpenAI SDK: reject with its abort error when the signal aborts. + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + if (!opts?.signal) { + reject(new Error("SDK create was called without a signal")) + return + } + opts.signal.addEventListener("abort", () => reject(sdkAbortError()), { once: true }) + }) + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await waitForCreateCall(mockCreate) + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should normalize an abort error thrown mid-stream", async () => { + const external = new AbortController() + // Simulate the OpenAI SDK stream: yield once, then reject with its + // abort error once the request-local signal is aborted. + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + await waitForSignalAbort(opts?.signal) + throw sdkAbortError() + })() + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const chunks: { type: string; text?: string }[] = [] + let caught: unknown + try { + for await (const chunk of stream) { + chunks.push(chunk) + if (chunk.type === "text") { + external.abort() + } + } + } catch (error) { + caught = error + } + + expect(chunks).toContainEqual({ type: "text", text: "partial" }) + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should rethrow non-abort stream errors unchanged", async () => { + const boom = new Error("boom") + mockCreate.mockImplementationOnce(() => { + return (async function* () { + yield { choices: [{ delta: { content: "x" } }] } + throw boom + })() + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1" }) + let caught: unknown + try { + await collectStream(stream) + } catch (error) { + caught = error + } + + expect(caught).toBe(boom) + }) + + it("should split tag boundaries across chunks into reasoning and text", async () => { + // Exercises the incremental think-tag parser: one chunk opens a + // thinking block (odd segment), the next one closes it (even + // segment) and continues as visible text. + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [{ delta: { content: "ab" } }] }, + { choices: [{ delta: { content: "c" } }] }, + ]), + ) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1" }) + const chunks = await collectStream(stream) + + expect(chunks).toContainEqual({ type: "reasoning", text: "b" }) + expect(chunks).toContainEqual({ type: "text", text: "c" }) + expect(chunks).not.toContainEqual(expect.objectContaining({ type: "text", text: "b" })) + }) + + it("should tolerate degenerate stream shapes (empty choice, repeated content, zero usage)", async () => { + // Changed-line coverage: exercises the defensive branches of the stream + // loop — a chunk with no choice, a delta that repeats the previous full + // content (empty after trimming), a think block that starts the text so the + // split yields an empty leading segment, and a usage payload of zeros. + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [] }, + { choices: [{ delta: { content: "hi" }, index: 0 }] }, + { choices: [{ delta: { content: "hi" }, index: 0 }] }, + { choices: [{ delta: { content: "thoughtout" }, index: 0 }] }, + { + choices: [{ delta: {}, index: 0, finish_reason: "stop" }], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }, + ]), + ) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1" }) + const chunks = await collectStream(stream) + + const hiChunks = chunks.filter((chunk) => chunk.type === "text" && chunk.text === "hi") + expect(hiChunks).toHaveLength(1) // the repeated content chunk yields nothing + expect(chunks).toContainEqual({ type: "reasoning", text: "thought" }) + expect(chunks).toContainEqual({ type: "text", text: "out" }) + expect(chunks).toContainEqual({ type: "usage", inputTokens: 0, outputTokens: 0 }) + }) + + it("should not retry after 401 when the abort signal fires during the refresh", async () => { + const external = new AbortController() + const fetchMock = vi.fn().mockImplementation(async () => { + external.abort() // simulate Stop pressed while the token refresh is in flight + return tokenResponse() + }) + vi.stubGlobal("fetch", fetchMock) + mockCreate.mockRejectedValueOnce(unauthorizedError()) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await collectStream(stream) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + expect(mockCreate).toHaveBeenCalledTimes(1) // the retried request was never sent + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it("should normalize an abort error from the 401 retry instead of exposing the raw SDK error", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const external = new AbortController() + mockCreate + .mockRejectedValueOnce(unauthorizedError()) + .mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + external.abort() // Stop pressed while the retried request is in flight + return Promise.reject(sdkAbortError()) + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await collectStream(stream) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect((caught as Error).message).not.toBe("Request was aborted.") // not the raw SDK error + expect(mockCreate).toHaveBeenCalledTimes(2) // first attempt 401, then the aborted retry + }) + + it("should rethrow a non-abort error from the 401 retry unchanged", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const apiError = new Error("boom") + mockCreate.mockRejectedValueOnce(unauthorizedError()).mockRejectedValueOnce(apiError) + + const stream = handler.createMessage("test prompt", [], { + taskId: "t1", + abortSignal: new AbortController().signal, + }) + let caught: unknown + try { + await collectStream(stream) + } catch (error) { + caught = error + } + + expect(caught).toBe(apiError) + expect(mockCreate).toHaveBeenCalledTimes(2) + }) + }) + + describe("completePrompt", () => { + it("should pass the external signal through, and nothing without a signal or with a zero timeout", async () => { + mockCreate + .mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + .mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + .mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + const external = new AbortController() + + expect(await handler.completePrompt("hi")).toBe("ok") + expect(mockCreate.mock.calls[0][1]).toBeUndefined() // no signal, no timeout: nothing reaches the SDK + + expect(await handler.completePrompt("hi", { abortSignal: external.signal })).toBe("ok") + // no timeout: the merged signal is the external signal itself + expect(mockCreate.mock.calls[1][1]?.signal).toBe(external.signal) + + // timeoutMs <= 0 means "no explicit timeout": nothing may reach the SDK + expect(await handler.completePrompt("hi", { timeoutMs: 0 })).toBe("ok") + expect(mockCreate.mock.calls[2][1]).toBeUndefined() + }) + + it("should merge the external signal with a positive timeoutMs", async () => { + const external = new AbortController() + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + return new Promise((_resolve, reject) => { + const signal = opts?.signal + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + signal.addEventListener("abort", () => reject(sdkAbortError()), { once: true }) + }) + }) + + const pending = handler.completePrompt("hi", { abortSignal: external.signal, timeoutMs: 60_000 }) + await waitForCreateCall(mockCreate) + const opts = mockCreate.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // merged via AbortSignal.any + expect(opts.signal.aborted).toBe(false) + + external.abort() + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should fast-fail with a normalized AbortError for a pre-aborted signal", async () => { + const external = new AbortController() + external.abort() + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + expect(fs.readFile).not.toHaveBeenCalled() // no work starts after a pre-aborted signal + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should retry after 401 and pass the same abort signal to the retry", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + mockCreate + .mockRejectedValueOnce(unauthorizedError()) + .mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] }) + + const result = await handler.completePrompt("hi", { abortSignal: new AbortController().signal }) + + expect(result).toBe("retried") + expect(mockCreate).toHaveBeenCalledTimes(2) + expect(mockCreate.mock.calls[1][1]?.signal).toBe(mockCreate.mock.calls[0][1]?.signal) + }) + + it("should not retry after 401 when the abort signal fires during the refresh", async () => { + const external = new AbortController() + const fetchMock = vi.fn().mockImplementation(async () => { + external.abort() // simulate Stop pressed while the token refresh is in flight + return tokenResponse() + }) + vi.stubGlobal("fetch", fetchMock) + mockCreate.mockRejectedValueOnce(unauthorizedError()) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + expect(mockCreate).toHaveBeenCalledTimes(1) // the retried request was never sent + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it("should normalize SDK abort errors instead of rethrowing them", async () => { + mockCreate.mockRejectedValueOnce(sdkAbortError()) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toMatch(/aborted$/) + }) + + it("should normalize an abort error from the 401 retry instead of exposing the raw SDK error", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const external = new AbortController() + mockCreate + .mockRejectedValueOnce(unauthorizedError()) + .mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + external.abort() // Stop pressed while the retried request is in flight + return Promise.reject(sdkAbortError()) + }) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect((caught as Error).message).not.toBe("Request was aborted.") // not the raw SDK error + expect(mockCreate).toHaveBeenCalledTimes(2) // first attempt 401, then the aborted retry + }) + + it("should rethrow a non-abort error from the 401 retry unchanged", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const apiError = new Error("boom") + mockCreate.mockRejectedValueOnce(unauthorizedError()).mockRejectedValueOnce(apiError) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: new AbortController().signal }) + } catch (error) { + caught = error + } + + expect(caught).toBe(apiError) + expect(mockCreate).toHaveBeenCalledTimes(2) + }) + }) + }) }) diff --git a/src/api/providers/__tests__/qwen-code.spec.ts b/src/api/providers/__tests__/qwen-code.spec.ts new file mode 100644 index 0000000000..c1834b32df --- /dev/null +++ b/src/api/providers/__tests__/qwen-code.spec.ts @@ -0,0 +1,668 @@ +// npx vitest run api/providers/__tests__/qwen-code.spec.ts + +// Mock filesystem - must come before other imports +vi.mock("node:fs", () => ({ + promises: { + readFile: vi.fn(), + writeFile: vi.fn(), + }, +})) + +const mockCreate = vi.fn() +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { clearAllMocks } from "../../../test-utils/reset" +vi.mock("openai", () => { + return { + __esModule: true, + default: vi.fn().mockImplementation(function () { + return { + apiKey: "test-key", + baseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1", + chat: { + completions: { + create: mockCreate, + }, + }, + } + }), + } +}) + +import { promises as fs } from "node:fs" +import { QwenCodeHandler } from "../qwen-code" +import { NativeToolCallParser } from "../../../core/assistant-message/NativeToolCallParser" +import type { ApiHandlerOptions } from "../../../shared/api" + +describe("QwenCodeHandler abort wiring", () => { + let handler: QwenCodeHandler + let mockOptions: ApiHandlerOptions + + // Mirror the OpenAI SDK's APIUserAbortError shape: name "Error", message + // "Request was aborted." It does not satisfy the Task.ts abort contract + // (message must end in "aborted"), so the provider must normalize it. + const sdkAbortError = (): Error => { + const err = new Error("Request was aborted.") + err.name = "Error" + return err + } + + const unauthorizedError = (): Error & { status: number } => + Object.assign(new Error("unauthorized"), { status: 401 }) + + const tokenResponse = (): { ok: boolean; json: () => Promise> } => ({ + ok: true, + json: async () => ({ + access_token: "new-access-token", + refresh_token: "new-refresh-token", + token_type: "Bearer", + expires_in: 3600, + }), + }) + + const waitForSignalAbort = (signal: AbortSignal | undefined): Promise => { + return new Promise((resolve, reject) => { + if (!signal) { + reject(new Error("SDK create was called without a signal")) + return + } + if (signal.aborted) { + resolve() + return + } + signal.addEventListener("abort", () => resolve(), { once: true }) + }) + } + + afterEach(() => { + vi.unstubAllGlobals() + }) + + beforeEach(() => { + clearAllMocks() + + // Mock credentials file + const mockCredentials = { + access_token: "test-access-token", + refresh_token: "test-refresh-token", + token_type: "Bearer", + expiry_date: Date.now() + 3600000, // 1 hour from now + resource_url: "https://dashscope.aliyuncs.com/compatible-mode/v1", + } + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(mockCredentials)) + vi.mocked(fs.writeFile).mockResolvedValue(undefined) + + mockOptions = { + apiModelId: "qwen3-coder-plus", + } + handler = new QwenCodeHandler(mockOptions) + + // Clear NativeToolCallParser state before each test + NativeToolCallParser.clearRawChunkState() + }) + + describe("callApiWithRetry", () => { + it("normalizes a first-attempt SDK abort error", async () => { + mockCreate.mockRejectedValueOnce(sdkAbortError()) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + }) + + it("rethrows a non-abort, non-401 error unchanged", async () => { + const apiError = new Error("server exploded") + Object.assign(apiError, { status: 500 }) + mockCreate.mockRejectedValueOnce(apiError) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBe(apiError) + }) + + it("rethrows a non-object rejection unchanged", async () => { + // A null rejection is not an object: the optional status read must + // yield undefined (not throw) and fall through to the rethrow. + mockCreate.mockRejectedValueOnce(null) + + let caught: unknown + try { + await handler.completePrompt("hi") + } catch (error) { + caught = error + } + + expect(caught).toBeNull() + }) + + it("retries once after 401 and succeeds", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + mockCreate + .mockRejectedValueOnce(unauthorizedError()) + .mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] }) + + const result = await handler.completePrompt("hi", { abortSignal: new AbortController().signal }) + + expect(result).toBe("retried") + expect(mockCreate).toHaveBeenCalledTimes(2) + }) + + it("retries after 401 when no external signal is provided", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + mockCreate + .mockRejectedValueOnce(unauthorizedError()) + .mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] }) + + const result = await handler.completePrompt("hi") + + expect(result).toBe("retried") + expect(mockCreate).toHaveBeenCalledTimes(2) + }) + + it("does not retry after 401 once the signal aborts during the refresh", async () => { + const external = new AbortController() + const fetchMock = vi.fn().mockImplementation(async () => { + external.abort() // simulate Stop pressed while the token refresh is in flight + return tokenResponse() + }) + vi.stubGlobal("fetch", fetchMock) + mockCreate.mockRejectedValueOnce(unauthorizedError()) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).toHaveBeenCalledTimes(1) // the retried request was never sent + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it("normalizes an abort error from the 401 retry", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const external = new AbortController() + mockCreate.mockRejectedValueOnce(unauthorizedError()).mockImplementationOnce(() => { + external.abort() // Stop pressed while the retried request is in flight + return Promise.reject(sdkAbortError()) + }) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).toHaveBeenCalledTimes(2) // first attempt 401, then the aborted retry + }) + + it("rethrows a non-abort error from the 401 retry unchanged", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(tokenResponse())) + const apiError = new Error("boom") + mockCreate.mockRejectedValueOnce(unauthorizedError()).mockRejectedValueOnce(apiError) + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: new AbortController().signal }) + } catch (error) { + caught = error + } + + expect(caught).toBe(apiError) + expect(mockCreate).toHaveBeenCalledTimes(2) + }) + }) + + describe("createMessage", () => { + it("passes a request-local signal and bridges the external abort", async () => { + const external = new AbortController() + const addSpy = vi.spyOn(external.signal, "addEventListener") + mockCreate.mockResolvedValueOnce(asyncStreamFrom([{ choices: [{ delta: { content: "x" } }] }])) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const first = await stream.next() + expect(first.value?.type).toBe("text") + + const opts = mockCreate.mock.calls[0][1] + expect(opts?.signal).toBeInstanceOf(AbortSignal) + expect(opts.signal).not.toBe(external.signal) // request-local, not the external signal + expect(addSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + + external.abort() + expect(opts.signal.aborted).toBe(true) // the external abort is bridged to the SDK signal + + await stream.next() // drain the generator + }) + + it("fast-fails with the abort contract error for a pre-aborted signal", async () => { + const external = new AbortController() + external.abort() + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + await stream.next() + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("This operation was aborted") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("settles with the abort contract when the signal aborts while the credential load is pending", async () => { + // The cached credential load never settles, so without the race the + // generator would wait for fs.readFile forever after a Stop. + vi.mocked(fs.readFile).mockImplementation(() => new Promise(() => {})) + const external = new AbortController() + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await new Promise((resolve) => setTimeout(resolve, 10)) // let the generator reach the pending load + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("settles with the abort contract when the signal aborts while the token refresh is pending, and the shared refresh still completes", async () => { + // Expired cached credentials force a refresh; its fetch stays + // pending until the Stop lands, then settles in the background — + // the request-local abort must cut only the wait, not the shared + // refresh that other requests dedupe against. + const expiredCredentials = { + access_token: "expired-access-token", + refresh_token: "test-refresh-token", + token_type: "Bearer", + expiry_date: Date.now() - 1000, // expired + resource_url: "https://dashscope.aliyuncs.com/compatible-mode/v1", + } + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expiredCredentials)) + let resolveFetch!: (response: { ok: boolean; json: () => Promise> }) => void + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation( + () => + new Promise<{ ok: boolean; json: () => Promise> }>((resolve) => { + resolveFetch = resolve + }), + ), + ) + const external = new AbortController() + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + const pending = stream.next() + await new Promise((resolve) => setTimeout(resolve, 10)) // let the generator reach the pending refresh + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).not.toHaveBeenCalled() + + // The shared refresh keeps running: once its fetch settles, the new + // credentials are persisted even though this request gave up. + resolveFetch(tokenResponse()) + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(fs.writeFile).toHaveBeenCalled() + }) + + it("streams content, strips repeated prefixes and splits think tags", async () => { + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([ + { choices: [{ delta: { content: "Hello" } }] }, + // A full duplicate chunk strips to empty content and must + // emit no chunk (it exercises the empty newText guard). + { choices: [{ delta: { content: "Hello" } }] }, + { choices: [{ delta: { content: "Hello world" } }] }, + { choices: [{ delta: { content: "bye" } }] }, + { choices: [{ delta: { content: "abc" } }] }, + { choices: [{ delta: { content: "onlytag" } }] }, + ]), + ) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toEqual([ + { type: "text", text: "Hello" }, + { type: "text", text: " world" }, + { type: "text", text: "bye" }, + { type: "text", text: "a" }, + { type: "reasoning", text: "b" }, + { type: "text", text: "c" }, + { type: "text", text: "only" }, + { type: "reasoning", text: "tag" }, + ]) + }) + + it("yields no chunks for empty thinking tags", async () => { + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([{ choices: [{ delta: { content: "" } }] }]), + ) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toEqual([]) + }) + + it("emits reasoning from reasoning_content and tolerates malformed deltas", async () => { + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([ + { choices: [{ delta: { reasoning_content: "thinking" } }] }, + { choices: [{}] }, // no delta at all + { choices: [] }, // empty choices + { choices: [{ delta: { content: "" } }] }, // empty content + ]), + ) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toEqual([{ type: "reasoning", text: "thinking" }]) + }) + + it("emits content equal to the Stryker sentinel as-is", async () => { + // fullContent must start empty: a first chunk that happens to begin + // with the sentinel string must not be treated as a repeated prefix. + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([{ choices: [{ delta: { content: "Stryker was here!done" } }] }]), + ) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toEqual([{ type: "text", text: "Stryker was here!done" }]) + }) + + it("builds the streaming request with strict defaults when no metadata is provided", async () => { + mockCreate.mockResolvedValueOnce(asyncStreamFrom([{ choices: [{ delta: { content: "hello" } }] }])) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toContainEqual({ type: "text", text: "hello" }) + expect(mockCreate.mock.calls[0][0]).toEqual( + expect.objectContaining({ + model: "qwen3-coder-plus", + temperature: 0, + stream: true, + stream_options: { include_usage: true }, + parallel_tool_calls: true, + messages: [{ role: "system", content: "test prompt" }], + }), + ) + }) + + it("emits tool_call_partial chunks and tool_call_end on finish_reason", async () => { + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([ + { + choices: [ + { + delta: { + tool_calls: [ + { index: 0, id: "call_1", function: { name: "tool_x", arguments: '{"a":1}' } }, + ], + }, + }, + ], + }, + { choices: [{ delta: { tool_calls: [{ index: 1, id: "call_2" }] } }] }, + { choices: [{ delta: {}, finish_reason: "tool_calls" }] }, + ]), + ) + + const stream = handler.createMessage("test prompt", []) + + // Collect the provider stream and process tool_call_partial chunks + // through NativeToolCallParser, exactly as Task.ts does. + const chunks = [] + for await (const chunk of stream) { + if (chunk.type === "tool_call_partial") { + NativeToolCallParser.processRawChunk({ + index: chunk.index, + id: chunk.id, + name: chunk.name, + arguments: chunk.arguments, + }) + } + chunks.push(chunk) + } + + expect(chunks).toEqual([ + { type: "tool_call_partial", index: 0, id: "call_1", name: "tool_x", arguments: '{"a":1}' }, + { type: "tool_call_partial", index: 1, id: "call_2", name: undefined, arguments: undefined }, + { type: "tool_call_end", id: "call_1" }, + { type: "tool_call_end", id: "call_2" }, + ]) + }) + + it("emits a usage chunk from the stream usage field", async () => { + mockCreate.mockResolvedValueOnce( + asyncStreamFrom([ + { choices: [{ delta: {} }], usage: { prompt_tokens: 42, completion_tokens: 7, total_tokens: 49 } }, + ]), + ) + + const chunks = await collectStream(handler.createMessage("test prompt", [])) + + expect(chunks).toEqual([{ type: "usage", inputTokens: 42, outputTokens: 7 }]) + }) + + it("normalizes an abort error thrown mid-stream", async () => { + const external = new AbortController() + // Simulate the OpenAI SDK stream: yield once, then reject with its + // abort error once the request-local signal is aborted. + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + await waitForSignalAbort(opts?.signal) + throw sdkAbortError() + })() + }) + + const stream = handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }) + let caught: unknown + try { + for await (const chunk of stream) { + if (chunk.type === "text") { + external.abort() + } + } + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + }) + + it("rethrows a non-abort stream error unchanged", async () => { + const apiError = new Error("stream exploded") + mockCreate.mockImplementationOnce(() => { + return (async function* () { + yield { choices: [{ delta: { content: "partial" } }] } + throw apiError + })() + }) + + let caught: unknown + try { + await collectStream(handler.createMessage("test prompt", [])) + } catch (error) { + caught = error + } + + expect(caught).toBe(apiError) + }) + + it("aborts the request-local signal when the consumer stops early", async () => { + mockCreate.mockImplementationOnce((_params: unknown, opts?: { signal?: AbortSignal }) => { + return (async function* () { + yield { choices: [{ delta: { content: "x" } }] } + await waitForSignalAbort(opts?.signal) // stay open until the SDK signal aborts + })() + }) + + const stream = handler.createMessage("test prompt", []) + const first = await stream.next() + expect(first.value?.type).toBe("text") + + const opts = mockCreate.mock.calls[0][1] + await stream.return(undefined) // consumer stops early; finally must cancel the request + + expect(opts?.signal.aborted).toBe(true) + }) + + it("removes the external abort listener when the stream completes", async () => { + const external = new AbortController() + const removeSpy = vi.spyOn(external.signal, "removeEventListener") + mockCreate.mockResolvedValueOnce(asyncStreamFrom([{ choices: [{ delta: { content: "ok" } }] }])) + + await collectStream( + handler.createMessage("test prompt", [], { taskId: "t1", abortSignal: external.signal }), + ) + + expect(removeSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + }) + + describe("completePrompt", () => { + it("passes the merged signal through and nothing without options", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "qwen ok" } }] }) + + expect(await handler.completePrompt("hi")).toBe("qwen ok") + expect(mockCreate.mock.calls[0][1]).toBeUndefined() // no signal, no timeout: nothing reaches the SDK + + const external = new AbortController() + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "qwen ok" } }] }) + + expect(await handler.completePrompt("hi", { abortSignal: external.signal })).toBe("qwen ok") + // no timeout: the merged signal is the external signal itself + expect(mockCreate.mock.calls[1][1]?.signal).toBe(external.signal) + }) + + it("fast-fails with the abort contract error for a pre-aborted signal", async () => { + const external = new AbortController() + external.abort() + + let caught: unknown + try { + await handler.completePrompt("hi", { abortSignal: external.signal }) + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("This operation was aborted") + expect(mockCreate).not.toHaveBeenCalled() + expect(fs.readFile).not.toHaveBeenCalled() // no work starts after a pre-aborted signal + }) + + it("settles with the abort contract when the signal aborts while the credential load is pending", async () => { + // The cached credential load never settles, so without the race the + // call would wait for fs.readFile forever after a Stop. + vi.mocked(fs.readFile).mockImplementation(() => new Promise(() => {})) + const external = new AbortController() + + const pending = handler.completePrompt("hi", { abortSignal: external.signal }) + await new Promise((resolve) => setTimeout(resolve, 10)) // let the call reach the pending load + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("settles with the abort contract when the signal aborts while the token refresh is pending, and the shared refresh still completes", async () => { + // Expired cached credentials force a refresh; its fetch stays + // pending until the Stop lands, then settles in the background — + // the abort must cut only the wait, not the shared refresh that + // other requests dedupe against. + const expiredCredentials = { + access_token: "expired-access-token", + refresh_token: "test-refresh-token", + token_type: "Bearer", + expiry_date: Date.now() - 1000, // expired + resource_url: "https://dashscope.aliyuncs.com/compatible-mode/v1", + } + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expiredCredentials)) + let resolveFetch!: (response: { ok: boolean; json: () => Promise> }) => void + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation( + () => + new Promise<{ ok: boolean; json: () => Promise> }>((resolve) => { + resolveFetch = resolve + }), + ), + ) + const external = new AbortController() + + const pending = handler.completePrompt("hi", { abortSignal: external.signal }) + await new Promise((resolve) => setTimeout(resolve, 10)) // let the call reach the pending refresh + external.abort() + + let caught: unknown + try { + await pending + } catch (error) { + caught = error + } + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + expect(mockCreate).not.toHaveBeenCalled() + + // The shared refresh keeps running: once its fetch settles, the new + // credentials are persisted even though this request gave up. + resolveFetch(tokenResponse()) + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(fs.writeFile).toHaveBeenCalled() + }) + }) +}) diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index 59f484829c..8a64500ba2 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -18,8 +18,17 @@ import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" +import { RequestConfigBuilder } from "./config-builder/request-config-builder" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { getModelsFromCache } from "./fetchers/modelCache" +import { + mergeAbortSignalAndTimeout, + throwIfAborted, + createAbortError, + isRequestAborted, + settleOnAbort, + type OpenAiRequestOptions, +} from "./utils/abort-signal" import { handleOpenAIError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -47,6 +56,9 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { + // Fast-fail if the caller’s stop signal already fired before we started. + throwIfAborted(metadata?.abortSignal) + const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [ { role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages), @@ -77,18 +89,51 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan return result } - let inputTokens = 0 - try { - inputTokens = await this.countTokens([{ type: "text", text: systemPrompt }, ...toContentBlocks(messages)]) - } catch (err) { - console.error("[LmStudio] Failed to count input tokens:", err) - inputTokens = 0 + // Request-local abort controller — a class field would outlive this + // request and let concurrent requests abort each other. It is created + // before the token counts so an abort landing while either count is + // pending settles this generator promptly instead of waiting for the + // count to finish. + const requestController = new AbortController() + const onExternalAbort = () => { + requestController.abort() + } + const externalSignal = metadata?.abortSignal + if (externalSignal) { + externalSignal.addEventListener("abort", onExternalAbort) + // A listener registered after the signal already aborted never + // fires, so bridge the aborted state into the request-local + // controller (this also covers an abort landing while a token + // count below is still pending). + // Stryker disable next-line ConditionalExpression: externalSignal.aborted can never be true here - a pre-aborted signal is already rejected by the entry throwIfAborted fast-fail above, and no await sits between that check and this bridge, so the guard is unreachable + if (externalSignal.aborted) { + // Stryker disable next-line CallExpression: unreachable branch body - a pre-aborted external signal is rejected by the entry throwIfAborted fast-fail before this bridge registers + requestController.abort() + } } - - let assistantText = "" - let reasoningOutput = "" try { + let inputTokens = 0 + try { + inputTokens = await settleOnAbort( + this.countTokens([{ type: "text", text: systemPrompt }, ...toContentBlocks(messages)]), + requestController.signal, + this.providerName, + ) + } catch (err) { + if (isRequestAborted(err, requestController.signal)) { + // An abort is not a count failure: let it propagate to the + // outer catch for normalization instead of falling back to + // zero tokens. + throw err + } + console.error("[LmStudio] Failed to count input tokens:", err) + inputTokens = 0 + } + + let assistantText = "" + let reasoningOutput = "" + const params: OpenAI.Chat.ChatCompletionCreateParamsStreaming & { draft_model?: string } = { model: this.getModel().id, messages: openAiMessages, @@ -103,10 +148,25 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan params.draft_model = this.options.lmStudioDraftModelId } + // Bridge the request-local signal into the SDK request options so the + // in-flight request can be cancelled. + const createOptions = new RequestConfigBuilder() + .setOption("signal", requestController.signal) + .build() + + // Fast-fail if the caller aborted while countTokens() above was + // pending — the request must not be issued once the request-local + // signal has aborted. + throwIfAborted(requestController.signal) + let results try { - results = await this.client.chat.completions.create(params) + results = await this.client.chat.completions.create(params, createOptions) } catch (error) { + if (isRequestAborted(error, externalSignal)) { + // Stryker disable next-line StringLiteral: inner abort throw is re-normalized by the outer catch's createAbortError (name "AbortError" always matches isRequestAborted), so this literal is unobservable + throw createAbortError("LM Studio") + } throw handleOpenAIError(error, this.providerName) } @@ -169,21 +229,49 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan try { // Reasoning tokens are billed as output, so count them alongside the // visible text — otherwise thinking models under-report usage entirely. - outputTokens = await this.countTokens([{ type: "text", text: reasoningOutput + assistantText }]) + outputTokens = await settleOnAbort( + this.countTokens([{ type: "text", text: reasoningOutput + assistantText }]), + requestController.signal, + this.providerName, + ) } catch (err) { + if (isRequestAborted(err, requestController.signal)) { + // Same as the input count above: an abort is not a count + // failure — propagate it instead of reporting zero output + // tokens. + throw err + } console.error("[LmStudio] Failed to count output tokens:", err) outputTokens = 0 } + // A stop can land in the microtask gap between the output count + // settling and this generator resuming; do not report usage for an + // aborted response. + if (requestController.signal.aborted) { + throw createAbortError(this.providerName) + } + yield { type: "usage", inputTokens, outputTokens, } as const } catch (error) { + if (isRequestAborted(error, externalSignal)) { + throw createAbortError("LM Studio") + } throw new Error( "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", ) + } finally { + // Cancel the in-flight SDK request if the consumer stopped iterating + // early (break or return): the external signal may never fire in that + // case, and only the request-local signal reaches the SDK. + requestController.abort() + if (externalSignal) { + externalSignal.removeEventListener("abort", onExternalAbort) + } } } @@ -206,6 +294,14 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Fast-fail if the caller’s stop signal already fired before we started. + throwIfAborted(options?.abortSignal) + + // Merge the external stop signal with an optional per-call timeout. A + // timeoutMs <= 0 means "no explicit timeout" inside the util, so zero + // never reaches the SDK as an explicit timeout. + const requestSignal = mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs) + try { // Create params object with optional draft model const params: any = { @@ -220,14 +316,28 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan params.draft_model = this.options.lmStudioDraftModelId } + // CompletePromptOptions is not createMessage metadata (no taskId), so + // the generic builder takes the merged signal via setOption instead of + // setAbortSignal(metadata). + const createOptions = new RequestConfigBuilder() + .setOption("signal", requestSignal) + .build() + let response try { - response = await this.client.chat.completions.create(params) + response = await this.client.chat.completions.create(params, createOptions) } catch (error) { + if (isRequestAborted(error, requestSignal)) { + // Stryker disable next-line StringLiteral: inner abort throw is re-normalized by the outer catch's createAbortError (name "AbortError" always matches isRequestAborted), so this literal is unobservable + throw createAbortError("LM Studio") + } throw handleOpenAIError(error, this.providerName) } return response.choices[0]?.message.content || "" } catch (error) { + if (isRequestAborted(error, requestSignal)) { + throw createAbortError("LM Studio") + } throw new Error( "Please check the LM Studio developer logs to debug what went wrong. You may need to load the model with a larger context length to work with Zoo Code's prompts.", ) diff --git a/src/api/providers/qwen-code.ts b/src/api/providers/qwen-code.ts index 7d98bcb77d..c8174b36dd 100644 --- a/src/api/providers/qwen-code.ts +++ b/src/api/providers/qwen-code.ts @@ -14,7 +14,16 @@ import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" +import { RequestConfigBuilder } from "./config-builder/request-config-builder" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { + mergeAbortSignalAndTimeout, + throwIfAborted, + createAbortError, + isRequestAborted, + settleOnAbort, + type OpenAiRequestOptions, +} from "./utils/abort-signal" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" const QWEN_OAUTH_BASE_URL = "https://chat.qwen.ai" @@ -194,17 +203,46 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan return baseUrl.endsWith("/v1") ? baseUrl : `${baseUrl}/v1` } - private async callApiWithRetry(apiCall: () => Promise): Promise { + private async callApiWithRetry(apiCall: () => Promise, externalSignal?: AbortSignal): Promise { try { return await apiCall() - } catch (error: any) { - if (error.status === 401) { - // Token expired, refresh and retry + } catch (error) { + // An aborted request must never be retried: normalize it to the + // Task.ts abort contract (name "AbortError", message ending in + // "aborted") instead of rethrowing the raw SDK abort error. + if (isRequestAborted(error, externalSignal)) { + throw createAbortError("Qwen Code") + } + // The catch binding is unknown, so read the OpenAI SDK's APIError + // `status` through a narrow shape; a non-object throw yields + // undefined and falls through to the rethrow below. + const status = (error as { status?: number } | null)?.status + if (status === 401) { + // Token expired, refresh and retry. The retry reuses apiCall’s + // captured request options, so it carries the same abort signal. + // (An already-aborted request is normalized above and never + // reaches this branch.) this.credentials = await this.refreshAccessToken(this.credentials!) + // A stop can land while the refresh await is in flight — re-check + // before the retried request goes out so it is not sent. + if (externalSignal?.aborted) { + throw createAbortError("Qwen Code") + } const client = this.ensureClient() client.apiKey = this.credentials.access_token client.baseURL = this.getBaseUrl(this.credentials) - return await apiCall() + // A stop can also land while the retried request itself is in + // flight; that rejection must go through the same abort + // normalization as the first attempt instead of escaping as the + // raw SDK abort error. + try { + return await apiCall() + } catch (retryError) { + if (isRequestAborted(retryError, externalSignal)) { + throw createAbortError("Qwen Code") + } + throw retryError + } } else { throw error } @@ -216,107 +254,154 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - await this.ensureAuthenticated() - const client = this.ensureClient() - const model = this.getModel() - - const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = { - role: "system", - content: systemPrompt, + // Fast-fail if the caller’s stop signal already fired before we started. + throwIfAborted(metadata?.abortSignal) + + // Request-local abort controller — a class field would outlive this + // request and let concurrent requests abort each other. + const requestController = new AbortController() + const onExternalAbort = () => { + requestController.abort() + } + const externalSignal = metadata?.abortSignal + if (externalSignal) { + externalSignal.addEventListener("abort", onExternalAbort) } - const convertedMessages = [systemMessage, ...convertToOpenAiMessages(messages)] + try { + // A stop can land while the credential load or the token refresh + // below is still in flight; race the auth wait against the + // request-local signal so the caller settles promptly instead of + // waiting for credential I/O. The shared refresh keeps running — + // only this wait is cut. + // Stryker disable next-line StringLiteral: the rejection from this wait is always re-normalized by the outer catch's createAbortError("Qwen Code") (name "AbortError" always matches isRequestAborted), so this provider-name literal is unobservable + await settleOnAbort(this.ensureAuthenticated(), requestController.signal, "Qwen Code") + const client = this.ensureClient() + const model = this.getModel() + + const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = { + role: "system", + content: systemPrompt, + } - const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { - model: model.id, - temperature: 0, - messages: convertedMessages, - stream: true, - stream_options: { include_usage: true }, - max_completion_tokens: model.info.maxTokens, - tools: this.convertToolsForOpenAI(metadata?.tools), - tool_choice: metadata?.tool_choice, - parallel_tool_calls: metadata?.parallelToolCalls ?? true, - } + const convertedMessages = [systemMessage, ...convertToOpenAiMessages(messages)] + + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: model.id, + temperature: 0, + messages: convertedMessages, + stream: true, + stream_options: { include_usage: true }, + max_completion_tokens: model.info.maxTokens, + tools: this.convertToolsForOpenAI(metadata?.tools), + tool_choice: metadata?.tool_choice, + parallel_tool_calls: metadata?.parallelToolCalls ?? true, + } - const stream = await this.callApiWithRetry(() => client.chat.completions.create(requestOptions)) + // Bridge the request-local signal into the SDK request options so + // the in-flight request (and any 401 retry) can be cancelled. + const createOptions = new RequestConfigBuilder() + .setOption("signal", requestController.signal) + .build() - let fullContent = "" + const stream = await this.callApiWithRetry( + () => client.chat.completions.create(requestOptions, createOptions), + externalSignal, + ) - for await (const apiChunk of stream) { - const delta = apiChunk.choices[0]?.delta ?? {} - const finishReason = apiChunk.choices[0]?.finish_reason + let fullContent = "" - const reasoningText = extractReasoningFromDelta(delta) - if (reasoningText) { - yield { type: "reasoning", text: reasoningText } - } + for await (const apiChunk of stream) { + const delta = apiChunk.choices[0]?.delta ?? {} + const finishReason = apiChunk.choices[0]?.finish_reason - if (delta.content) { - let newText = delta.content - if (newText.startsWith(fullContent)) { - newText = newText.substring(fullContent.length) + const reasoningText = extractReasoningFromDelta(delta) + if (reasoningText) { + yield { type: "reasoning", text: reasoningText } } - fullContent = delta.content - - if (newText) { - // Check for thinking blocks - if (newText.includes("") || newText.includes("")) { - // Simple parsing for thinking blocks - const parts = newText.split(/<\/?think>/g) - for (let i = 0; i < parts.length; i++) { - if (parts[i]) { - if (i % 2 === 0) { - // Outside thinking block - yield { - type: "text", - text: parts[i], - } - } else { - // Inside thinking block - yield { - type: "reasoning", - text: parts[i], + + if (delta.content) { + let newText = delta.content + if (newText.startsWith(fullContent)) { + newText = newText.substring(fullContent.length) + } + fullContent = delta.content + + if (newText) { + // Check for thinking blocks + // Stryker disable next-line ConditionalExpression,StringLiteral: for tag-free content split(/<\/?think>/g) yields a single even-indexed part producing the identical text chunk as the else branch, and think-tagged content already takes this branch + if (newText.includes("") || newText.includes("")) { + // Simple parsing for thinking blocks + const parts = newText.split(/<\/?think>/g) + // Stryker disable next-line EqualityOperator: the extra i === parts.length iteration reads undefined, which the existing parts[i] guard on the next line skips + for (let i = 0; i < parts.length; i++) { + if (parts[i]) { + if (i % 2 === 0) { + // Outside thinking block + yield { + type: "text", + text: parts[i], + } + } else { + // Inside thinking block + yield { + type: "reasoning", + text: parts[i], + } } } } + } else { + yield { + type: "text", + text: newText, + } } - } else { + } + } + + // Handle tool calls in stream - emit partial chunks for NativeToolCallParser + if (delta.tool_calls) { + for (const toolCall of delta.tool_calls) { yield { - type: "text", - text: newText, + type: "tool_call_partial", + index: toolCall.index, + id: toolCall.id, + name: toolCall.function?.name, + arguments: toolCall.function?.arguments, } } } - } - // Handle tool calls in stream - emit partial chunks for NativeToolCallParser - if (delta.tool_calls) { - for (const toolCall of delta.tool_calls) { - yield { - type: "tool_call_partial", - index: toolCall.index, - id: toolCall.id, - name: toolCall.function?.name, - arguments: toolCall.function?.arguments, + // Process finish_reason to emit tool_call_end events + // Stryker disable next-line ConditionalExpression: processFinishReason only emits end events when finishReason is exactly "tool_calls" and the raw chunk tracker is non-empty, so entering the block with a falsy finishReason yields nothing + if (finishReason) { + const endEvents = NativeToolCallParser.processFinishReason(finishReason) + for (const event of endEvents) { + yield event } } - } - // Process finish_reason to emit tool_call_end events - if (finishReason) { - const endEvents = NativeToolCallParser.processFinishReason(finishReason) - for (const event of endEvents) { - yield event + if (apiChunk.usage) { + yield { + type: "usage", + inputTokens: apiChunk.usage.prompt_tokens || 0, + outputTokens: apiChunk.usage.completion_tokens || 0, + } } } - - if (apiChunk.usage) { - yield { - type: "usage", - inputTokens: apiChunk.usage.prompt_tokens || 0, - outputTokens: apiChunk.usage.completion_tokens || 0, - } + } catch (error) { + if (isRequestAborted(error, externalSignal)) { + throw createAbortError("Qwen Code") + } + throw error + } finally { + // Cancel the in-flight SDK request if the consumer stopped iterating + // early (break or return): the external signal may never fire in that + // case, and only the request-local signal reaches the SDK. + requestController.abort() + if (externalSignal) { + externalSignal.removeEventListener("abort", onExternalAbort) } } } @@ -328,7 +413,27 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { - await this.ensureAuthenticated() + // Fast-fail if the caller’s stop signal already fired before we started. + throwIfAborted(options?.abortSignal) + + // Merge the external stop signal with an optional per-call timeout. A + // timeoutMs <= 0 means "no explicit timeout" inside the util, so zero + // never reaches the SDK as an explicit timeout. + const requestSignal = mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs) + + // CompletePromptOptions is not createMessage metadata (no taskId), so + // the generic builder takes the merged signal via setOption instead of + // setAbortSignal(metadata). + const createOptions = new RequestConfigBuilder() + .setOption("signal", requestSignal) + .build() + + // A stop or timeout can land while the credential load or the token + // refresh below is still in flight; race the auth wait against the + // merged signal so the caller settles promptly instead of waiting for + // credential I/O. The shared refresh keeps running — only this wait + // is cut. + await settleOnAbort(this.ensureAuthenticated(), requestSignal, "Qwen Code") const client = this.ensureClient() const model = this.getModel() @@ -338,7 +443,13 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan max_completion_tokens: model.info.maxTokens, } - const response = await this.callApiWithRetry(() => client.chat.completions.create(requestOptions)) + // The retry reuses the captured request options, so it carries the same + // merged signal — and the guard in callApiWithRetry refuses to retry + // once this signal has aborted. + const response = await this.callApiWithRetry( + () => client.chat.completions.create(requestOptions, createOptions), + requestSignal, + ) return response.choices[0]?.message.content || "" } diff --git a/src/api/providers/utils/__tests__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index 1692f71e63..2a92bcd47f 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -3,6 +3,7 @@ import { isRequestAborted, mergeAbortSignalAndTimeout, mergeAbortSignals, + settleOnAbort, throwIfAborted, } from "../abort-signal" @@ -130,6 +131,8 @@ describe("abort-signal utilities", () => { expect(caught).toBeInstanceOf(Error) expect((caught as Error).name).toBe("AbortError") + // The exact message is part of the abort contract: callers (Task.ts, + // provider guards) must be able to recognize this error shape. expect((caught as Error).message).toBe("This operation was aborted") }) }) @@ -186,4 +189,124 @@ describe("abort-signal utilities", () => { expect(createAbortError("X")).not.toBe(createAbortError("X")) }) }) + + describe("settleOnAbort", () => { + it("returns the pending promise unchanged when the signal is undefined", async () => { + expect(await settleOnAbort(Promise.resolve(42), undefined, "Test")).toBe(42) + }) + + it("rejects immediately when the signal is already aborted", async () => { + // A listener registered on an already-aborted signal never fires, + // so the aborted state must be checked up front; the pending work + // keeps running (a helper that skipped the check would resolve + // with the sentinel instead of rejecting). + const controller = new AbortController() + controller.abort() + const pending = new Promise((resolve) => { + setTimeout(() => resolve(1), 20) + }) + + let caught: unknown + try { + await settleOnAbort(pending, controller.signal, "Qwen Code") + } catch (error) { + caught = error + } + await new Promise((resolve) => setTimeout(resolve, 30)) // let the sentinel arrive + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The Qwen Code request was aborted") + }) + + it("resolves with the pending value when it settles before the signal aborts", async () => { + const controller = new AbortController() + const pending = new Promise((resolve) => { + setTimeout(() => resolve(7), 10) + }) + + expect(await settleOnAbort(pending, controller.signal, "Test")).toBe(7) + }) + + it("rejects with the abort contract error when the signal aborts before the pending promise settles", async () => { + const controller = new AbortController() + let resolvePending!: (value: number) => void + const pending = new Promise((resolve) => { + resolvePending = resolve + }) + const racing = settleOnAbort(pending, controller.signal, "LM Studio") + controller.abort() + + let caught: unknown + try { + await racing + } catch (error) { + caught = error + } + resolvePending(1) // the underlying work still settles; it must not leak a rejection + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect((caught as Error).message).toBe("The LM Studio request was aborted") + }) + + it("removes the abort listener once the pending promise settles", async () => { + const controller = new AbortController() + const addSpy = vi.spyOn(controller.signal, "addEventListener") + const removeSpy = vi.spyOn(controller.signal, "removeEventListener") + + expect(await settleOnAbort(Promise.resolve("done"), controller.signal, "Test")).toBe("done") + + expect(addSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeSpy).toHaveBeenCalledTimes(1) + expect(removeSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + + it("removes the abort listener when the pending promise settles after an abort", async () => { + const controller = new AbortController() + let resolvePending!: (value: number) => void + const pending = new Promise((resolve) => { + resolvePending = resolve + }) + const addSpy = vi.spyOn(controller.signal, "addEventListener") + const removeSpy = vi.spyOn(controller.signal, "removeEventListener") + const racing = settleOnAbort(pending, controller.signal, "Test") + controller.abort() + + let caught: unknown + try { + await racing + } catch (error) { + caught = error + } + resolvePending(1) // settles the underlying work, which triggers the cleanup + await Promise.resolve() + + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).name).toBe("AbortError") + expect(addSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeSpy).toHaveBeenCalledTimes(1) + }) + + it("propagates a pending promise rejection unchanged and detaches the abort listener", async () => { + const controller = new AbortController() + const addSpy = vi.spyOn(controller.signal, "addEventListener") + const removeSpy = vi.spyOn(controller.signal, "removeEventListener") + const failure = new Error("count failed") + + let caught: unknown + try { + await settleOnAbort(Promise.reject(failure), controller.signal, "Test") + } catch (error) { + caught = error + } + + expect(caught).toBe(failure) + // The rejection path must detach the listener too: a leaked listener + // would keep this helper's closure alive for the life of the signal. + expect(addSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + expect(removeSpy).toHaveBeenCalledTimes(1) + expect(removeSpy).toHaveBeenCalledWith("abort", expect.any(Function)) + }) + }) }) diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index 26f57c3e9a..70c0940463 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -93,3 +93,42 @@ export function createAbortError(providerName: string): Error { abortError.name = "AbortError" return abortError } + +/** + * Await pending work, but reject with a request abort error the moment the + * signal aborts. The underlying promise keeps running to completion — only + * the wait is cut — so shared work (a deduped token refresh, for example) is + * never cancelled by one request's signal. + * + * An undefined signal skips the race entirely, and an already-aborted signal + * rejects immediately: listeners registered on an already-aborted signal + * never fire, so the aborted state must be checked up front. + */ +export function settleOnAbort( + pending: Promise, + signal: AbortSignal | undefined, + providerName: string, +): Promise { + if (!signal) { + return pending + } + if (signal.aborted) { + return Promise.reject(createAbortError(providerName)) + } + return new Promise((resolve, reject) => { + const onAbort = () => reject(createAbortError(providerName)) + signal.addEventListener("abort", onAbort) + void pending.then( + (value) => { + // Stryker disable next-line StringLiteral: mirrors the event name registered above, whose path is covered by the abort tests; a mutated removal event is unobservable because a signal cannot re-dispatch "abort" + signal.removeEventListener("abort", onAbort) + resolve(value) + }, + (error) => { + // Stryker disable next-line StringLiteral: mirrors the event name registered above, whose path is covered by the abort tests; a mutated removal event is unobservable because a signal cannot re-dispatch "abort" + signal.removeEventListener("abort", onAbort) + reject(error) + }, + ) + }) +} diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index 393e108645..084163229d 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -169,11 +169,6 @@ "count": 36 } }, - "api/providers/__tests__/lm-studio-timeout.spec.ts": { - "@typescript-eslint/no-explicit-any": { - "count": 3 - } - }, "api/providers/__tests__/mimo.spec.ts": { "@typescript-eslint/no-explicit-any": { "count": 18 @@ -414,11 +409,6 @@ "count": 3 } }, - "api/providers/qwen-code.ts": { - "@typescript-eslint/no-explicit-any": { - "count": 1 - } - }, "api/providers/requesty.ts": { "@typescript-eslint/no-explicit-any": { "count": 3