From cb594ae3704bcc687a62e4c48a0d2e23fbbad852 Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Sat, 18 Jul 2026 20:13:01 +0200 Subject: [PATCH] Reject invalid output token budgets (#7379) * fix(context): reject invalid output token budgets * fix(context): enforce output budgets across request formats * fix(context): enforce default Claude output budgets * fix(context): include Responses API input in token budget * ci: rerun pull request checks --- open-sse/handlers/chatCore.ts | 52 ++++++++ .../handlers/chatCore/outputTokenBudget.ts | 117 ++++++++++++++++++ tests/unit/chatcore-sanitization.test.ts | 2 +- tests/unit/output-token-budget.test.ts | 66 ++++++++++ 4 files changed, 236 insertions(+), 1 deletion(-) create mode 100644 open-sse/handlers/chatCore/outputTokenBudget.ts create mode 100644 tests/unit/output-token-budget.test.ts diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index bf6aa9c6a8..abcd3a9c0b 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -14,6 +14,7 @@ import { buildPostCallGuardrailContext } from "./chatCore/postCallGuardrailConte import { storeSemanticCacheResponse } from "./chatCore/semanticCacheStore.ts"; import { buildNonStreamingResponseHeaders } from "./chatCore/nonStreamingResponseHeaders.ts"; import { buildNonStreamingJsonResponse } from "./chatCore/nonStreamingJsonResponse.ts"; +import { enforceOutputTokenBudget } from "./chatCore/outputTokenBudget.ts"; import { maybeConvertJsonBodyToSse } from "./chatCore/jsonBodyToSse.ts"; import { assembleStreamingResponseHeaders } from "./chatCore/streamingResponseHeaders.ts"; import { storeStreamingSemanticCacheResponse } from "./chatCore/streamingSemanticCacheStore.ts"; @@ -144,6 +145,7 @@ import { STREAM_READINESS_TIMEOUT_MS, ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, STREAM_RECOVERY, + DEFAULT_MAX_TOKENS, } from "../config/constants.ts"; import { createRecoverableStream, makeContinuationBody } from "../services/streamRecovery.ts"; import { @@ -1691,6 +1693,56 @@ export async function handleChatCore({ ); } + // Re-check the concrete target after all compression passes. Combo compatibility + // filtering is advisory and may preserve an all-incompatible pool; this is the + // hard boundary that prevents a too-large prompt (or a negative token budget) + // from reaching an OpenAI-compatible upstream such as NVIDIA NIM. + const finalCompressionBody = body + ? adaptBodyForCompression(body as Record).body + : null; + const finalMessages = + finalCompressionBody?.messages || + body?.contents || + body?.request?.contents || + (body?.input && typeof body.input === "object" && !Array.isArray(body.input) + ? body.input + : []); + const finalEstimatedInputTokens = + estimateTokens(finalMessages) + + (Array.isArray(body?.tools) ? estimateTokens(body.tools) : 0) + + estimateTokens(body?.system) + + estimateTokens(body?.instructions); + const finalContextLimit = getTokenLimit(provider, effectiveModel); + const outputBudget = enforceOutputTokenBudget( + body as Record, + finalEstimatedInputTokens, + finalContextLimit, + targetFormat === FORMATS.CLAUDE && sourceFormat !== FORMATS.CLAUDE ? DEFAULT_MAX_TOKENS : 0 + ); + if (!outputBudget.ok) { + const message = + `Input exceeds the context window for ${provider}/${effectiveModel}: ` + + `estimated ${outputBudget.estimatedInputTokens} input tokens, limit ${outputBudget.contextLimit}. ` + + "Reduce the prompt or route to a model with a larger context window."; + log?.warn?.("CONTEXT", message); + trackPendingRequest(model, provider, connectionId, false); + return createErrorResult( + HTTP_STATUS.BAD_REQUEST, + message, + null, + "context_length_exceeded", + "invalid_request_error" + ); + } + if (outputBudget.adjustedFields.length > 0) { + log?.info?.( + "CONTEXT", + `Adjusted invalid or oversized output token fields (${outputBudget.adjustedFields.join(", ")}); ` + + `${outputBudget.availableOutputTokens} tokens remain for output` + ); + } + body = outputBudget.body; + let translatedBody = body; const isClaudePassthrough = sourceFormat === FORMATS.CLAUDE && targetFormat === FORMATS.CLAUDE; const isClaudeCodeCompatible = isClaudeCodeCompatibleProvider(provider); diff --git a/open-sse/handlers/chatCore/outputTokenBudget.ts b/open-sse/handlers/chatCore/outputTokenBudget.ts new file mode 100644 index 0000000000..2d26d488c2 --- /dev/null +++ b/open-sse/handlers/chatCore/outputTokenBudget.ts @@ -0,0 +1,117 @@ +export const OUTPUT_TOKEN_FIELDS = [ + "max_tokens", + "max_completion_tokens", + "max_output_tokens", +] as const; + +export type OutputTokenBudgetResult = + | { + ok: true; + body: Record; + availableOutputTokens: number; + adjustedFields: string[]; + } + | { + ok: false; + estimatedInputTokens: number; + contextLimit: number; + }; + +type OutputTokenAdjustment = { field: string; value?: number; remove?: boolean }; + +function getOutputTokenAdjustment( + field: string, + value: unknown, + availableOutputTokens: number +): OutputTokenAdjustment | null { + if (typeof value !== "number") return null; + if (!Number.isFinite(value) || value <= 0) return { field, remove: true }; + + const capped = Math.min(Math.floor(value), availableOutputTokens); + return capped === value ? null : { field, value: capped }; +} + +function hasTranslatorOutputTokenLimit(body: Record): boolean { + return ["max_tokens", "max_completion_tokens"].some((field) => { + const value = body[field]; + return typeof value === "number" && Number.isFinite(value) && value > 0; + }); +} + +function adjustOutputTokenFields( + body: Record, + availableOutputTokens: number +): Pick, "body" | "adjustedFields"> { + const adjustments = OUTPUT_TOKEN_FIELDS.map((field) => + getOutputTokenAdjustment(field, body[field], availableOutputTokens) + ).filter((adjustment): adjustment is OutputTokenAdjustment => adjustment !== null); + if (adjustments.length === 0) return { body, adjustedFields: [] }; + + const nextBody = { ...body }; + for (const adjustment of adjustments) { + if (adjustment.remove) delete nextBody[adjustment.field]; + else nextBody[adjustment.field] = adjustment.value; + } + + return { body: nextBody, adjustedFields: adjustments.map(({ field }) => field) }; +} + +/** + * Enforce the target model's context budget immediately before translation. + * + * Compression and combo selection are best-effort: a request may still be too + * large for a concrete target, and some OpenAI-compatible gateways derive an + * internal max_tokens value by subtracting the prompt from the context window. + * Reject that target locally instead of allowing the derived value to become + * negative upstream. Positive client limits are capped to the remaining room; + * invalid numeric limits are removed. + */ +export function enforceOutputTokenBudget( + body: Record | null | undefined, + estimatedInputTokens: number, + contextLimit: number, + defaultOutputTokens = 0 +): OutputTokenBudgetResult { + const normalizedInputTokens = Math.max(0, Math.ceil(estimatedInputTokens)); + const normalizedContextLimit = Math.max(1, Math.floor(contextLimit)); + const normalizedDefaultOutputTokens = Math.max(0, Math.floor(defaultOutputTokens)); + const availableOutputTokens = normalizedContextLimit - normalizedInputTokens; + + if (availableOutputTokens < 1) { + return { + ok: false, + estimatedInputTokens: normalizedInputTokens, + contextLimit: normalizedContextLimit, + }; + } + + if (!body) { + if (normalizedDefaultOutputTokens > availableOutputTokens) { + return { + ok: false, + estimatedInputTokens: normalizedInputTokens, + contextLimit: normalizedContextLimit, + }; + } + return { + ok: true, + body: {}, + availableOutputTokens, + adjustedFields: [], + }; + } + + if ( + normalizedDefaultOutputTokens > availableOutputTokens && + !hasTranslatorOutputTokenLimit(body) + ) { + return { + ok: false, + estimatedInputTokens: normalizedInputTokens, + contextLimit: normalizedContextLimit, + }; + } + + const adjusted = adjustOutputTokenFields(body, availableOutputTokens); + return { ok: true, ...adjusted, availableOutputTokens }; +} diff --git a/tests/unit/chatcore-sanitization.test.ts b/tests/unit/chatcore-sanitization.test.ts index b92a098abf..50cf0fe12b 100644 --- a/tests/unit/chatcore-sanitization.test.ts +++ b/tests/unit/chatcore-sanitization.test.ts @@ -178,7 +178,7 @@ test("chatCore sanitization normalizes max_output_tokens into max_tokens", async }, }); - assert.equal(copied.call.body.max_tokens, 0); + assert.equal(copied.call.body.max_tokens, undefined); assert.equal("max_output_tokens" in copied.call.body, false); assert.equal(preserved.call.body.max_tokens, 7); assert.equal("max_output_tokens" in preserved.call.body, false); diff --git a/tests/unit/output-token-budget.test.ts b/tests/unit/output-token-budget.test.ts new file mode 100644 index 0000000000..d4a5aec53a --- /dev/null +++ b/tests/unit/output-token-budget.test.ts @@ -0,0 +1,66 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { enforceOutputTokenBudget } from "../../open-sse/handlers/chatCore/outputTokenBudget.ts"; + +test("rejects a prompt that cannot leave one output token", () => { + const result = enforceOutputTokenBudget({ max_tokens: 8192 }, 527_058, 128_000); + + assert.deepEqual(result, { + ok: false, + estimatedInputTokens: 527_058, + contextLimit: 128_000, + }); +}); + +test("caps a positive output budget to the target's remaining context", () => { + const input = { messages: [], max_tokens: 12_000 }; + const result = enforceOutputTokenBudget(input, 127_000, 128_000); + + assert.equal(result.ok, true); + if (!result.ok) return; + assert.equal(result.body.max_tokens, 1_000); + assert.equal(input.max_tokens, 12_000, "must not mutate the shared combo request body"); +}); + +test("removes non-positive numeric output limits before upstream dispatch", () => { + const result = enforceOutputTokenBudget( + { max_tokens: -398_464, max_completion_tokens: 0 }, + 1_000, + 128_000 + ); + + assert.equal(result.ok, true); + if (!result.ok) return; + assert.equal("max_tokens" in result.body, false); + assert.equal("max_completion_tokens" in result.body, false); +}); + +test("caps max_output_tokens to the target's remaining context", () => { + const result = enforceOutputTokenBudget({ max_output_tokens: 12_000 }, 127_000, 128_000); + + assert.equal(result.ok, true); + if (!result.ok) return; + assert.equal(result.body.max_output_tokens, 1_000); +}); + +test("accepts a missing request body when output budget remains", () => { + const result = enforceOutputTokenBudget(null, 1_000, 128_000); + + assert.deepEqual(result, { + ok: true, + body: {}, + availableOutputTokens: 127_000, + adjustedFields: [], + }); +}); + +test("rejects when a Claude target's default output budget does not fit", () => { + const result = enforceOutputTokenBudget({}, 70_000, 128_000, 64_000); + + assert.deepEqual(result, { + ok: false, + estimatedInputTokens: 70_000, + contextLimit: 128_000, + }); +});