Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// npx vitest run api/providers/__tests__/base-openai-compatible-provider.spec.ts

import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import OpenAI, { APIUserAbortError } from "openai"

import type { ModelInfo } from "@roo-code/types"

Expand All @@ -14,6 +14,8 @@ const mockCreate = vi.fn()

// Mock OpenAI module
vi.mock("openai", () => ({
// Named export consumed by the provider for abort-error normalization
APIUserAbortError: class extends Error {},
default: vi.fn(function () {
return {
chat: {
Expand Down Expand Up @@ -49,6 +51,20 @@ class TestOpenAiCompatibleProvider extends BaseOpenAiCompatibleProvider<"test-mo
}
}

/**
* Captures the rejection of an operation as an Error. The abort contract always
* throws an Error; the guard keeps strict typing without a cast. Fails the test
* if the operation resolves.
*/
async function captureError(operation: Promise<unknown>): Promise<Error> {
try {
await operation
} catch (error) {
return error instanceof Error ? error : new Error(String(error))
}
throw new Error("Expected the operation to reject")
}

describe("BaseOpenAiCompatibleProvider", () => {
let handler: TestOpenAiCompatibleProvider

Expand Down Expand Up @@ -278,6 +294,128 @@ describe("BaseOpenAiCompatibleProvider", () => {
})
})

describe("abort signal wiring", () => {
it("should pass the metadata abort signal to the client request", async () => {
const controller = new AbortController()
mockCreate.mockImplementationOnce(() => asyncStreamFrom([]))

const stream = handler.createMessage("system prompt", [], {
taskId: "test-task",
abortSignal: controller.signal,
})
await stream.next()

expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), {
signal: controller.signal,
})
})

it("should reject before issuing any request when the abort signal is already aborted", async () => {
const controller = new AbortController()
controller.abort()

await expect(async () => {
for await (const _ of handler.createMessage("system prompt", [], {
taskId: "test-task",
abortSignal: controller.signal,
})) {
// consume
}
}).rejects.toMatchObject({ name: "AbortError", message: "This operation was aborted" })
expect(mockCreate).not.toHaveBeenCalled()
})

it("should normalize the SDK APIUserAbortError from the stream path into the abort contract", async () => {
mockCreate.mockImplementationOnce(() => {
throw new APIUserAbortError()
})

const result = await captureError(
(async () => {
for await (const _ of handler.createMessage("system prompt", [])) {
// consume
}
})(),
)

// The SDK error has name "Error" and a message ending in a period; the
// provider must rethrow the Task.ts contract shape instead.
expect(result.name).toBe("AbortError")
expect(result.message).toBe("TestProvider request aborted")
expect(result.message.endsWith("aborted")).toBe(true)
})

it("should still wrap non-abort request errors with the provider prefix", async () => {
mockCreate.mockImplementationOnce(() => {
throw new Error("boom")
})

const result = await captureError(
(async () => {
for await (const _ of handler.createMessage("system prompt", [])) {
// consume
}
})(),
)

expect(result.message).toBe("TestProvider completion error: boom")
})

it("should pass the completePrompt abort signal to the client request", async () => {
const controller = new AbortController()
mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })

const result = await handler.completePrompt("test prompt", { abortSignal: controller.signal })

expect(result).toBe("response")
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), {
signal: controller.signal,
})
})

it("should merge completePrompt timeoutMs into the request signal", async () => {
mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })

await handler.completePrompt("test prompt", { timeoutMs: 5000 })

const requestOptions = mockCreate.mock.calls.at(-1)?.[1]
expect(requestOptions?.signal).toBeInstanceOf(AbortSignal)
expect(requestOptions?.signal.aborted).toBe(false)
})

it("should not set a request signal for zero completePrompt timeoutMs", async () => {
mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "response" } }] })

await handler.completePrompt("test prompt", { timeoutMs: 0 })

expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: "test-model" }), undefined)
})

it("should reject before any request when the completePrompt signal is already aborted", async () => {
const controller = new AbortController()
controller.abort()

await expect(
handler.completePrompt("test prompt", { abortSignal: controller.signal }),
).rejects.toMatchObject({
name: "AbortError",
message: "This operation was aborted",
})
expect(mockCreate).not.toHaveBeenCalled()
})

it("should normalize the SDK APIUserAbortError from completePrompt", async () => {
mockCreate.mockImplementationOnce(() => {
throw new APIUserAbortError()
})

const result = await captureError(handler.completePrompt("test prompt"))

expect(result.name).toBe("AbortError")
expect(result.message).toBe("TestProvider request aborted")
})
})

