diff --git a/src/main/services/provider-adapters.ts b/src/main/services/provider-adapters.ts index 394ed2c4..4ee0cb17 100644 --- a/src/main/services/provider-adapters.ts +++ b/src/main/services/provider-adapters.ts @@ -22,7 +22,9 @@ import { } from "./provider-auth-profiles"; import { backgroundUnsupported, retrieveUnsupported } from "./provider-background"; import { authorizeWithParser, refreshUnsupported, revokeLocalProfile } from "./provider-auth-lifecycle"; -import { getJson, postEventStream, postJson, type JsonRecord } from "./provider-http"; +import { getJson, postEventStream, postJson, WHOLE_RESPONSE_WAIT, type JsonRecord } from "./provider-http"; +import { OUTPUT_RESERVE_TOKENS } from "./context-budget"; +import { chatMaxTokensField } from "./provider-tool-turns/chat-format"; import { fetchAnthropicModels, fetchOpenAiModels } from "./provider-models-fetch"; import { KIMI_CODE_BASE_URL } from "./oauth/kimi-code-flow"; import { FINISHED_STOP_REASONS, incompleteTurn, requireFinished } from "./provider-tool-turn"; @@ -70,6 +72,7 @@ export const providerAdapters: ProviderAdapterRegistry = { model: input.model, input: input.prompt, stream: true, + max_output_tokens: OUTPUT_RESERVE_TOKENS, ...(input.instructions ? { instructions: input.instructions } : {}), }); return { events: mapOpenAiResponseStream(response.messages), rateLimitState: response.rateLimitState }; @@ -91,6 +94,7 @@ export const providerAdapters: ProviderAdapterRegistry = { input: input.prompt, background: true, store: true, + max_output_tokens: OUTPUT_RESERVE_TOKENS, ...(input.instructions ? { instructions: input.instructions } : {}), }); const data = response.data; @@ -123,8 +127,9 @@ export const providerAdapters: ProviderAdapterRegistry = { }, { model: input.model, input: input.prompt, + max_output_tokens: OUTPUT_RESERVE_TOKENS, ...(input.instructions ? { instructions: input.instructions } : {}), - }); + }, WHOLE_RESPONSE_WAIT); const data = response.data; if (data.status !== "completed") throw incompleteTurn(incompleteReason(data)); const text = requireText(collectText(data.output_text ?? data.output), "OpenAI"); @@ -158,7 +163,7 @@ export const providerAdapters: ProviderAdapterRegistry = { "anthropic-version": "2023-06-01", }, { model: input.model, - max_tokens: 4096, + max_tokens: OUTPUT_RESERVE_TOKENS, stream: true, ...(input.instructions ? { system: input.instructions } : {}), messages: [{ role: "user", content: input.prompt }], @@ -175,10 +180,10 @@ export const providerAdapters: ProviderAdapterRegistry = { "anthropic-version": "2023-06-01", }, { model: input.model, - max_tokens: 4096, + max_tokens: OUTPUT_RESERVE_TOKENS, ...(input.instructions ? { system: input.instructions } : {}), messages: [{ role: "user", content: input.prompt }], - }); + }, WHOLE_RESPONSE_WAIT); const data = response.data; requireFinished(data.stop_reason, FINISHED_STOP_REASONS.anthropic); const text = requireText(collectText(data.content), "Anthropic"); @@ -403,6 +408,7 @@ function openAiCompatibleAdapter({ model: input.model, messages, stream: true, + [chatMaxTokensField(providerKey)]: OUTPUT_RESERVE_TOKENS, ...(requestStreamUsage ? { stream_options: { include_usage: true } } : {}), }); return { events: mapOpenAiChatCompletionStream(response.messages), rateLimitState: response.rateLimitState }; @@ -418,7 +424,8 @@ function openAiCompatibleAdapter({ }, { model: input.model, messages, - }); + [chatMaxTokensField(providerKey)]: OUTPUT_RESERVE_TOKENS, + }, WHOLE_RESPONSE_WAIT); const data = response.data; requireFinished(Array.isArray(data.choices) && isRecord(data.choices[0]) ? data.choices[0].finish_reason : null, FINISHED_STOP_REASONS.chat); const text = requireText(collectText(Array.isArray(data.choices) ? data.choices.filter(isRecord).map((choice) => choice.message) : []), label); diff --git a/src/main/services/provider-http.ts b/src/main/services/provider-http.ts index 5eac66ea..338e07ea 100644 --- a/src/main/services/provider-http.ts +++ b/src/main/services/provider-http.ts @@ -24,6 +24,9 @@ const HEADERS_TIMEOUT_MS = 60_000; const BODY_TIMEOUT_MS = 300_000; const STREAM_IDLE_TIMEOUT_MS = 60_000; +/** A non-streamed request sends headers only after the whole answer; wait as long as Node's fetch itself allows (five minutes). */ +export const WHOLE_RESPONSE_WAIT: ProviderRequestTimeouts = { headersTimeoutMs: 300_000 }; + export interface ProviderRequestTimeouts { headersTimeoutMs?: number; bodyTimeoutMs?: number; diff --git a/src/main/services/provider-tool-turns/chat-format.ts b/src/main/services/provider-tool-turns/chat-format.ts index dd5f1e90..2b1c98c4 100644 --- a/src/main/services/provider-tool-turns/chat-format.ts +++ b/src/main/services/provider-tool-turns/chat-format.ts @@ -2,6 +2,18 @@ import { isRecord, ProviderHttpError, type JsonRecord, type ProviderEventStreamMessage } from "../provider-http"; import { FINISHED_STOP_REASONS, replayIn, requireFinished, type AgentTurnMessage, type AgentTurnResult, type AssistantReplay, toolCallFromJson, type ToolCall, type ToolDefinition } from "../provider-tool-turn"; +/** How each Chat Completions provider names the output cap: OpenAI's current field, or the older one the others document. */ +const MAX_TOKENS_FIELD: Record = { + xai: "max_completion_tokens", groq: "max_completion_tokens", cerebras: "max_completion_tokens", "azure-openai": "max_completion_tokens", + mistral: "max_tokens", openrouter: "max_tokens", kimi: "max_tokens", "kimi-cn": "max_tokens", "kimi-code": "max_tokens", +}; + +export function chatMaxTokensField(providerKey: string): "max_completion_tokens" | "max_tokens" { + const field = MAX_TOKENS_FIELD[providerKey]; + if (!field) throw new Error(`No output-cap field is known for the ${providerKey} Chat Completions API.`); + return field; +} + /** Fields providers use for reasoning text; the first non-empty one is the provider's, as in Pi. */ const REASONING_FIELDS = ["reasoning_content", "reasoning", "reasoning_text"] as const; diff --git a/src/main/services/provider-tool-turns/gemini-format.ts b/src/main/services/provider-tool-turns/gemini-format.ts index 71da8146..6633991c 100644 --- a/src/main/services/provider-tool-turns/gemini-format.ts +++ b/src/main/services/provider-tool-turns/gemini-format.ts @@ -3,6 +3,7 @@ * Ported from Pi's google-shared `convertMessages` / `convertTools` and the * Antigravity request wrapper in google-gemini-cli. */ +import { OUTPUT_RESERVE_TOKENS } from "../context-budget"; import { isRecord, ProviderHttpError, type JsonRecord, type ProviderEventStreamMessage } from "../provider-http"; import { FINISHED_STOP_REASONS, replayIn, requireFinished, toolCallFromObject, type AgentTurnMessage, type AgentTurnResult, type AssistantReplay, type ToolCall, type ToolDefinition } from "../provider-tool-turn"; @@ -74,6 +75,7 @@ export function antigravityRequestBody(input: { }): JsonRecord { const request: JsonRecord = { contents: toGeminiContents(input.model, input.messages), + generationConfig: { maxOutputTokens: OUTPUT_RESERVE_TOKENS }, systemInstruction: { role: "user", parts: [ @@ -97,7 +99,7 @@ export function vertexRequestBody(input: { messages: AgentTurnMessage[]; tools: ToolDefinition[]; }): JsonRecord { - const body: JsonRecord = { contents: toGeminiContents(input.model, input.messages) }; + const body: JsonRecord = { contents: toGeminiContents(input.model, input.messages), generationConfig: { maxOutputTokens: OUTPUT_RESERVE_TOKENS } }; if (input.instructions) body.systemInstruction = { parts: [{ text: input.instructions }] }; if (input.tools.length > 0) { body.tools = geminiToolDefs(input.model, input.tools); diff --git a/src/main/services/provider-tool-turns/router.ts b/src/main/services/provider-tool-turns/router.ts index ee83c90a..40400f2f 100644 --- a/src/main/services/provider-tool-turns/router.ts +++ b/src/main/services/provider-tool-turns/router.ts @@ -19,13 +19,13 @@ import { const TURNS: Record = { openai: runOpenAiResponsesTurn, anthropic: runAnthropicTurn, - xai: makeChatCompletionsTurn("https://api.x.ai/v1/chat/completions", "xAI", true), - groq: makeChatCompletionsTurn("https://api.groq.com/openai/v1/chat/completions", "Groq", true), - mistral: makeChatCompletionsTurn("https://api.mistral.ai/v1/chat/completions", "Mistral"), - cerebras: makeChatCompletionsTurn("https://api.cerebras.ai/v1/chat/completions", "Cerebras"), - openrouter: makeChatCompletionsTurn("https://openrouter.ai/api/v1/chat/completions", "OpenRouter"), - kimi: makeChatCompletionsTurn("https://api.moonshot.ai/v1/chat/completions", "Kimi"), - "kimi-cn": makeChatCompletionsTurn("https://api.moonshot.cn/v1/chat/completions", "Kimi CN"), + xai: makeChatCompletionsTurn("xai", "https://api.x.ai/v1/chat/completions", "xAI", true), + groq: makeChatCompletionsTurn("groq", "https://api.groq.com/openai/v1/chat/completions", "Groq", true), + mistral: makeChatCompletionsTurn("mistral", "https://api.mistral.ai/v1/chat/completions", "Mistral"), + cerebras: makeChatCompletionsTurn("cerebras", "https://api.cerebras.ai/v1/chat/completions", "Cerebras"), + openrouter: makeChatCompletionsTurn("openrouter", "https://openrouter.ai/api/v1/chat/completions", "OpenRouter"), + kimi: makeChatCompletionsTurn("kimi", "https://api.moonshot.ai/v1/chat/completions", "Kimi"), + "kimi-cn": makeChatCompletionsTurn("kimi-cn", "https://api.moonshot.cn/v1/chat/completions", "Kimi CN"), "gpt-account": runCodexTurn, "github-copilot": runCopilotTurn, "google-antigravity": runAntigravityTurn, diff --git a/src/main/services/provider-tool-turns/turns.ts b/src/main/services/provider-tool-turns/turns.ts index 3dcc9039..39cc5339 100644 --- a/src/main/services/provider-tool-turns/turns.ts +++ b/src/main/services/provider-tool-turns/turns.ts @@ -4,17 +4,18 @@ * {@link AgentTurnResult}. Faithful ports of Pi's provider request handlers. */ import { release as osRelease, arch as osArch, platform as osPlatform } from "node:os"; -import { postEventStream, postJson, type JsonRecord } from "../provider-http"; +import { postEventStream, postJson, WHOLE_RESPONSE_WAIT, type JsonRecord } from "../provider-http"; import type { AgentTurnMessage, AgentTurnResult, ProviderToolTurnInput, TurnCredential } from "../provider-tool-turn"; import { accumulateResponsesStream, parseResponsesOutput, responsesToolDefs, toResponsesInput } from "./responses-format"; import { anthropicToolDefs, parseAnthropicContent, toAnthropicMessages } from "./anthropic-format"; -import { accumulateChatStream, chatToolDefs, parseChatChoice, toChatMessages } from "./chat-format"; +import { accumulateChatStream, chatMaxTokensField, chatToolDefs, parseChatChoice, toChatMessages } from "./chat-format"; import { accumulateGeminiStream, antigravityRequestBody, parseGeminiResponse, vertexRequestBody } from "./gemini-format"; import { bedrockToolConfig, parseBedrockOutput, toBedrockMessages } from "./bedrock-format"; import { signAwsRequest } from "./aws-sigv4"; import { getGcpAccessToken } from "./gcp-auth"; import { COPILOT_HEADERS, copilotBaseUrl } from "../oauth/github-copilot-flow"; import { KIMI_CODE_BASE_URL } from "../oauth/kimi-code-flow"; +import { OUTPUT_RESERVE_TOKENS } from "../context-budget"; const ANTIGRAVITY_VERSION = "1.21.9"; const ANTIGRAVITY_ENDPOINT = "https://cloudcode-pa.googleapis.com"; @@ -44,29 +45,25 @@ export async function runOpenAiResponsesTurn(input: ProviderToolTurnInput): Prom tools: responsesToolDefs(input.tools), tool_choice: "auto", store: false, + max_output_tokens: OUTPUT_RESERVE_TOKENS, ...(input.instructions ? { instructions: input.instructions } : {}), }; - const response = await postJson("https://api.openai.com/v1/responses", { authorization: `Bearer ${requireApiKey(input.credential, "OpenAI")}` }, body); + const response = await postJson("https://api.openai.com/v1/responses", { authorization: `Bearer ${requireApiKey(input.credential, "OpenAI")}` }, body, WHOLE_RESPONSE_WAIT); const parsed = parseResponsesOutput(response.data); return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState }; } // ---- OpenAI Chat Completions, streamed (xAI, Groq, Mistral, Cerebras, OpenRouter, Kimi, Kimi CN) ---- -/** - * Streaming is required, not an optimization: a non-streaming completion - * sends its response headers only after the WHOLE generation finishes, so a - * reasoning model that thinks past the 60s header deadline gets aborted with - * "Provider request timed out." even though it is working. Streamed turns - * emit chunks continuously and only a truly dead connection times out. - */ -async function runChatCompletionsStream(endpoint: string, headers: Record, input: ProviderToolTurnInput, requestUsage = false): Promise { +/** Streamed so a long think never meets a header deadline; only a stream that goes idle times out. */ +async function runChatCompletionsStream(providerKey: string, endpoint: string, headers: Record, input: ProviderToolTurnInput, requestUsage = false): Promise { const body: JsonRecord = { model: input.model, messages: toChatMessages(input.instructions, input.messages), tools: chatToolDefs(input.tools), tool_choice: "auto", stream: true, + [chatMaxTokensField(providerKey)]: OUTPUT_RESERVE_TOKENS, // Some providers only report a stream's token usage when asked; the others send it unasked and may reject the field. ...(requestUsage ? { stream_options: { include_usage: true } } : {}), }; @@ -74,8 +71,8 @@ async function runChatCompletionsStream(endpoint: string, headers: Record Promise { - return async (input) => runChatCompletionsStream(endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, input, requestUsage); +export function makeChatCompletionsTurn(providerKey: string, endpoint: string, label: string, requestUsage = false): (input: ProviderToolTurnInput) => Promise { + return async (input) => runChatCompletionsStream(providerKey, endpoint, { authorization: `Bearer ${requireApiKey(input.credential, label)}` }, input, requestUsage); } // ---- Anthropic Messages (api key) ---- @@ -83,12 +80,12 @@ export function makeChatCompletionsTurn(endpoint: string, label: string, request export async function runAnthropicTurn(input: ProviderToolTurnInput): Promise { const body: JsonRecord = { model: input.model, - max_tokens: 4096, + max_tokens: OUTPUT_RESERVE_TOKENS, tools: anthropicToolDefs(input.tools), messages: toAnthropicMessages(input.messages), ...(input.instructions ? { system: input.instructions } : {}), }; - const response = await postJson("https://api.anthropic.com/v1/messages", { "x-api-key": requireApiKey(input.credential, "Anthropic"), "anthropic-version": "2023-06-01" }, body); + const response = await postJson("https://api.anthropic.com/v1/messages", { "x-api-key": requireApiKey(input.credential, "Anthropic"), "anthropic-version": "2023-06-01" }, body, WHOLE_RESPONSE_WAIT); const parsed = parseAnthropicContent(response.data); return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState }; } @@ -102,7 +99,7 @@ export async function runClaudeProTurn(input: ProviderToolTurnInput): Promise { const access = requireAccessToken(input.credential, "Kimi Code"); - return runChatCompletionsStream(`${KIMI_CODE_BASE_URL}/chat/completions`, { authorization: `Bearer ${access}` }, input); + return runChatCompletionsStream("kimi-code", `${KIMI_CODE_BASE_URL}/chat/completions`, { authorization: `Bearer ${access}` }, input); } // ---- GPT account / OpenAI Codex (OAuth, SSE) ---- @@ -142,6 +139,7 @@ export async function runCodexTurn(input: ProviderToolTurnInput): Promise { for await (const event of result.events) void event; } expect(bodies.map((body) => body.stream_options ?? null)).toEqual([{ include_usage: true }, { include_usage: true }, null, null]); + expect(bodies.map((body) => body.max_completion_tokens ?? body.max_tokens)).toEqual([16_000, 16_000, 16_000, 16_000]); + expect(bodies.map((body) => "max_completion_tokens" in body)).toEqual([true, true, false, true]); - await makeChatCompletionsTurn("https://api.x.ai/v1/chat/completions", "xAI", true)({ credential: { apiKey: "k" }, model: "m", instructions: null, messages: [{ role: "user", text: "hi" }], tools: [] }); + await makeChatCompletionsTurn("xai", "https://api.x.ai/v1/chat/completions", "xAI", true)({ credential: { apiKey: "k" }, model: "m", instructions: null, messages: [{ role: "user", text: "hi" }], tools: [] }); expect(bodies.at(-1)?.stream_options).toEqual({ include_usage: true }); }); }); diff --git a/tests/unit/turn-request-bounds.test.ts b/tests/unit/turn-request-bounds.test.ts new file mode 100644 index 00000000..8ba52214 --- /dev/null +++ b/tests/unit/turn-request-bounds.test.ts @@ -0,0 +1,72 @@ +import { beforeEach, describe, expect, test, vi } from "vitest"; +import type { ProviderRequestTimeouts } from "../../src/main/services/provider-http"; + +const sent: Array<{ streamed: boolean; body: Record; timeouts: ProviderRequestTimeouts | undefined }> = []; +const STOP = new Error("captured"); + +vi.mock("../../src/main/services/provider-http", async (original) => ({ + ...(await original()), + postJson: async (_url: string, _headers: unknown, body: Record, timeouts?: ProviderRequestTimeouts) => { sent.push({ streamed: false, body, timeouts }); throw STOP; }, + postEventStream: async (_url: string, _headers: unknown, body: Record, timeouts?: ProviderRequestTimeouts) => { sent.push({ streamed: true, body, timeouts }); throw STOP; }, +})); +vi.mock("../../src/main/services/provider-tool-turns/gcp-auth", () => ({ getGcpAccessToken: async () => "token" })); + +const { selectToolTurn } = await import("../../src/main/services/provider-tool-turns/router"); +const { OUTPUT_RESERVE_TOKENS } = await import("../../src/main/services/context-budget"); + +const credential = { + apiKey: "k", accessToken: "tid=1;exp=9999999999", accountId: "acct", projectId: "proj", baseUrl: "https://example.test", + deployment: "dep", region: "us-east-1", accessKeyId: "AKIA", secretAccessKey: "secret", serviceAccountJson: "{}", +}; + +async function capture(providerKey: string, model = "m"): Promise<(typeof sent)[number]> { + await expect(selectToolTurn(providerKey)({ credential, model, instructions: null, messages: [{ role: "user", text: "hi" }], tools: [] })).rejects.toBe(STOP); + const request = sent.at(-1); + if (!request) throw new Error(`${providerKey} sent nothing`); + return request; +} + +/** The output cap a request carries, whatever the API calls it. */ +function outputCap(body: Record): unknown { + const request = (body.request ?? body) as Record; + const generation = request.generationConfig as Record | undefined; + const inference = body.inferenceConfig as Record | undefined; + return body.max_output_tokens ?? body.max_tokens ?? body.max_completion_tokens ?? generation?.maxOutputTokens ?? inference?.maxTokens ?? null; +} + +beforeEach(() => { sent.length = 0; }); + +describe("every agent turn has the same output cap and a wait that fits how it is sent", () => { + test("each provider caps output at the reserve the context budget keeps, in the field its API names", async () => { + const fields: Record = { + openai: "max_output_tokens", "github-copilot": "max_output_tokens", + anthropic: "max_tokens", "claude-pro": "max_tokens", "custom-anthropic": "max_tokens", + xai: "max_completion_tokens", groq: "max_completion_tokens", cerebras: "max_completion_tokens", "azure-openai": "max_completion_tokens", + mistral: "max_tokens", openrouter: "max_tokens", kimi: "max_tokens", "kimi-cn": "max_tokens", "kimi-code": "max_tokens", + }; + for (const [key, field] of Object.entries(fields)) { + const { body } = await capture(key); + expect([key, body[field]]).toEqual([key, OUTPUT_RESERVE_TOKENS]); + } + expect(outputCap((await capture("google-antigravity", "gemini-3-pro")).body)).toBe(OUTPUT_RESERVE_TOKENS); + expect(outputCap((await capture("vertex-ai", "gemini-3-pro")).body)).toBe(OUTPUT_RESERVE_TOKENS); + expect(outputCap((await capture("aws-bedrock", "anthropic.claude-sonnet-4-5")).body)).toBe(OUTPUT_RESERVE_TOKENS); + }); + + test("Codex, custom OpenAI endpoints and non-Claude Bedrock models get no cap: none taken, field unknown, or limits that differ per model", async () => { + expect(outputCap((await capture("gpt-account")).body)).toBeNull(); + expect(outputCap((await capture("custom-openai")).body)).toBeNull(); + expect(outputCap((await capture("aws-bedrock", "meta.llama3-70b-instruct-v1:0")).body)).toBeNull(); + }); + + test("turns that wait for a whole response allow a long think before headers; streamed turns keep the idle deadline", async () => { + for (const key of ["openai", "anthropic", "claude-pro", "github-copilot", "custom-openai", "custom-anthropic", "azure-openai", "aws-bedrock", "vertex-ai"]) { + const { streamed, timeouts } = await capture(key); + expect([key, streamed, timeouts?.headersTimeoutMs]).toEqual([key, false, 300_000]); + } + for (const key of ["xai", "kimi-code", "gpt-account", "google-antigravity"]) { + const { streamed, timeouts } = await capture(key); + expect([key, streamed, timeouts?.headersTimeoutMs]).toEqual([key, true, undefined]); + } + }); +});