Q82: one output cap and a whole-response wait across provider paths
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
c3458a2eb5
commit
4e232b241f
8 files changed
+141
-41
No files matched your search
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<string, "max_completion_tokens" | "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",
|
||||
};
|
||||
|
||||
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;
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -19,13 +19,13 @@ import {
|
||||
const TURNS: Record<string, ToolTurnFn> = {
|
||||
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,
|
||||
|
||||
@@ -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<string, string>, input: ProviderToolTurnInput, requestUsage = false): Promise<AgentTurnResult> {
|
||||
/** 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<string, string>, input: ProviderToolTurnInput, requestUsage = false): Promise<AgentTurnResult> {
|
||||
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<string
|
||||
return { ...(await accumulateChatStream(response.messages)), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
|
||||
export function makeChatCompletionsTurn(endpoint: string, label: string, requestUsage = false): (input: ProviderToolTurnInput) => Promise<AgentTurnResult> {
|
||||
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<AgentTurnResult> {
|
||||
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<AgentTurnResult> {
|
||||
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<Ag
|
||||
];
|
||||
const body: JsonRecord = {
|
||||
model: input.model,
|
||||
max_tokens: 4096,
|
||||
max_tokens: OUTPUT_RESERVE_TOKENS,
|
||||
tools: anthropicToolDefs(input.tools),
|
||||
messages: toAnthropicMessages(input.messages),
|
||||
system,
|
||||
@@ -114,7 +111,7 @@ export async function runClaudeProTurn(input: ProviderToolTurnInput): Promise<Ag
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
"user-agent": "claude-cli/2.1.75",
|
||||
"x-app": "cli",
|
||||
}, body);
|
||||
}, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseAnthropicContent(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -123,7 +120,7 @@ export async function runClaudeProTurn(input: ProviderToolTurnInput): Promise<Ag
|
||||
|
||||
export async function runKimiCodeTurn(input: ProviderToolTurnInput): Promise<AgentTurnResult> {
|
||||
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<AgentT
|
||||
tool_choice: "auto",
|
||||
parallel_tool_calls: true,
|
||||
tools: responsesToolDefs(input.tools, null),
|
||||
// No output cap: the ChatGPT backend takes none, and neither Codex CLI nor Pi sends one.
|
||||
};
|
||||
const response = await postEventStream(`${CODEX_BASE_URL}/codex/responses`, {
|
||||
Authorization: `Bearer ${access}`,
|
||||
@@ -164,6 +162,7 @@ export async function runCopilotTurn(input: ProviderToolTurnInput): Promise<Agen
|
||||
tools: responsesToolDefs(input.tools),
|
||||
tool_choice: "auto",
|
||||
store: false,
|
||||
max_output_tokens: OUTPUT_RESERVE_TOKENS,
|
||||
...(input.instructions ? { instructions: input.instructions } : {}),
|
||||
};
|
||||
const response = await postJson(`${baseUrl}/responses`, {
|
||||
@@ -171,7 +170,7 @@ export async function runCopilotTurn(input: ProviderToolTurnInput): Promise<Agen
|
||||
...COPILOT_HEADERS,
|
||||
"X-Initiator": lastMessageIsUser(input.messages) ? "user" : "agent",
|
||||
"Openai-Intent": "conversation-edits",
|
||||
}, body);
|
||||
}, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseResponsesOutput(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -186,8 +185,9 @@ export async function runCustomOpenAiTurn(input: ProviderToolTurnInput): Promise
|
||||
tools: chatToolDefs(input.tools),
|
||||
tool_choice: "auto",
|
||||
stream: false,
|
||||
// No output cap: whatever serves a custom endpoint decides which cap field, if any, it accepts.
|
||||
};
|
||||
const response = await postJson(`${input.credential.baseUrl}/chat/completions`, { authorization: `Bearer ${requireApiKey(input.credential, "Custom OpenAI-compatible provider")}` }, body);
|
||||
const response = await postJson(`${input.credential.baseUrl}/chat/completions`, { authorization: `Bearer ${requireApiKey(input.credential, "Custom OpenAI-compatible provider")}` }, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseChatChoice(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -198,12 +198,12 @@ export async function runCustomAnthropicTurn(input: ProviderToolTurnInput): Prom
|
||||
if (!input.credential.baseUrl) throw new Error("Custom Anthropic-compatible provider is missing its base URL.");
|
||||
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(`${input.credential.baseUrl}/v1/messages`, { "x-api-key": requireApiKey(input.credential, "Custom Anthropic-compatible provider"), "anthropic-version": "2023-06-01" }, body);
|
||||
const response = await postJson(`${input.credential.baseUrl}/v1/messages`, { "x-api-key": requireApiKey(input.credential, "Custom Anthropic-compatible provider"), "anthropic-version": "2023-06-01" }, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseAnthropicContent(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -241,9 +241,10 @@ export async function runAzureTurn(input: ProviderToolTurnInput): Promise<AgentT
|
||||
tools: chatToolDefs(input.tools),
|
||||
tool_choice: "auto",
|
||||
stream: false,
|
||||
[chatMaxTokensField("azure-openai")]: OUTPUT_RESERVE_TOKENS,
|
||||
};
|
||||
const url = `${input.credential.baseUrl}/openai/deployments/${input.credential.deployment}/chat/completions?api-version=2024-10-21`;
|
||||
const response = await postJson(url, { "api-key": apiKey }, body);
|
||||
const response = await postJson(url, { "api-key": apiKey }, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseChatChoice(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -255,7 +256,8 @@ export async function runBedrockTurn(input: ProviderToolTurnInput): Promise<Agen
|
||||
if (!region || !accessKeyId || !secretAccessKey) throw new Error("AWS Bedrock credential is incomplete.");
|
||||
const converse: JsonRecord = {
|
||||
messages: toBedrockMessages(input.messages),
|
||||
inferenceConfig: { maxTokens: 4096 },
|
||||
// Output limits differ per Bedrock model family; only Claude's are known to cover the reserve, as in Pi.
|
||||
...(input.model.includes("anthropic.") ? { inferenceConfig: { maxTokens: OUTPUT_RESERVE_TOKENS } } : {}),
|
||||
toolConfig: bedrockToolConfig(input.tools),
|
||||
...(input.instructions ? { system: [{ text: input.instructions }] } : {}),
|
||||
};
|
||||
@@ -267,7 +269,7 @@ export async function runBedrockTurn(input: ProviderToolTurnInput): Promise<Agen
|
||||
body: bodyString,
|
||||
signer: { region, service: "bedrock", accessKeyId, secretAccessKey, sessionToken },
|
||||
});
|
||||
const response = await postJson(signed.url, signed.headers, converse);
|
||||
const response = await postJson(signed.url, signed.headers, converse, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseBedrockOutput(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -280,7 +282,7 @@ export async function runVertexTurn(input: ProviderToolTurnInput): Promise<Agent
|
||||
const token = await getGcpAccessToken(serviceAccountJson);
|
||||
const body = vertexRequestBody({ model: input.model, instructions: input.instructions, messages: input.messages, tools: input.tools });
|
||||
const url = `https://${region}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${region}/publishers/google/models/${input.model}:generateContent`;
|
||||
const response = await postJson(url, { authorization: `Bearer ${token}` }, body);
|
||||
const response = await postJson(url, { authorization: `Bearer ${token}` }, body, WHOLE_RESPONSE_WAIT);
|
||||
const parsed = parseGeminiResponse(response.data);
|
||||
return { ...parsed, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState };
|
||||
}
|
||||
@@ -25,8 +25,10 @@ describe("stream usage requests", () => {
|
||||
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 });
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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<string, unknown>; timeouts: ProviderRequestTimeouts | undefined }> = [];
|
||||
const STOP = new Error("captured");
|
||||
|
||||
vi.mock("../../src/main/services/provider-http", async (original) => ({
|
||||
...(await original<typeof import("../../src/main/services/provider-http")>()),
|
||||
postJson: async (_url: string, _headers: unknown, body: Record<string, unknown>, timeouts?: ProviderRequestTimeouts) => { sent.push({ streamed: false, body, timeouts }); throw STOP; },
|
||||
postEventStream: async (_url: string, _headers: unknown, body: Record<string, unknown>, 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<string, unknown>): unknown {
|
||||
const request = (body.request ?? body) as Record<string, unknown>;
|
||||
const generation = request.generationConfig as Record<string, unknown> | undefined;
|
||||
const inference = body.inferenceConfig as Record<string, unknown> | 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<string, string> = {
|
||||
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]);
|
||||
}
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user