describe("Tool call handling", () => {
it("should yield tool_call_end events when finish_reason is tool_calls", async () => {
mockCreate.mockImplementationOnce(() =>
Expand Down
29 changes: 29 additions & 0 deletions src/api/providers/__tests__/complete-prompt-options.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
import { describe, it, expect } from "vitest"

import type { CompletePromptOptions } from "../../index"

describe("CompletePromptOptions", () => {
it("should allow abortSignal property", () => {
const controller = new AbortController()
const options: CompletePromptOptions = { abortSignal: controller.signal }
expect(options.abortSignal).toBe(controller.signal)
})

it("should allow timeoutMs property", () => {
const options: CompletePromptOptions = { timeoutMs: 5000 }
expect(options.timeoutMs).toBe(5000)
})

it("should allow both abortSignal and timeoutMs together", () => {
const controller = new AbortController()
const options: CompletePromptOptions = { abortSignal: controller.signal, timeoutMs: 10000 }
expect(options.abortSignal).toBe(controller.signal)
expect(options.timeoutMs).toBe(10000)
})

it("should allow empty options object", () => {
const options: CompletePromptOptions = {}
expect(options.abortSignal).toBeUndefined()
expect(options.timeoutMs).toBeUndefined()
})
})
114 changes: 112 additions & 2 deletions src/api/providers/__tests__/kimi-code.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,32 @@ vi.mock("../../../integrations/kimi-code/oauth", () => ({

vi.mock("../fetchers/modelCache", () => ({ getModels: mockGetModels }))

/**
* Spies on the inherited OpenAI client's chat.completions.create. `client` is
* protected on the OpenAiHandler base (not on the public interface), so it is
* reached through a documented `as unknown as` double assertion (AGENTS.md
* last resort; no `as any`).
*/
function completionsCreate(handler: KimiCodeHandler): ReturnType<typeof vi.fn> {
const client = (
handler as unknown as {
client: { chat: { completions: Record<string, (...args: never[]) => never> } }
}
).client
return vi.spyOn(client.chat.completions, "create") as unknown as ReturnType<typeof vi.fn>
}

/** Captures the rejection of an operation; fails the test if it resolves. */
async function captureError(operation: Promise<unknown>): Promise<Error> {
try {
await operation
} catch (error) {
// Abort normalization always throws an Error; the guard keeps strict typing without casts.
return error instanceof Error ? error : new Error(String(error))
}
throw new Error("Expected the operation to reject")
}

describe("KimiCodeHandler", () => {
beforeEach(() => {
clearAllMocks()
Expand Down Expand Up @@ -117,8 +143,7 @@ describe("KimiCodeHandler", () => {
it("force-refreshes and retries exactly once after a non-streaming OAuth 401", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 })
const createCompletion = vi
.spyOn((handler as any).client.chat.completions, "create")
const createCompletion = completionsCreate(handler)
.mockRejectedValueOnce(unauthorized)
.mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] })

Expand Down Expand Up @@ -235,4 +260,89 @@ describe("KimiCodeHandler", () => {
})
expect(handler.getModel().reasoning).toEqual({ reasoning_effort: "max" })
})

it("forwards the metadata abort signal to the inherited OpenAI SDK request", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
const controller = new AbortController()
const streamChunks = (async function* () {
yield { choices: [{ delta: { content: "hi" } }] }
})()
const createCompletion = completionsCreate(handler).mockResolvedValueOnce(streamChunks)

const gen = handler.createMessage("system", [{ role: "user", content: "test" }], {
taskId: "test-task",
abortSignal: controller.signal,
})
const first = await gen.next()

expect(first.value).toEqual({ type: "text", text: "hi" })
expect(createCompletion).toHaveBeenCalledWith(expect.anything(), { signal: controller.signal })
})

it("rejects before any request when the createMessage abort signal is already aborted", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
const controller = new AbortController()
controller.abort()
const createCompletion = completionsCreate(handler)

const gen = handler.createMessage("system", [{ role: "user", content: "test" }], {
taskId: "test-task",
abortSignal: controller.signal,
})

await expect(async () => {
for await (const _ of gen) {
// consume
}
}).rejects.toMatchObject({ name: "AbortError", message: "This operation was aborted" })
expect(createCompletion).not.toHaveBeenCalled()
})

it("forwards completePrompt abort options through the override on both 401 retry attempts", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" })
const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 })
const createCompletion = completionsCreate(handler)
.mockRejectedValueOnce(unauthorized)
.mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] })
const controller = new AbortController()

await expect(handler.completePrompt("test", { abortSignal: controller.signal })).resolves.toBe("retried")
expect(mockForceRefreshAccessToken).toHaveBeenCalledOnce()
expect(createCompletion).toHaveBeenCalledTimes(2)
for (const call of createCompletion.mock.calls) {
expect(call[1]).toEqual({ signal: controller.signal })
}
})

it("rejects before any request when the completePrompt signal is already aborted", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
const controller = new AbortController()
controller.abort()
const createCompletion = completionsCreate(handler)

await expect(handler.completePrompt("test", { abortSignal: controller.signal })).rejects.toMatchObject({
name: "AbortError",
message: "This operation was aborted",
})
expect(createCompletion).not.toHaveBeenCalled()
})

it("surfaces a normalized AbortError when the SDK aborts a streaming request", async () => {
const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" })
// The real SDK class: this spec does not mock the openai module.
const { APIUserAbortError } = await import("openai")
completionsCreate(handler).mockRejectedValueOnce(new APIUserAbortError())

const gen = handler.createMessage("system", [{ role: "user", content: "test" }], { taskId: "test-task" })
const result = await captureError(
(async () => {
for await (const _ of gen) {
// consume
}
})(),
)

expect(result.name).toBe("AbortError")
expect(result.message).toBe("OpenAI request aborted")
})
})
Loading
Loading