From 09f10595755040963796cbd4d925a0f56a7f4b97 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 25 Jun 2026 14:10:05 +0200 Subject: [PATCH] fix(ai): clamp streamSimple max tokens Clamps streamSimple max-token defaults against estimated context, addressing #5595. closes #6061 --- packages/ai/CHANGELOG.md | 1 + packages/ai/src/api/anthropic-messages.ts | 10 +- packages/ai/src/api/azure-openai-responses.ts | 2 +- .../ai/src/api/bedrock-converse-stream.ts | 15 ++- packages/ai/src/api/google-generative-ai.ts | 2 +- packages/ai/src/api/google-vertex.ts | 2 +- packages/ai/src/api/mistral-conversations.ts | 2 +- packages/ai/src/api/openai-codex-responses.ts | 2 +- packages/ai/src/api/openai-completions.ts | 2 +- packages/ai/src/api/openai-responses.ts | 2 +- packages/ai/src/api/simple-options.ts | 29 ++++- packages/ai/src/utils/estimate.ts | 111 ++++++++++++++++++ .../openai-completions-empty-tools.test.ts | 38 +++++- 13 files changed, 198 insertions(+), 20 deletions(-) create mode 100644 packages/ai/src/utils/estimate.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 646e20af1..597bba9bd 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed +- Fixed `streamSimple()` to send a context-aware max-token cap so providers that count input and output against one context window do not reject long requests ([#5595](https://github.com/earendil-works/pi/issues/5595)). - Fixed OpenAI Responses streams to preserve reasoning replay state when output items finish out of order ([#6009](https://github.com/earendil-works/pi/issues/6009)). - Fixed retry classification for provider errors that explicitly tell callers to retry the request ([#6019](https://github.com/earendil-works/pi/issues/6019)). diff --git a/packages/ai/src/api/anthropic-messages.ts b/packages/ai/src/api/anthropic-messages.ts index 39fe4eea1..486be77c8 100644 --- a/packages/ai/src/api/anthropic-messages.ts +++ b/packages/ai/src/api/anthropic-messages.ts @@ -36,7 +36,7 @@ import { getProviderEnvValue } from "../utils/provider-env.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts"; -import { adjustMaxTokensForThinking, buildBaseOptions } from "./simple-options.ts"; +import { adjustMaxTokensForThinking, buildBaseOptions, clampMaxTokensToContext } from "./simple-options.ts"; import { transformMessages } from "./transform-messages.ts"; /** @@ -771,7 +771,7 @@ export const streamSimple: StreamFunction<"anthropic-messages", SimpleStreamOpti ): AssistantMessageEventStream => { assertRequestAuth(model.provider, options?.apiKey, options?.headers); - const base = buildBaseOptions(model, options, options?.apiKey); + const base = buildBaseOptions(model, context, options, options?.apiKey); if (!options?.reasoning) { return stream(model, context, { ...base, thinkingEnabled: false } satisfies AnthropicOptions); } @@ -796,11 +796,13 @@ export const streamSimple: StreamFunction<"anthropic-messages", SimpleStreamOpti options.thinkingBudgets, ); + const maxTokens = clampMaxTokensToContext(model, context, adjusted.maxTokens); + return stream(model, context, { ...base, - maxTokens: adjusted.maxTokens, + maxTokens, thinkingEnabled: true, - thinkingBudgetTokens: adjusted.thinkingBudget, + thinkingBudgetTokens: Math.min(adjusted.thinkingBudget, Math.max(0, maxTokens - 1024)), } satisfies AnthropicOptions); }; diff --git a/packages/ai/src/api/azure-openai-responses.ts b/packages/ai/src/api/azure-openai-responses.ts index 137c43c91..a5438c194 100644 --- a/packages/ai/src/api/azure-openai-responses.ts +++ b/packages/ai/src/api/azure-openai-responses.ts @@ -160,7 +160,7 @@ export const streamSimple: StreamFunction<"azure-openai-responses", SimpleStream throw new Error(`No API key for provider: ${model.provider}`); } - const base = buildBaseOptions(model, options, apiKey); + const base = buildBaseOptions(model, context, options, apiKey); const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning; diff --git a/packages/ai/src/api/bedrock-converse-stream.ts b/packages/ai/src/api/bedrock-converse-stream.ts index 68af206de..fd20828e9 100644 --- a/packages/ai/src/api/bedrock-converse-stream.ts +++ b/packages/ai/src/api/bedrock-converse-stream.ts @@ -53,7 +53,12 @@ import { parseStreamingJson } from "../utils/json-parse.ts"; import { resolveHttpProxyUrlForTarget } from "../utils/node-http-proxy.ts"; import { getProviderEnvValue } from "../utils/provider-env.ts"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts"; -import { adjustMaxTokensForThinking, buildBaseOptions, clampReasoning } from "./simple-options.ts"; +import { + adjustMaxTokensForThinking, + buildBaseOptions, + clampMaxTokensToContext, + clampReasoning, +} from "./simple-options.ts"; import { transformMessages } from "./transform-messages.ts"; export type BedrockThinkingDisplay = "summarized" | "omitted"; @@ -374,7 +379,7 @@ export const streamSimple: StreamFunction<"bedrock-converse-stream", SimpleStrea context: Context, options?: SimpleStreamOptions, ): AssistantMessageEventStream => { - const base = buildBaseOptions(model, options, undefined); + const base = buildBaseOptions(model, context, options, undefined); if (!options?.reasoning) { return stream(model, context, { ...base, reasoning: undefined } satisfies BedrockOptions); } @@ -397,13 +402,15 @@ export const streamSimple: StreamFunction<"bedrock-converse-stream", SimpleStrea options.thinkingBudgets, ); + const maxTokens = clampMaxTokensToContext(model, context, adjusted.maxTokens); + return stream(model, context, { ...base, - maxTokens: adjusted.maxTokens, + maxTokens, reasoning: options.reasoning, thinkingBudgets: { ...(options.thinkingBudgets || {}), - [clampReasoning(options.reasoning)!]: adjusted.thinkingBudget, + [clampReasoning(options.reasoning)!]: Math.min(adjusted.thinkingBudget, Math.max(0, maxTokens - 1024)), }, } satisfies BedrockOptions); } diff --git a/packages/ai/src/api/google-generative-ai.ts b/packages/ai/src/api/google-generative-ai.ts index 1e90e53ae..623fc3ea3 100644 --- a/packages/ai/src/api/google-generative-ai.ts +++ b/packages/ai/src/api/google-generative-ai.ts @@ -290,7 +290,7 @@ export const streamSimple: StreamFunction<"google-generative-ai", SimpleStreamOp throw new Error(`No API key for provider: ${model.provider}`); } - const base = buildBaseOptions(model, options, apiKey); + const base = buildBaseOptions(model, context, options, apiKey); if (!options?.reasoning) { return stream(model, context, { ...base, thinking: { enabled: false } } satisfies GoogleOptions); } diff --git a/packages/ai/src/api/google-vertex.ts b/packages/ai/src/api/google-vertex.ts index f4ea7b5c0..4c7ab9806 100644 --- a/packages/ai/src/api/google-vertex.ts +++ b/packages/ai/src/api/google-vertex.ts @@ -302,7 +302,7 @@ export const streamSimple: StreamFunction<"google-vertex", SimpleStreamOptions> context: Context, options?: SimpleStreamOptions, ): AssistantMessageEventStream => { - const base = buildBaseOptions(model, options, undefined); + const base = buildBaseOptions(model, context, options, undefined); if (!options?.reasoning) { return stream(model, context, { ...base, diff --git a/packages/ai/src/api/mistral-conversations.ts b/packages/ai/src/api/mistral-conversations.ts index 66519e5ec..f1e4fc4c4 100644 --- a/packages/ai/src/api/mistral-conversations.ts +++ b/packages/ai/src/api/mistral-conversations.ts @@ -117,7 +117,7 @@ export const streamSimple: StreamFunction<"mistral-conversations", SimpleStreamO throw new Error(`No API key for provider: ${model.provider}`); } - const base = buildBaseOptions(model, options, apiKey); + const base = buildBaseOptions(model, context, options, apiKey); const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; const reasoning = clampedReasoning === "off" ? undefined : clampedReasoning; const shouldUseReasoning = model.reasoning && reasoning !== undefined; diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index 5db519a7e..267fa90b1 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -430,7 +430,7 @@ export const streamSimple: StreamFunction<"openai-codex-responses", SimpleStream throw new Error(`No API key for provider: ${model.provider}`); } - const base = buildBaseOptions(model, options, apiKey); + const base = buildBaseOptions(model, context, options, apiKey); const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning; diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 27f44b2cf..caeb8638e 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -482,7 +482,7 @@ export const streamSimple: StreamFunction<"openai-completions", SimpleStreamOpti ): AssistantMessageEventStream => { getClientApiKey(model.provider, options?.apiKey, options?.headers); - const base = buildBaseOptions(model, options, options?.apiKey); + const base = buildBaseOptions(model, context, options, options?.apiKey); const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning; const toolChoice = (options as OpenAICompletionsOptions | undefined)?.toolChoice; diff --git a/packages/ai/src/api/openai-responses.ts b/packages/ai/src/api/openai-responses.ts index 40998eaa3..52868c655 100644 --- a/packages/ai/src/api/openai-responses.ts +++ b/packages/ai/src/api/openai-responses.ts @@ -180,7 +180,7 @@ export const streamSimple: StreamFunction<"openai-responses", SimpleStreamOption ): AssistantMessageEventStream => { getClientApiKey(model.provider, options?.apiKey, options?.headers); - const base = buildBaseOptions(model, options, options?.apiKey); + const base = buildBaseOptions(model, context, options, options?.apiKey); const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning; diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index 773e8d95a..f1bb29fd2 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -1,9 +1,32 @@ -import type { Api, Model, SimpleStreamOptions, StreamOptions, ThinkingBudgets, ThinkingLevel } from "../types.ts"; +import type { + Api, + Context, + Model, + SimpleStreamOptions, + StreamOptions, + ThinkingBudgets, + ThinkingLevel, +} from "../types.ts"; +import { estimateContextTokens } from "../utils/estimate.ts"; -export function buildBaseOptions(_model: Model, options?: SimpleStreamOptions, apiKey?: string): StreamOptions { +const CONTEXT_SAFETY_TOKENS = 4096; +const MIN_MAX_TOKENS = 1; + +export function clampMaxTokensToContext(model: Model, context: Context, maxTokens: number): number { + if (model.contextWindow <= 0) return Math.max(MIN_MAX_TOKENS, maxTokens); + const available = model.contextWindow - estimateContextTokens(context).tokens - CONTEXT_SAFETY_TOKENS; + return Math.min(maxTokens, Math.max(MIN_MAX_TOKENS, available)); +} + +export function buildBaseOptions( + model: Model, + context: Context, + options?: SimpleStreamOptions, + apiKey?: string, +): StreamOptions { return { temperature: options?.temperature, - maxTokens: options?.maxTokens, + maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens), signal: options?.signal, apiKey: apiKey || options?.apiKey, transport: options?.transport, diff --git a/packages/ai/src/utils/estimate.ts b/packages/ai/src/utils/estimate.ts new file mode 100644 index 000000000..711bf7830 --- /dev/null +++ b/packages/ai/src/utils/estimate.ts @@ -0,0 +1,111 @@ +import type { AssistantMessage, Context, ImageContent, Message, TextContent, Usage } from "../types.ts"; + +export interface ContextUsageEstimate { + /** Estimated total context tokens. */ + tokens: number; + /** Tokens reported by the most recent assistant usage block. */ + usageTokens: number; + /** Estimated tokens after the most recent assistant usage block. */ + trailingTokens: number; + /** Index of the message that provided usage, or null when none exists. */ + lastUsageIndex: number | null; +} + +const CHARS_PER_TOKEN = 4; +const ESTIMATED_IMAGE_CHARS = 4800; + +export function calculateContextTokens(usage: Usage): number { + return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; +} + +function safeJsonStringify(value: unknown): string { + try { + return JSON.stringify(value) ?? "undefined"; + } catch { + return "[unserializable]"; + } +} + +function estimateTextAndImageContentChars(content: string | Array): number { + if (typeof content === "string") return content.length; + + let chars = 0; + for (const block of content) chars += block.type === "text" ? block.text.length : ESTIMATED_IMAGE_CHARS; + return chars; +} + +export function estimateTextTokens(text: string): number { + return Math.ceil(text.length / CHARS_PER_TOKEN); +} + +export function estimateTextAndImageContentTokens(content: string | Array): number { + return Math.ceil(estimateTextAndImageContentChars(content) / CHARS_PER_TOKEN); +} + +export function estimateMessageTokens(message: Message): number { + let chars = 0; + + if (message.role === "user") return estimateTextAndImageContentTokens(message.content); + if (message.role === "toolResult") return estimateTextAndImageContentTokens(message.content); + + for (const block of message.content) { + if (block.type === "text") { + chars += block.text.length; + } else if (block.type === "thinking") { + chars += block.thinking.length; + } else { + chars += block.name.length + safeJsonStringify(block.arguments).length; + } + } + return Math.ceil(chars / CHARS_PER_TOKEN); +} + +function getLastAssistantUsageInfo(messages: readonly Message[]): { usage: Usage; index: number } | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + if (message.role !== "assistant") continue; + const assistant = message as AssistantMessage; + if (assistant.stopReason === "aborted" || assistant.stopReason === "error") continue; + if (calculateContextTokens(assistant.usage) > 0) return { usage: assistant.usage, index: i }; + } + return undefined; +} + +function estimateMessages(messages: readonly Message[]): ContextUsageEstimate { + const usageInfo = getLastAssistantUsageInfo(messages); + if (usageInfo) { + const usageTokens = calculateContextTokens(usageInfo.usage); + let trailingTokens = 0; + for (let i = usageInfo.index + 1; i < messages.length; i++) { + trailingTokens += estimateMessageTokens(messages[i]); + } + return { tokens: usageTokens + trailingTokens, usageTokens, trailingTokens, lastUsageIndex: usageInfo.index }; + } + + let tokens = 0; + for (const message of messages) tokens += estimateMessageTokens(message); + return { tokens, usageTokens: 0, trailingTokens: tokens, lastUsageIndex: null }; +} + +function isMessageArray(value: Context | readonly Message[]): value is readonly Message[] { + return Array.isArray(value); +} + +export function estimateContextTokens(context: Context | readonly Message[]): ContextUsageEstimate { + if (isMessageArray(context)) return estimateMessages(context); + + const estimate = estimateMessages(context.messages); + if (estimate.lastUsageIndex !== null) return estimate; + + let prefixTokens = context.systemPrompt ? estimateTextTokens(context.systemPrompt) : 0; + if (context.tools && context.tools.length > 0) { + prefixTokens += estimateTextTokens(safeJsonStringify(context.tools)); + } + + return { + tokens: estimate.tokens + prefixTokens, + usageTokens: estimate.usageTokens, + trailingTokens: estimate.trailingTokens + prefixTokens, + lastUsageIndex: estimate.lastUsageIndex, + }; +} diff --git a/packages/ai/test/openai-completions-empty-tools.test.ts b/packages/ai/test/openai-completions-empty-tools.test.ts index e83f63ae7..9fd9fc5f8 100644 --- a/packages/ai/test/openai-completions-empty-tools.test.ts +++ b/packages/ai/test/openai-completions-empty-tools.test.ts @@ -92,7 +92,7 @@ describe("openai-completions empty tools handling", () => { expect("tools" in (params as object)).toBe(false); }); - it("does not send default max token fields", async () => { + it("sends default maxTokens", async () => { const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!; const model = { ...baseModel, api: "openai-completions" } as const; @@ -106,7 +106,7 @@ describe("openai-completions empty tools handling", () => { const params = mockState.lastParams as { max_tokens?: number; max_completion_tokens?: number }; expect(params.max_tokens).toBeUndefined(); - expect(params.max_completion_tokens).toBeUndefined(); + expect(params.max_completion_tokens).toBe(model.maxTokens); }); it("sends explicit maxTokens", async () => { @@ -126,6 +126,40 @@ describe("openai-completions empty tools handling", () => { expect(params.max_completion_tokens).toBe(1234); }); + it("clamps default maxTokens to remaining context", async () => { + const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!; + const model = { ...baseModel, api: "openai-completions", contextWindow: 10000, maxTokens: 8000 } as const; + + await streamSimple( + model, + { + messages: [{ role: "user", content: "x".repeat(8000), timestamp: Date.now() }], + }, + { apiKey: "test" }, + ).result(); + + const params = mockState.lastParams as { max_tokens?: number; max_completion_tokens?: number }; + expect(params.max_tokens).toBeUndefined(); + expect(params.max_completion_tokens).toBe(3904); + }); + + it("clamps explicit maxTokens to remaining context", async () => { + const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!; + const model = { ...baseModel, api: "openai-completions", contextWindow: 10000, maxTokens: 8000 } as const; + + await streamSimple( + model, + { + messages: [{ role: "user", content: "x".repeat(8000), timestamp: Date.now() }], + }, + { apiKey: "test", maxTokens: 7000 }, + ).result(); + + const params = mockState.lastParams as { max_tokens?: number; max_completion_tokens?: number }; + expect(params.max_tokens).toBeUndefined(); + expect(params.max_completion_tokens).toBe(3904); + }); + it("uses conservative OpenAI-compatible fields for Cloudflare AI Gateway /compat models", async () => { process.env.CLOUDFLARE_API_KEY = "cf-token"; process.env.CLOUDFLARE_ACCOUNT_ID = "account-id";