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); });