From bcedf5037399ed35d632f62db6a8f854ffabea17 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Wed, 15 Jul 2026 22:12:08 +0800 Subject: [PATCH 1/3] 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 --- src/api/index.ts | 11 ++- src/api/providers/anthropic-vertex.ts | 4 +- src/api/providers/anthropic.ts | 4 +- .../base-openai-compatible-provider.ts | 4 +- src/api/providers/bedrock.ts | 4 +- src/api/providers/fake-ai.ts | 13 ++- src/api/providers/gemini.ts | 4 +- src/api/providers/kenari.ts | 4 +- src/api/providers/lite-llm.ts | 4 +- src/api/providers/lm-studio.ts | 4 +- src/api/providers/minimax.ts | 4 +- src/api/providers/mistral.ts | 5 +- src/api/providers/native-ollama.ts | 4 +- src/api/providers/openai-codex.ts | 4 +- src/api/providers/openai-compatible.ts | 4 +- src/api/providers/openai-native.ts | 12 +-- src/api/providers/openai.ts | 2 +- src/api/providers/opencode-go.ts | 4 +- src/api/providers/openrouter.ts | 4 +- src/api/providers/poe.ts | 4 +- src/api/providers/qwen-code.ts | 4 +- src/api/providers/requesty.ts | 4 +- src/api/providers/unbound.ts | 4 +- src/api/providers/vercel-ai-gateway.ts | 4 +- src/api/providers/vscode-lm.ts | 4 +- src/api/providers/xai.ts | 4 +- src/api/providers/zoo-gateway.ts | 4 +- src/core/task/__tests__/Task.spec.ts | 97 ++++++++++++++++++- src/utils/__tests__/enhance-prompt.spec.ts | 4 +- src/utils/single-completion-handler.ts | 10 +- 30 files changed, 173 insertions(+), 69 deletions(-) diff --git a/src/api/index.ts b/src/api/index.ts index 5a9532b43d..098ab9e540 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -42,8 +42,17 @@ import { } from "./providers" import { NativeOllamaHandler } from "./providers/native-ollama" +/** + * Options for completePrompt — unified with ApiHandlerCreateMessageMetadata. + * Uses abortSignal (not signal) to match the metadata pattern used in stream path. + */ +export interface CompletePromptOptions extends Pick { + /** Optional timeout override (ms) — falls back to provider default if omitted */ + timeoutMs?: number +} + export interface SingleCompletionHandler { - completePrompt(prompt: string): Promise + completePrompt(prompt: string, metadata?: CompletePromptOptions): Promise } export interface ApiHandlerCreateMessageMetadata { diff --git a/src/api/providers/anthropic-vertex.ts b/src/api/providers/anthropic-vertex.ts index 4f62fe0b1e..9d91b25fe5 100644 --- a/src/api/providers/anthropic-vertex.ts +++ b/src/api/providers/anthropic-vertex.ts @@ -25,7 +25,7 @@ import { import { BaseProvider } from "./base-provider" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" // https://docs.anthropic.com/en/api/claude-on-vertex-ai export class AnthropicVertexHandler extends BaseProvider implements SingleCompletionHandler { @@ -270,7 +270,7 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple } } - async completePrompt(prompt: string) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { try { const { id, diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index 2728dae1ad..f69c6eb29c 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -22,7 +22,7 @@ import { getAnthropicProviderReasoning } from "../transform/reasoning" import { handleProviderError } from "./utils/error-handler" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { calculateApiCostAnthropic } from "../../shared/cost" import { convertOpenAIToolsToAnthropic, @@ -429,7 +429,7 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa } } - async completePrompt(prompt: string) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { const { id: model, temperature } = this.getModel() let message diff --git a/src/api/providers/base-openai-compatible-provider.ts b/src/api/providers/base-openai-compatible-provider.ts index d6406f2a67..f4928b0b0a 100644 --- a/src/api/providers/base-openai-compatible-provider.ts +++ b/src/api/providers/base-openai-compatible-provider.ts @@ -8,7 +8,7 @@ import { TagMatcher } from "../../utils/tag-matcher" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { convertToOpenAiMessages } from "../transform/openai-format" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" import { handleOpenAIError } from "./utils/error-handler" @@ -212,7 +212,7 @@ export abstract class BaseOpenAiCompatibleProvider } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, info: modelInfo } = this.getModel() const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = { diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 8c2b5ace68..ae072b3559 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -44,7 +44,7 @@ import { convertToBedrockConverseMessages as sharedConverter } from "../transfor import { getModelParams } from "../transform/model-params" import { shouldUseReasoningBudget } from "../../shared/api" import { normalizeToolSchema } from "../../utils/json-schema" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /************************************************************************************ * @@ -798,7 +798,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { const modelConfig = this.getModel() diff --git a/src/api/providers/fake-ai.ts b/src/api/providers/fake-ai.ts index e69a1c84e8..f6c466b454 100644 --- a/src/api/providers/fake-ai.ts +++ b/src/api/providers/fake-ai.ts @@ -2,7 +2,12 @@ import { Anthropic } from "@anthropic-ai/sdk" import type { ModelInfo } from "@roo-code/types" -import type { ApiHandler, SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { + ApiHandler, + SingleCompletionHandler, + ApiHandlerCreateMessageMetadata, + CompletePromptOptions, +} from "../index" import type { ApiHandlerOptions } from "../../shared/api" import { ApiStream } from "../transform/stream" @@ -28,7 +33,7 @@ interface FakeAI { ): ApiStream getModel(): { id: string; info: ModelInfo } countTokens(content: Array): Promise - completePrompt(prompt: string): Promise + completePrompt(prompt: string, options?: CompletePromptOptions): Promise } /** @@ -75,7 +80,7 @@ export class FakeAIHandler implements ApiHandler, SingleCompletionHandler { return this.ai.countTokens(content) } - completePrompt(prompt: string): Promise { - return this.ai.completePrompt(prompt) + completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + return this.ai.completePrompt(prompt, options) } } diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts index 6c8168cae2..44065d8801 100644 --- a/src/api/providers/gemini.ts +++ b/src/api/providers/gemini.ts @@ -23,7 +23,7 @@ import { t } from "i18next" import type { ApiStream, GroundingSource } from "../transform/stream" import { getModelParams } from "../transform/model-params" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { BaseProvider } from "./base-provider" import { parseVertexJsonCredentials } from "./utils/vertex-credentials" @@ -576,7 +576,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl return citationLinks.join(", ") } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: model, info } = this.getModel() try { diff --git a/src/api/providers/kenari.ts b/src/api/providers/kenari.ts index 75c02ff164..7895a9452c 100644 --- a/src/api/providers/kenari.ts +++ b/src/api/providers/kenari.ts @@ -13,7 +13,7 @@ import { ApiHandlerOptions } from "../../shared/api" import { ApiStream } from "../transform/stream" import { convertToOpenAiMessages } from "../transform/openai-format" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -123,7 +123,7 @@ export class KenariHandler extends RouterProvider implements SingleCompletionHan * @returns The model's reply text, or an empty string if no content is returned. * @throws Error with a Kenari-specific prefix if the request fails. */ - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, info } = await this.fetchModel() try { diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index 3f9b3732e5..49cc1b115b 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -12,7 +12,7 @@ import { convertToOpenAiMessages } from "../transform/openai-format" import { GEMINI_THOUGHT_SIGNATURE_BYPASS } from "../transform/gemini-format" import { sanitizeOpenAiCallId } from "../../utils/tool-id" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -311,7 +311,7 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, info } = await this.fetchModel() // Check if this is a GPT-5 model that requires max_completion_tokens instead of max_tokens diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index 891695b5fc..39216418db 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -13,7 +13,7 @@ import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { getModelsFromCache } from "./fetchers/modelCache" import { handleOpenAIError } from "./utils/error-handler" @@ -187,7 +187,7 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { // Create params object with optional draft model const params: any = { diff --git a/src/api/providers/minimax.ts b/src/api/providers/minimax.ts index 93aa7ea8f1..e209add72d 100644 --- a/src/api/providers/minimax.ts +++ b/src/api/providers/minimax.ts @@ -12,7 +12,7 @@ import { getModelParams } from "../transform/model-params" import { mergeEnvironmentDetailsForMiniMax } from "../transform/minimax-format" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { calculateApiCostAnthropic } from "../../shared/cost" import { convertOpenAIToolsToAnthropic } from "../../core/prompts/tools/native-tools/converters" @@ -289,7 +289,7 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand } } - async completePrompt(prompt: string) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { const { id: model, temperature } = this.getModel() const message = await this.client.messages.create({ diff --git a/src/api/providers/mistral.ts b/src/api/providers/mistral.ts index e0e19298f4..c7816feaa2 100644 --- a/src/api/providers/mistral.ts +++ b/src/api/providers/mistral.ts @@ -18,7 +18,7 @@ import { ApiStream } from "../transform/stream" import { handleProviderError } from "./utils/error-handler" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" // Type helper to handle thinking chunks from Mistral API // The SDK includes ThinkChunk but TypeScript has trouble with the discriminated union @@ -192,8 +192,7 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand return { id, info, maxTokens, temperature } } - - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: model, temperature } = this.getModel() try { diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index 2a14885e11..eb380d1eb2 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -7,7 +7,7 @@ import { BaseProvider } from "./base-provider" import type { ApiHandlerOptions } from "../../shared/api" import { getOllamaModels } from "./fetchers/ollama" import { TagMatcher } from "../../utils/tag-matcher" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" interface OllamaChatOptions { temperature: number @@ -550,7 +550,7 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { const client = this.ensureClient() const { id: modelId } = await this.fetchModel() diff --git a/src/api/providers/openai-codex.ts b/src/api/providers/openai-codex.ts index 2b2a599c0f..374a715b34 100644 --- a/src/api/providers/openai-codex.ts +++ b/src/api/providers/openai-codex.ts @@ -21,7 +21,7 @@ import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { isMcpTool } from "../../utils/mcp-name" import { sanitizeOpenAiCallId } from "../../utils/tool-id" import { openAiCodexOAuthManager } from "../../integrations/openai-codex/oauth" @@ -1154,7 +1154,7 @@ export class OpenAiCodexHandler extends BaseProvider implements SingleCompletion return this.lastResponseId } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { this.abortController = new AbortController() try { diff --git a/src/api/providers/openai-compatible.ts b/src/api/providers/openai-compatible.ts index d129e72452..308652804a 100644 --- a/src/api/providers/openai-compatible.ts +++ b/src/api/providers/openai-compatible.ts @@ -17,7 +17,7 @@ import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** * Configuration options for creating an OpenAI-compatible provider. @@ -197,7 +197,7 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si /** * Complete a prompt using the AI SDK generateText. */ - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const languageModel = this.getLanguageModel() const { text } = await generateText({ diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index 8f86f46751..a1ce2d89d0 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -26,7 +26,7 @@ import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { isMcpTool } from "../../utils/mcp-name" import { sanitizeOpenAiCallId } from "../../utils/tool-id" @@ -1484,14 +1484,12 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio getResponseId(): string | undefined { return this.lastResponseId } - - async completePrompt(prompt: string): Promise { - // Create AbortController for cancellation - this.abortController = new AbortController() - + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { + this.abortController = new AbortController() + const model = this.getModel() - const { verbosity, reasoning } = model + const { verbosity } = model // Resolve reasoning effort for models that support it const reasoningEffort = this.getReasoningEffort(model) diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 07427c7dbe..fd49b13557 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -296,7 +296,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl return { id, info, ...params } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: import("../../api").CompletePromptOptions): Promise { try { const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl) const model = this.getModel() diff --git a/src/api/providers/opencode-go.ts b/src/api/providers/opencode-go.ts index 27d8ab3f7e..be53dc1c02 100644 --- a/src/api/providers/opencode-go.ts +++ b/src/api/providers/opencode-go.ts @@ -18,7 +18,7 @@ import { convertToR1Format } from "../transform/r1-format" import { filterNonAnthropicBlocks } from "../transform/anthropic-filter" import { getModelParams } from "../transform/model-params" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" import { DEFAULT_HEADERS } from "./constants" @@ -485,7 +485,7 @@ export class OpencodeGoHandler extends RouterProvider implements SingleCompletio * @returns The model's reply text, or an empty string if no content is returned. * @throws Error with an Opencode Go-specific prefix if the request fails. */ - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, format, temperature, reasoningEffort, maxTokens } = await this.resolveModel() if (format === "anthropic") { diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 360d4c7f81..3e59b4360b 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -35,7 +35,7 @@ import { getModelEndpoints } from "./fetchers/modelEndpointCache" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" -import type { ApiHandlerCreateMessageMetadata, SingleCompletionHandler } from "../index" +import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions, SingleCompletionHandler } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { generateImageWithProvider, ImageGenerationResult } from "./utils/image-generation" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" @@ -574,7 +574,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH return { id, info, topP: isDeepSeekR1 ? 0.95 : undefined, ...params } } - async completePrompt(prompt: string) { + async completePrompt(prompt: string, options?: CompletePromptOptions) { const { id: modelId, maxTokens, temperature, reasoning } = await this.fetchModel() const completionParams: OpenRouterChatCompletionParams = { diff --git a/src/api/providers/poe.ts b/src/api/providers/poe.ts index 9513a1445d..1e5315b1ba 100644 --- a/src/api/providers/poe.ts +++ b/src/api/providers/poe.ts @@ -18,7 +18,7 @@ import { convertToAiSdkMessages, convertToolsForAiSdk, processAiSdkStreamPart } import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { getModelsFromCache } from "./fetchers/modelCache" const DEFAULT_THINKING_BUDGET = 8192 @@ -138,7 +138,7 @@ export class PoeHandler extends BaseProvider implements SingleCompletionHandler } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id } = this.getModel() try { const { text } = await generateText({ diff --git a/src/api/providers/qwen-code.ts b/src/api/providers/qwen-code.ts index cdf0f88e4b..5001b4c8ed 100644 --- a/src/api/providers/qwen-code.ts +++ b/src/api/providers/qwen-code.ts @@ -15,7 +15,7 @@ import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" import { extractReasoningFromDelta } from "./utils/extract-reasoning" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" const QWEN_OAUTH_BASE_URL = "https://chat.qwen.ai" const QWEN_OAUTH_TOKEN_ENDPOINT = `${QWEN_OAUTH_BASE_URL}/api/v1/oauth2/token` @@ -327,7 +327,7 @@ export class QwenCodeHandler extends BaseProvider implements SingleCompletionHan return { id, info } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { await this.ensureAuthenticated() const client = this.ensureClient() const model = this.getModel() diff --git a/src/api/providers/requesty.ts b/src/api/providers/requesty.ts index 006cd9cdfb..5753660de5 100644 --- a/src/api/providers/requesty.ts +++ b/src/api/providers/requesty.ts @@ -14,7 +14,7 @@ import { AnthropicProviderReasoningParams, getAnthropicProviderReasoning } from import { DEFAULT_HEADERS } from "./constants" import { getModels } from "./fetchers/modelCache" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { toRequestyServiceUrl } from "../../shared/utils/requesty" import { handleOpenAIError } from "./utils/error-handler" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" @@ -204,7 +204,7 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: model, maxTokens: max_tokens, temperature } = await this.fetchModel() const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [{ role: "system", content: prompt }] diff --git a/src/api/providers/unbound.ts b/src/api/providers/unbound.ts index cb41591634..c3ec9c44fc 100644 --- a/src/api/providers/unbound.ts +++ b/src/api/providers/unbound.ts @@ -14,7 +14,7 @@ import { OpenAiReasoningParams } from "../transform/reasoning" import { DEFAULT_HEADERS } from "./constants" import { getModels } from "./fetchers/modelCache" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -192,7 +192,7 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: model, maxTokens: max_tokens, temperature } = await this.fetchModel() const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [{ role: "system", content: prompt }] diff --git a/src/api/providers/vercel-ai-gateway.ts b/src/api/providers/vercel-ai-gateway.ts index 0c7bd1d485..0820e1b2d8 100644 --- a/src/api/providers/vercel-ai-gateway.ts +++ b/src/api/providers/vercel-ai-gateway.ts @@ -14,7 +14,7 @@ import { ApiStream } from "../transform/stream" import { convertToOpenAiMessages } from "../transform/openai-format" import { addCacheBreakpoints } from "../transform/caching/vercel-ai-gateway" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" // Extend OpenAI's CompletionUsage to include Vercel AI Gateway specific fields @@ -117,7 +117,7 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, info } = await this.fetchModel() try { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 5b9064363f..5c88a947de 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -12,7 +12,7 @@ import { ApiStream } from "../transform/stream" import { convertToVsCodeLmMessages, extractTextCountFromMessage } from "../transform/vscode-lm-format" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** * Converts OpenAI-format tools to VSCode Language Model tools. @@ -582,7 +582,7 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan return this.getModel().info.contextWindow } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { const client = await this.getClient() const response = await client.sendRequest( diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index 17ac40c8bf..110269cba9 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -13,7 +13,7 @@ import { getModelParams } from "../transform/model-params" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { isMcpTool } from "../../utils/mcp-name" @@ -142,7 +142,7 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler yield* processResponsesApiStream(stream, normalizeUsage) } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const model = this.getModel() try { diff --git a/src/api/providers/zoo-gateway.ts b/src/api/providers/zoo-gateway.ts index 4724464ff3..cea410bb2e 100644 --- a/src/api/providers/zoo-gateway.ts +++ b/src/api/providers/zoo-gateway.ts @@ -18,7 +18,7 @@ import { ApiStream } from "../transform/stream" import { convertToOpenAiMessages } from "../transform/openai-format" import { addCacheBreakpoints } from "../transform/caching/vercel-ai-gateway" -import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { RouterProvider } from "./router-provider" function getApiErrorStatus(error: unknown): number | undefined { @@ -276,7 +276,7 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio } } - async completePrompt(prompt: string): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { this.ensureAuthenticated() const { id: modelId, info } = await this.fetchModel() diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 27ba5ce8ff..fbc01cb78f 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -1812,7 +1812,7 @@ describe("Cline", () => { expect(summarizeConversation).toHaveBeenCalled() const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)! - expect(options.metadata?.abortSignal).toBeInstanceOf(AbortSignal) + expect(options.metadata?.abortSignal).toBe(task.currentRequestAbortController!.signal) }) it("should omit abortSignal from condenseContext metadata when no current request exists", async () => { @@ -1899,7 +1899,7 @@ describe("Cline", () => { const [, , metadata] = createMessageSpy.mock.calls[0]! expect(metadata).toBeDefined() - expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal) + expect(metadata!.abortSignal).toBe(task.currentRequestAbortController!.signal) }) it("should invoke abort on currentRequestAbortController during first-chunk wait", async () => { @@ -2136,7 +2136,7 @@ describe("Cline", () => { const [, , metadata] = createMessageSpy.mock.calls[0]! expect(metadata).toBeDefined() expect("abortSignal" in metadata!).toBe(true) - expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal) + expect(metadata!.abortSignal).toBe(task.currentRequestAbortController!.signal) }) it("should keep createMessage abortSignal metadata unaborted before cancellation", async () => { @@ -2197,11 +2197,100 @@ describe("Cline", () => { await iterator.next() const [, , metadata] = createMessageSpy.mock.calls[0]! - expect(metadata?.abortSignal).toBeInstanceOf(AbortSignal) + expect(metadata?.abortSignal).toBe(task.currentRequestAbortController!.signal) expect(metadata?.abortSignal?.aborted).toBe(false) }) }) + it("should create a fresh AbortController for each sequential request", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + let callCount = 0 + const mockStreamFactory = () => { + return { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: `response ${callCount}` } + }, + async next() { + callCount++ + return { done: true, value: { type: "text", text: `response ${callCount - 1}` } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + } + + const createMessageSpy = vi + .spyOn(task.api, "createMessage") + .mockImplementation(() => mockStreamFactory()) + + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + // First request + const iterator1 = task.attemptApiRequest(0) + await iterator1.next() + + expect(createMessageSpy).toHaveBeenCalledTimes(1) + const [, , metadata1] = createMessageSpy.mock.calls[0]! + const signal1 = metadata1!.abortSignal + expect(signal1).toBeDefined() + expect(signal1!.aborted).toBe(false) + + // Simulate request completion and cancellation to clear the controller + task.cancelCurrentRequest() + + // Second request should create a fresh AbortController with a new signal + callCount = 0 + const iterator2 = task.attemptApiRequest(0) + await iterator2.next() + + expect(createMessageSpy).toHaveBeenCalledTimes(2) + const [, , metadata2] = createMessageSpy.mock.calls[1]! + const signal2 = metadata2!.abortSignal + + // Signals should be different instances (fresh controller per request) + expect(signal2).not.toBe(signal1) + expect(signal2).toBe(task.currentRequestAbortController!.signal) + expect(signal2!.aborted).toBe(false) + }) + it("should propagate AbortController signal through attemptApiRequest context-window retry path", async () => { const task = new Task({ provider: mockProvider, diff --git a/src/utils/__tests__/enhance-prompt.spec.ts b/src/utils/__tests__/enhance-prompt.spec.ts index 2546878d8c..7e8c702984 100644 --- a/src/utils/__tests__/enhance-prompt.spec.ts +++ b/src/utils/__tests__/enhance-prompt.spec.ts @@ -42,7 +42,7 @@ describe("enhancePrompt", () => { expect(result).toBe("Enhanced prompt") const handler = buildApiHandler(mockApiConfig) - expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`) + expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`, undefined) }) it("enhances prompt using custom enhancement prompt when provided", async () => { @@ -64,7 +64,7 @@ describe("enhancePrompt", () => { expect(result).toBe("Enhanced prompt") const handler = buildApiHandler(mockApiConfig) - expect((handler as any).completePrompt).toHaveBeenCalledWith(`${customEnhancePrompt}\n\nTest prompt`) + expect((handler as any).completePrompt).toHaveBeenCalledWith(`${customEnhancePrompt}\n\nTest prompt`, undefined) }) it("throws error for empty prompt input", async () => { diff --git a/src/utils/single-completion-handler.ts b/src/utils/single-completion-handler.ts index 4606a17bab..4890b38471 100644 --- a/src/utils/single-completion-handler.ts +++ b/src/utils/single-completion-handler.ts @@ -1,12 +1,16 @@ import type { ProviderSettings } from "@roo-code/types" -import { buildApiHandler, SingleCompletionHandler } from "../api" +import { buildApiHandler, SingleCompletionHandler, type CompletePromptOptions } from "../api" /** * Enhances a prompt using the configured API without creating a full Cline instance or task history. * This is a lightweight alternative that only uses the API's completion functionality. */ -export async function singleCompletionHandler(apiConfiguration: ProviderSettings, promptText: string): Promise { +export async function singleCompletionHandler( + apiConfiguration: ProviderSettings, + promptText: string, + options?: CompletePromptOptions, +): Promise { if (!promptText) { throw new Error("No prompt text provided") } @@ -21,5 +25,5 @@ export async function singleCompletionHandler(apiConfiguration: ProviderSettings throw new Error("The selected API provider does not support prompt enhancement") } - return (handler as SingleCompletionHandler).completePrompt(promptText) + return (handler as SingleCompletionHandler).completePrompt(promptText, options) } From 7785da5a2e0a04ec0e2ed4e47a8eab141d2bc5ab Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Wed, 15 Jul 2026 22:44:44 +0800 Subject: [PATCH 2/3] test(fake-ai): add test coverage for countTokens and completePrompt delegation --- src/api/providers/__tests__/fake-ai.spec.ts | 79 +++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 src/api/providers/__tests__/fake-ai.spec.ts diff --git a/src/api/providers/__tests__/fake-ai.spec.ts b/src/api/providers/__tests__/fake-ai.spec.ts new file mode 100644 index 0000000000..955264c68f --- /dev/null +++ b/src/api/providers/__tests__/fake-ai.spec.ts @@ -0,0 +1,79 @@ +import type { Anthropic } from "@anthropic-ai/sdk" + +import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../../index" +import { FakeAIHandler } from "../fake-ai" + +const modelInfo = { + contextWindow: 8192, + maxTokens: 4096, + supportsImages: false, + supportsPromptCache: false, +} + +describe("FakeAIHandler", () => { + it("should delegate completePrompt with options to the cached FakeAI instance", async () => { + const completePrompt = vitest.fn().mockResolvedValue("delegated response") + const fakeAi: { + id: string + createMessage: () => AsyncGenerator + getModel: () => { id: string; info: typeof modelInfo } + countTokens: ReturnType + completePrompt: typeof completePrompt + removeFromCache?: () => void + } = { + id: "fake-ai-completePrompt-delegation", + createMessage: async function* () {}, + getModel: () => ({ id: "fake-model", info: modelInfo }), + countTokens: vitest.fn().mockResolvedValue(0), + completePrompt, + } + const controller = new AbortController() + const options: CompletePromptOptions = { abortSignal: controller.signal, timeoutMs: 1234 } + + const handler = new FakeAIHandler({ fakeAi }) + const result = await handler.completePrompt("Test prompt", options) + + expect(result).toBe("delegated response") + expect(completePrompt).toHaveBeenCalledWith("Test prompt", options) + fakeAi.removeFromCache?.() + }) + + it("should delegate createMessage, getModel, and countTokens to FakeAI", async () => { + const metadata = { taskId: "task-1" } as ApiHandlerCreateMessageMetadata + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello" }] + const createMessage = vitest.fn(async function* () { + yield { type: "text" as const, text: "Hello" } + }) + const getModel = vitest.fn(() => ({ id: "fake-model", info: modelInfo })) + const countTokens = vitest.fn().mockResolvedValue(7) + const fakeAi: { + id: string + createMessage: typeof createMessage + getModel: typeof getModel + countTokens: typeof countTokens + completePrompt: ReturnType + removeFromCache?: () => void + } = { + id: "fake-ai-handler-delegation", + createMessage, + getModel, + countTokens, + completePrompt: vitest.fn().mockResolvedValue("complete"), + } + + const handler = new FakeAIHandler({ fakeAi }) + const chunks = [] + for await (const chunk of handler.createMessage("System", messages, metadata)) { + chunks.push(chunk) + } + + expect(chunks).toEqual([{ type: "text", text: "Hello" }]) + expect(createMessage).toHaveBeenCalledWith("System", messages, metadata) + expect(handler.getModel()).toEqual({ id: "fake-model", info: modelInfo }) + expect(getModel).toHaveBeenCalledTimes(1) + await expect(handler.countTokens(content)).resolves.toBe(7) + expect(countTokens).toHaveBeenCalledWith(content) + fakeAi.removeFromCache?.() + }) +}) From 62ed671e165d5b38f909d8957cb12aeeebf709c4 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Wed, 15 Jul 2026 23:40:26 +0800 Subject: [PATCH 3/3] refactor(Task.spec): simplify mockStreamFactory to standard async generator --- src/core/task/__tests__/Task.spec.ts | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index fbc01cb78f..0c784821c1 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -2232,25 +2232,10 @@ describe("Cline", () => { }) let callCount = 0 - const mockStreamFactory = () => { - return { - async *[Symbol.asyncIterator]() { - yield { type: "text", text: `response ${callCount}` } - }, - async next() { - callCount++ - return { done: true, value: { type: "text", text: `response ${callCount - 1}` } } - }, - async return() { - return { done: true, value: undefined } - }, - async throw(e: any) { - throw e - }, - [Symbol.asyncDispose]: async () => {}, - } as AsyncGenerator + const mockStreamFactory = async function* (): AsyncGenerator { + yield { type: "text", text: `response ${callCount}` } + callCount++ } - const createMessageSpy = vi .spyOn(task.api, "createMessage") .mockImplementation(() => mockStreamFactory())