Skip to content

Commit 5f582a5

Browse files
easonLiangWorldedtecheasonliang28edelauna
authored
feat(api): add CompletePromptOptions parameter to completePrompt method (Zoo-Code-Org#901)
* feat(api): add CompletePromptOptions parameter to completePrompt method Add optional second parameter to all provider completePrompt methods for future abort signal and timeout support. - Add CompletePromptOptions interface with abortSignal and timeoutMs fields - Update 25+ provider implementations to accept options?: CompletePromptOptions parameter - Simplify providers to use basic request calls without signal/timeout logic * test(fake-ai): add test coverage for countTokens and completePrompt delegation * refactor(Task.spec): simplify mockStreamFactory to standard async generator * fix: CompletePromptOptions naming align Co-authored-by: edelauna <54631123+edelauna@users.noreply.github.com> * fix: remove inline import in providers\openai.ts Co-authored-by: edelauna <54631123+edelauna@users.noreply.github.com> --------- Co-authored-by: Eason Liang <easonliang28@gmail.com> Co-authored-by: edelauna <54631123+edelauna@users.noreply.github.com>
1 parent 3a4006a commit 5f582a5

31 files changed

Lines changed: 238 additions & 70 deletions

‎src/api/index.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,17 @@ import {
4242
} from "./providers"
4343
import { NativeOllamaHandler } from "./providers/native-ollama"
4444

45+
/**
46+
* Options for completePrompt — unified with ApiHandlerCreateMessageMetadata.
47+
* Uses abortSignal (not signal) to match the metadata pattern used in stream path.
48+
*/
49+
export interface CompletePromptOptions extends Pick<ApiHandlerCreateMessageMetadata, "abortSignal"> {
50+
/** Optional timeout override (ms) — falls back to provider default if omitted */
51+
timeoutMs?: number
52+
}
53+
4554
export interface SingleCompletionHandler {
46-
completePrompt(prompt: string): Promise<string>
55+
completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string>
4756
}
4857

4958
export interface ApiHandlerCreateMessageMetadata {

‎src/api/providers/__tests__/fake-ai.spec.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
import type { Anthropic } from "@anthropic-ai/sdk"
2+
3+
import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../../index"
4+
import { FakeAIHandler } from "../fake-ai"
5+
6+
const modelInfo = {
7+
contextWindow: 8192,
8+
maxTokens: 4096,
9+
supportsImages: false,
10+
supportsPromptCache: false,
11+
}
12+
13+
describe("FakeAIHandler", () => {
14+
it("should delegate completePrompt with options to the cached FakeAI instance", async () => {
15+
const completePrompt = vitest.fn().mockResolvedValue("delegated response")
16+
const fakeAi: {
17+
id: string
18+
createMessage: () => AsyncGenerator<never, void, unknown>
19+
getModel: () => { id: string; info: typeof modelInfo }
20+
countTokens: ReturnType<typeof vitest.fn>
21+
completePrompt: typeof completePrompt
22+
removeFromCache?: () => void
23+
} = {
24+
id: "fake-ai-completePrompt-delegation",
25+
createMessage: async function* () {},
26+
getModel: () => ({ id: "fake-model", info: modelInfo }),
27+
countTokens: vitest.fn().mockResolvedValue(0),
28+
completePrompt,
29+
}
30+
const controller = new AbortController()
31+
const options: CompletePromptOptions = { abortSignal: controller.signal, timeoutMs: 1234 }
32+
33+
const handler = new FakeAIHandler({ fakeAi })
34+
const result = await handler.completePrompt("Test prompt", options)
35+
36+
expect(result).toBe("delegated response")
37+
expect(completePrompt).toHaveBeenCalledWith("Test prompt", options)
38+
fakeAi.removeFromCache?.()
39+
})
40+
41+
it("should delegate createMessage, getModel, and countTokens to FakeAI", async () => {
42+
const metadata = { taskId: "task-1" } as ApiHandlerCreateMessageMetadata
43+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
44+
const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello" }]
45+
const createMessage = vitest.fn(async function* () {
46+
yield { type: "text" as const, text: "Hello" }
47+
})
48+
const getModel = vitest.fn(() => ({ id: "fake-model", info: modelInfo }))
49+
const countTokens = vitest.fn().mockResolvedValue(7)
50+
const fakeAi: {
51+
id: string
52+
createMessage: typeof createMessage
53+
getModel: typeof getModel
54+
countTokens: typeof countTokens
55+
completePrompt: ReturnType<typeof vitest.fn>
56+
removeFromCache?: () => void
57+
} = {
58+
id: "fake-ai-handler-delegation",
59+
createMessage,
60+
getModel,
61+
countTokens,
62+
completePrompt: vitest.fn().mockResolvedValue("complete"),
63+
}
64+
65+
const handler = new FakeAIHandler({ fakeAi })
66+
const chunks = []
67+
for await (const chunk of handler.createMessage("System", messages, metadata)) {
68+
chunks.push(chunk)
69+
}
70+
71+
expect(chunks).toEqual([{ type: "text", text: "Hello" }])
72+
expect(createMessage).toHaveBeenCalledWith("System", messages, metadata)
73+
expect(handler.getModel()).toEqual({ id: "fake-model", info: modelInfo })
74+
expect(getModel).toHaveBeenCalledTimes(1)
75+
await expect(handler.countTokens(content)).resolves.toBe(7)
76+
expect(countTokens).toHaveBeenCalledWith(content)
77+
fakeAi.removeFromCache?.()
78+
})
79+
})

‎src/api/providers/anthropic-vertex.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import {
2525

2626
import { BaseProvider } from "./base-provider"
2727
import { parseVertexJsonCredentials } from "./utils/vertex-credentials"
28-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
28+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
2929

3030
// https://docs.anthropic.com/en/api/claude-on-vertex-ai
3131
export class AnthropicVertexHandler extends BaseProvider implements SingleCompletionHandler {
@@ -270,7 +270,7 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
270270
}
271271
}
272272

273-
async completePrompt(prompt: string) {
273+
async completePrompt(prompt: string, options?: CompletePromptOptions) {
274274
try {
275275
const {
276276
id,

‎src/api/providers/anthropic.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import { getAnthropicProviderReasoning } from "../transform/reasoning"
2222
import { handleProviderError } from "./utils/error-handler"
2323

2424
import { BaseProvider } from "./base-provider"
25-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
25+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
2626
import { calculateApiCostAnthropic } from "../../shared/cost"
2727
import {
2828
convertOpenAIToolsToAnthropic,
@@ -429,7 +429,7 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
429429
}
430430
}
431431

432-
async completePrompt(prompt: string) {
432+
async completePrompt(prompt: string, options?: CompletePromptOptions) {
433433
const { id: model, temperature } = this.getModel()
434434

435435
let message

‎src/api/providers/base-openai-compatible-provider.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { TagMatcher } from "../../utils/tag-matcher"
88
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
99
import { convertToOpenAiMessages } from "../transform/openai-format"
1010

11-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
11+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
1212
import { DEFAULT_HEADERS } from "./constants"
1313
import { BaseProvider } from "./base-provider"
1414
import { handleOpenAIError } from "./utils/error-handler"
@@ -212,7 +212,7 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
212212
}
213213
}
214214

215-
async completePrompt(prompt: string): Promise<string> {
215+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
216216
const { id: modelId, info: modelInfo } = this.getModel()
217217

218218
const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = {

‎src/api/providers/bedrock.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ import { convertToBedrockConverseMessages as sharedConverter } from "../transfor
4444
import { getModelParams } from "../transform/model-params"
4545
import { shouldUseReasoningBudget } from "../../shared/api"
4646
import { normalizeToolSchema } from "../../utils/json-schema"
47-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
47+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
4848

4949
/************************************************************************************
5050
*
@@ -798,7 +798,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
798798
}
799799
}
800800

801-
async completePrompt(prompt: string): Promise<string> {
801+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
802802
try {
803803
const modelConfig = this.getModel()
804804

‎src/api/providers/fake-ai.ts‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { Anthropic } from "@anthropic-ai/sdk"
22

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

5-
import type { ApiHandler, SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
5+
import type {
6+
ApiHandler,
7+
SingleCompletionHandler,
8+
ApiHandlerCreateMessageMetadata,
9+
CompletePromptOptions,
10+
} from "../index"
611
import type { ApiHandlerOptions } from "../../shared/api"
712
import { ApiStream } from "../transform/stream"
813

@@ -28,7 +33,7 @@ interface FakeAI {
2833
): ApiStream
2934
getModel(): { id: string; info: ModelInfo }
3035
countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number>
31-
completePrompt(prompt: string): Promise<string>
36+
completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string>
3237
}
3338

3439
/**
@@ -75,7 +80,7 @@ export class FakeAIHandler implements ApiHandler, SingleCompletionHandler {
7580
return this.ai.countTokens(content)
7681
}
7782

78-
completePrompt(prompt: string): Promise<string> {
79-
return this.ai.completePrompt(prompt)
83+
completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
84+
return this.ai.completePrompt(prompt, options)
8085
}
8186
}

‎src/api/providers/gemini.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { t } from "i18next"
2323
import type { ApiStream, GroundingSource } from "../transform/stream"
2424
import { getModelParams } from "../transform/model-params"
2525

26-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
26+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
2727
import { BaseProvider } from "./base-provider"
2828
import { parseVertexJsonCredentials } from "./utils/vertex-credentials"
2929

@@ -576,7 +576,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
576576
return citationLinks.join(", ")
577577
}
578578

579-
async completePrompt(prompt: string): Promise<string> {
579+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
580580
const { id: model, info } = this.getModel()
581581

582582
try {

‎src/api/providers/kenari.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { ApiHandlerOptions } from "../../shared/api"
1313
import { ApiStream } from "../transform/stream"
1414
import { convertToOpenAiMessages } from "../transform/openai-format"
1515

16-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
16+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
1717
import { RouterProvider } from "./router-provider"
1818
import { extractReasoningFromDelta } from "./utils/extract-reasoning"
1919

@@ -123,7 +123,7 @@ export class KenariHandler extends RouterProvider implements SingleCompletionHan
123123
* @returns The model's reply text, or an empty string if no content is returned.
124124
* @throws Error with a Kenari-specific prefix if the request fails.
125125
*/
126-
async completePrompt(prompt: string): Promise<string> {
126+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
127127
const { id: modelId, info } = await this.fetchModel()
128128

129129
try {

‎src/api/providers/lite-llm.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { convertToOpenAiMessages } from "../transform/openai-format"
1212
import { GEMINI_THOUGHT_SIGNATURE_BYPASS } from "../transform/gemini-format"
1313
import { sanitizeOpenAiCallId } from "../../utils/tool-id"
1414

15-
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
15+
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
1616
import { RouterProvider } from "./router-provider"
1717
import { extractReasoningFromDelta } from "./utils/extract-reasoning"
1818

@@ -311,7 +311,7 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa
311311
}
312312
}
313313

314-
async completePrompt(prompt: string): Promise<string> {
314+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
315315
const { id: modelId, info } = await this.fetchModel()
316316

317317
// Check if this is a GPT-5 model that requires max_completion_tokens instead of max_tokens

0 commit comments

Comments
 (0)