From e8719783ef8a9d9c8a745c631ee848482a06da44 Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza Date: Fri, 24 Jul 2026 12:10:06 -0300 Subject: [PATCH] fix(api): stop the 2000-token safety buffer from inflating usage.prompt_tokens in the client response (#8331) (#8356) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(api): stop the 2000-token safety buffer from inflating usage.prompt_tokens in the client response (#8331) * fix(sse): scope #8331's usage-buffer fix around Claude-Code-compatible providers The #8331 fix correctly stopped folding the 2000-token context-window safety margin into client-visible prompt_tokens/input_tokens/total_tokens for normal API metering clients. But it also silently changed the response shape for Claude-Code-compatible providers, whose own context accounting reads the buffered number straight out of usage — regressing tests/unit/cc-compatible-provider.test.ts (expected 2007, got 7). Fold the computed context_budget_* fields back into the visible usage fields for that one path only (applyClientUsageBuffer's new preserveContextBudgetInVisibleUsage option, gated on the existing isClaudeCodeCompatible flag in chatCore.ts). Every other caller keeps the real, unbuffered #8331 numbers. --- .../fixes/8331-usage-buffer-inflation.md | 1 + open-sse/handlers/chatCore.ts | 7 +- .../handlers/chatCore/clientUsageBuffer.ts | 39 ++++- open-sse/utils/usageTracking.ts | 24 ++- .../unit/8331-usage-buffer-inflation.test.ts | 137 ++++++++++++++++++ tests/unit/usage-token-buffer.test.ts | 52 ++++--- 6 files changed, 232 insertions(+), 28 deletions(-) create mode 100644 changelog.d/fixes/8331-usage-buffer-inflation.md create mode 100644 tests/unit/8331-usage-buffer-inflation.test.ts diff --git a/changelog.d/fixes/8331-usage-buffer-inflation.md b/changelog.d/fixes/8331-usage-buffer-inflation.md new file mode 100644 index 0000000000..1ea77942a4 --- /dev/null +++ b/changelog.d/fixes/8331-usage-buffer-inflation.md @@ -0,0 +1 @@ +- fix(api): stop the 2000-token safety buffer from inflating usage.prompt_tokens in the client response (#8331) diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 368cb70095..c1231c79f6 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -4238,7 +4238,12 @@ export async function handleChatCore({ }); } - applyClientUsageBuffer(translatedResponse, body, clientResponseFormat); + // #8331: keep the client-visible metering fields real everywhere except Claude-Code-compatible + // providers, where Claude Code's own context accounting relies on the buffered number — see + // clientUsageBuffer.ts module docstring. + applyClientUsageBuffer(translatedResponse, body, clientResponseFormat, { + preserveContextBudgetInVisibleUsage: isClaudeCodeCompatible, + }); if (memoryOwnerId && memorySettings?.enabled && memorySettings.maxTokens > 0) { const requestMemoryText = extractMemoryTextFromRequestBody(body as Record); diff --git a/open-sse/handlers/chatCore/clientUsageBuffer.ts b/open-sse/handlers/chatCore/clientUsageBuffer.ts index 0f25a830cc..754c4f0811 100644 --- a/open-sse/handlers/chatCore/clientUsageBuffer.ts +++ b/open-sse/handlers/chatCore/clientUsageBuffer.ts @@ -7,6 +7,16 @@ * usage block, fall back to estimating from the serialized content length. Mutates * `translatedResponse.usage` in place — byte-identical to the previous inline block, including the * `?.usage` guard, the `JSON.stringify(... || "")` content-length, and the `> 0` estimate gate. + * + * #8331 scoping: `addBufferToUsage()` now keeps the safety margin OUT of the client-visible + * metering fields (prompt_tokens/input_tokens/total_tokens) for normal API clients, so billing + * reflects real upstream usage. Claude-Code-compatible providers are the one exception — the + * buffer's original purpose (see `usageTracking.ts` module docstring) is CLI context-window + * headroom, and Claude Code's own context accounting reads the buffered number straight out of + * the response `usage` block. `preserveContextBudgetInVisibleUsage` re-folds the computed + * `context_budget_*` fields back into the visible fields for that one path only, before + * `filterUsageForFormat()` strips the internal fields — every other caller keeps the real, + * unbuffered #8331 numbers. */ import { addBufferToUsage as defaultAddBuffer, @@ -60,15 +70,42 @@ function isEmptyUsage(usage: unknown): boolean { return true; } +/** context_budget_* → visible-field mapping folded back in for Claude-Code-compatible + * responses only (see module docstring above). */ +const CONTEXT_BUDGET_TO_VISIBLE_FIELD: Record = { + context_budget_prompt_tokens: "prompt_tokens", + context_budget_input_tokens: "input_tokens", + context_budget_total_tokens: "total_tokens", +}; + +function foldContextBudgetIntoVisibleUsage(usage: Record): void { + for (const [budgetField, visibleField] of Object.entries(CONTEXT_BUDGET_TO_VISIBLE_FIELD)) { + const value = usage[budgetField]; + if (typeof value === "number") { + usage[visibleField] = value; + } + } +} + +export interface ApplyClientUsageBufferOptions { + /** Claude-Code-compatible providers only (#8331 scoping) — see module docstring. */ + preserveContextBudgetInVisibleUsage?: boolean; +} + export function applyClientUsageBuffer( translatedResponse: ResponseLike, body: unknown, clientResponseFormat: unknown, + options: ApplyClientUsageBufferOptions = {}, deps: ClientUsageBufferDeps = DEFAULT_DEPS ): void { + const { preserveContextBudgetInVisibleUsage = false } = options; // Add buffer and filter usage for client (to prevent CLI context errors) if (translatedResponse?.usage && !isEmptyUsage(translatedResponse.usage)) { - const buffered = deps.addBufferToUsage(translatedResponse.usage); + const buffered = deps.addBufferToUsage(translatedResponse.usage) as Record; + if (preserveContextBudgetInVisibleUsage) { + foldContextBudgetIntoVisibleUsage(buffered); + } translatedResponse.usage = deps.filterUsageForFormat(buffered, clientResponseFormat); } else { // Fallback: estimate usage when provider returned no usage block diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index caef2c8f1b..90f457b77e 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -113,9 +113,18 @@ function getTimeString() { } /** - * Add buffer tokens to usage to prevent context errors + * Compute the context-window safety margin for a usage object, WITHOUT touching the + * client-visible/metering fields (prompt_tokens/input_tokens/total_tokens). + * + * #8331: the buffer used to be added directly into those fields, so a real 69-token + * request was reported to the client as 2069 while call_logs/the raw upstream body kept + * the true 69 — a metering/billing discrepancy. The safety margin this function computes + * is intentionally scoped to context-fit/CLI-headroom use only (see module docstring + * above); it is surfaced here as separate `context_budget_*` fields so it never gets + * confused with reported token accounting again. `filterUsageForFormat()` does not + * allow-list these fields, so they are stripped before any response reaches a client. * @param {object} usage - Usage object (supported format) - * @returns {object} Usage with buffer added + * @returns {object} Usage with context_budget_* fields added (metering fields unchanged) */ export function addBufferToUsage(usage) { if (!usage || typeof usage !== "object") return usage; @@ -134,20 +143,21 @@ export function addBufferToUsage(usage) { // Claude format if (result.input_tokens !== undefined) { - result.input_tokens += buffer; + result.context_budget_input_tokens = result.input_tokens + buffer; } // OpenAI format if (result.prompt_tokens !== undefined) { - result.prompt_tokens += buffer; + result.context_budget_prompt_tokens = result.prompt_tokens + buffer; } - // Calculate or update total_tokens + // Calculate or update the context-budget total if (result.total_tokens !== undefined) { - result.total_tokens += buffer; + result.context_budget_total_tokens = result.total_tokens + buffer; } else if (result.prompt_tokens !== undefined && result.completion_tokens !== undefined) { - // Calculate total_tokens if not exists + // Calculate total_tokens if not exists (real value — not buffered) result.total_tokens = result.prompt_tokens + result.completion_tokens; + result.context_budget_total_tokens = result.total_tokens + buffer; } return result; diff --git a/tests/unit/8331-usage-buffer-inflation.test.ts b/tests/unit/8331-usage-buffer-inflation.test.ts new file mode 100644 index 0000000000..28161f5311 --- /dev/null +++ b/tests/unit/8331-usage-buffer-inflation.test.ts @@ -0,0 +1,137 @@ +// Regression guard for #8331 — the DEFAULT_BUFFER_TOKENS (2000) context-window safety margin +// was leaking straight into the CLIENT-VISIBLE usage.prompt_tokens/input_tokens/total_tokens, +// so a real 69-token upstream request was reported to the client as 2069. call_logs.tokens_in +// and the raw upstream body always showed the true 69 — only the client-facing response was +// inflated. +// +// Fix: addBufferToUsage() (open-sse/utils/usageTracking.ts) no longer mutates the metering +// fields. The buffer is still computed (context-fit/CLI-headroom use is preserved) but is now +// surfaced only via separate context_budget_* fields, which filterUsageForFormat() does not +// allow-list — so it never reaches a client response, streaming or non-streaming. +// +// Covers both client-visible seams: +// 1. Non-streaming JSON usage — via applyClientUsageBuffer() (chatCore.ts:4236 call site). +// 2. Streaming SSE usage frame — via the same addBufferToUsage()+filterUsageForFormat() chain +// used verbatim by open-sse/utils/stream.ts:1018 and :1897. +import test from "node:test"; +import assert from "node:assert/strict"; +import { + addBufferToUsage, + filterUsageForFormat, + invalidateBufferTokensCache, +} from "../../open-sse/utils/usageTracking.ts"; +import { applyClientUsageBuffer } from "../../open-sse/handlers/chatCore/clientUsageBuffer.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; + +function resetEnv(saved: string | undefined) { + if (saved === undefined) { + delete process.env.USAGE_TOKEN_BUFFER; + } else { + process.env.USAGE_TOKEN_BUFFER = saved; + } +} + +test("#8331 non-streaming: client-facing usage.prompt_tokens equals raw upstream value, not +2000", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + // Reporter's real-world shape: 69 prompt tokens, 5 completion tokens. + const upstreamUsage = { prompt_tokens: 69, completion_tokens: 5, total_tokens: 74 }; + const translatedResponse: { + usage?: unknown; + choices?: Array<{ message?: { content?: unknown } }>; + } = { usage: upstreamUsage }; + + applyClientUsageBuffer(translatedResponse, { messages: [] }, "openai"); + + const clientUsage = translatedResponse.usage as Record; + assert.equal(clientUsage.prompt_tokens, 69, "must report the real upstream prompt_tokens"); + assert.equal(clientUsage.completion_tokens, 5); + assert.equal(clientUsage.total_tokens, 74, "must report the real upstream total_tokens"); + // The safety buffer must never leak into the client payload under any field name. + assert.equal("context_budget_prompt_tokens" in clientUsage, false); + assert.equal("context_budget_total_tokens" in clientUsage, false); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("#8331 streaming: SSE usage frame (addBufferToUsage + filterUsageForFormat chain) is not inflated", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + // Mirrors the exact call chain used at stream.ts:1018/:1897 to build the final SSE usage chunk. + const upstreamUsage = { prompt_tokens: 69, completion_tokens: 5, total_tokens: 74 }; + const buffered = addBufferToUsage(upstreamUsage); + const clientFrameUsage = filterUsageForFormat(buffered, FORMATS.OPENAI) as Record< + string, + unknown + >; + + assert.equal(clientFrameUsage.prompt_tokens, 69, "SSE usage frame must not be buffer-inflated"); + assert.equal(clientFrameUsage.total_tokens, 74); + assert.equal("context_budget_prompt_tokens" in clientFrameUsage, false); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("#8331 Claude-format streaming frame: input_tokens stays real, not buffered", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + const upstreamUsage = { input_tokens: 69, output_tokens: 5 }; + const buffered = addBufferToUsage(upstreamUsage); + const clientFrameUsage = filterUsageForFormat(buffered, FORMATS.CLAUDE) as Record< + string, + unknown + >; + + assert.equal(clientFrameUsage.input_tokens, 69); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("#8331 addBufferToUsage still computes the safety margin internally (buffer not deleted)", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + // The context-fit safety margin must still be computable for internal/future consumers — + // this fix decouples it from client metering, it does not remove the feature. + const result = addBufferToUsage({ + prompt_tokens: 69, + completion_tokens: 5, + total_tokens: 74, + }) as Record; + + assert.equal(result.prompt_tokens, 69); + assert.equal(result.context_budget_prompt_tokens, 2069, "buffer must still be computed"); + assert.equal(result.context_budget_total_tokens, 2074); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("#8331 estimated usage remains fully unbuffered (unaffected by the fix)", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + const result = addBufferToUsage({ + prompt_tokens: 6, + completion_tokens: 1, + total_tokens: 7, + estimated: true, + }) as Record; + + assert.equal(result.prompt_tokens, 6); + assert.equal("context_budget_prompt_tokens" in result, false); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); diff --git a/tests/unit/usage-token-buffer.test.ts b/tests/unit/usage-token-buffer.test.ts index ca06ac6882..27588b8a70 100644 --- a/tests/unit/usage-token-buffer.test.ts +++ b/tests/unit/usage-token-buffer.test.ts @@ -18,22 +18,25 @@ function resetEnv(saved: string | undefined) { // ─── addBufferToUsage — baseline / env-var path ─────────────────────────── -test("addBufferToUsage — adds DEFAULT 2000 when no env var and cache is null", () => { +test("addBufferToUsage — #8331: keeps DEFAULT 2000 out of client-visible prompt_tokens, exposes it via context_budget_prompt_tokens", () => { const saved = process.env.USAGE_TOKEN_BUFFER; delete process.env.USAGE_TOKEN_BUFFER; invalidateBufferTokensCache(); const result = addBufferToUsage({ prompt_tokens: 25, completion_tokens: 24, total_tokens: 49 }); - // Documents the race: after invalidation the sync path falls back to DEFAULT=2000 - assert.equal(result.prompt_tokens, 2025); + // Metering/client-visible fields stay the real upstream values (#8331 fix). + assert.equal(result.prompt_tokens, 25); assert.equal(result.completion_tokens, 24); - assert.equal(result.total_tokens, 2049); + assert.equal(result.total_tokens, 49); + // The safety margin is still computed, just decoupled from client metering. + assert.equal(result.context_budget_prompt_tokens, 2025); + assert.equal(result.context_budget_total_tokens, 2049); resetEnv(saved); }); -test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=0 env override", () => { +test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=0 env override (no context_budget_* fields added)", () => { const saved = process.env.USAGE_TOKEN_BUFFER; process.env.USAGE_TOKEN_BUFFER = "0"; invalidateBufferTokensCache(); @@ -42,31 +45,35 @@ test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=0 env override", () => { assert.equal(result.prompt_tokens, 25); assert.equal(result.total_tokens, 49); + assert.equal("context_budget_prompt_tokens" in result, false); resetEnv(saved); }); -test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=500 env override", () => { +test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=500 env override via context_budget_* fields", () => { const saved = process.env.USAGE_TOKEN_BUFFER; process.env.USAGE_TOKEN_BUFFER = "500"; invalidateBufferTokensCache(); const result = addBufferToUsage({ prompt_tokens: 86, completion_tokens: 52, total_tokens: 138 }); - assert.equal(result.prompt_tokens, 586); - assert.equal(result.total_tokens, 638); + assert.equal(result.prompt_tokens, 86); + assert.equal(result.total_tokens, 138); + assert.equal(result.context_budget_prompt_tokens, 586); + assert.equal(result.context_budget_total_tokens, 638); resetEnv(saved); }); -test("addBufferToUsage — also adds to Claude-format input_tokens", () => { +test("addBufferToUsage — also computes context_budget_input_tokens for Claude-format input_tokens", () => { const saved = process.env.USAGE_TOKEN_BUFFER; process.env.USAGE_TOKEN_BUFFER = "100"; invalidateBufferTokensCache(); const result = addBufferToUsage({ input_tokens: 40, output_tokens: 20 }); - assert.equal(result.input_tokens, 140); + assert.equal(result.input_tokens, 40); + assert.equal(result.context_budget_input_tokens, 140); resetEnv(saved); }); @@ -117,7 +124,7 @@ test("addBufferToUsage — skips safety buffer for estimated usage (web/heuristi // The fix: runtimeSettings.ts calls setBufferTokensCache(newValue) instead of // invalidateBufferTokensCache() so the correct value is available synchronously. -test("setBufferTokensCache(0) — immediately prevents buffer addition (no race window)", () => { +test("setBufferTokensCache(0) — immediately prevents context_budget_* addition (no race window)", () => { const saved = process.env.USAGE_TOKEN_BUFFER; delete process.env.USAGE_TOKEN_BUFFER; @@ -126,16 +133,17 @@ test("setBufferTokensCache(0) — immediately prevents buffer addition (no race const result = addBufferToUsage({ prompt_tokens: 25, completion_tokens: 24, total_tokens: 49 }); - // With the fix: 0 is applied synchronously — no 2000-token race window + // With the fix: 0 is applied synchronously — no context_budget_* fields at all assert.equal(result.prompt_tokens, 25); assert.equal(result.completion_tokens, 24); assert.equal(result.total_tokens, 49); + assert.equal("context_budget_prompt_tokens" in result, false); resetEnv(saved); invalidateBufferTokensCache(); }); -test("setBufferTokensCache(500) — immediately sets custom buffer value", () => { +test("setBufferTokensCache(500) — immediately sets custom buffer value in context_budget_* fields", () => { const saved = process.env.USAGE_TOKEN_BUFFER; delete process.env.USAGE_TOKEN_BUFFER; @@ -143,14 +151,16 @@ test("setBufferTokensCache(500) — immediately sets custom buffer value", () => const result = addBufferToUsage({ prompt_tokens: 86, completion_tokens: 52, total_tokens: 138 }); - assert.equal(result.prompt_tokens, 586); - assert.equal(result.total_tokens, 638); + assert.equal(result.prompt_tokens, 86); + assert.equal(result.total_tokens, 138); + assert.equal(result.context_budget_prompt_tokens, 586); + assert.equal(result.context_budget_total_tokens, 638); resetEnv(saved); invalidateBufferTokensCache(); }); -test("setBufferTokensCache(0) — works for Claude-format (input_tokens)", () => { +test("setBufferTokensCache(0) — works for Claude-format (input_tokens), no context_budget_input_tokens", () => { const saved = process.env.USAGE_TOKEN_BUFFER; delete process.env.USAGE_TOKEN_BUFFER; @@ -159,6 +169,7 @@ test("setBufferTokensCache(0) — works for Claude-format (input_tokens)", () => const result = addBufferToUsage({ input_tokens: 40, output_tokens: 20 }); assert.equal(result.input_tokens, 40); + assert.equal("context_budget_input_tokens" in result, false); resetEnv(saved); invalidateBufferTokensCache(); @@ -171,12 +182,15 @@ test("invalidateBufferTokensCache — still resets to null (returns DEFAULT on n // First prime the cache with a custom value setBufferTokensCache(0); const afterSet = addBufferToUsage({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); - assert.equal(afterSet.prompt_tokens, 10); // 0 buffer + assert.equal(afterSet.prompt_tokens, 10); // 0 buffer, real value either way - // Then invalidate — next sync call reverts to DEFAULT (2000) while async reload happens + // Then invalidate — next sync call reverts to DEFAULT (2000) while async reload happens. + // Client-visible prompt_tokens stays real (#8331); only the internal context_budget_* + // field carries the DEFAULT=2000 margin during the race window. invalidateBufferTokensCache(); const afterInvalidate = addBufferToUsage({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); - assert.equal(afterInvalidate.prompt_tokens, 2010); // DEFAULT=2000 applied (race window) + assert.equal(afterInvalidate.prompt_tokens, 10); + assert.equal(afterInvalidate.context_budget_prompt_tokens, 2010); // DEFAULT=2000 applied (race window) resetEnv(saved); });