diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 4b229e24a8..c06a97639e 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -99,7 +99,7 @@ import { resolveAgentGoalPolicy } from "../utils/agentGoalPolicy.ts"; import { createStreamController } from "../utils/streamHandler.ts"; import * as streamFailure from "../utils/streamFailureFinalization.ts"; import { createSseHeartbeatTransform, shapeForClientFormat } from "../utils/sseHeartbeat.ts"; -import { addBufferToUsage, filterUsageForFormat, estimateUsage } from "../utils/usageTracking.ts"; +import { addBufferToUsage, filterUsageForFormat, estimateUsage, sanitizeUsagePayloadForRequest } from "../utils/usageTracking.ts"; import { refreshWithRetry, isUnrecoverableRefreshError, @@ -4257,8 +4257,8 @@ export async function handleChatCore({ } : responseBody ); + sanitizeUsagePayloadForRequest(responseBody, finalBody || translatedBody || body, responsePayloadFormat); effectiveServiceTier = resolveReportedServiceTier(responseBody) ?? effectiveServiceTier; - // Notify success - caller can clear error status if needed if (onRequestSuccess) { await onRequestSuccess(); @@ -4395,7 +4395,7 @@ export async function handleChatCore({ // #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, { + applyClientUsageBuffer(translatedResponse, finalBody || translatedBody || body, clientResponseFormat, { preserveContextBudgetInVisibleUsage: isClaudeCodeCompatible, }); diff --git a/open-sse/handlers/chatCore/clientUsageBuffer.ts b/open-sse/handlers/chatCore/clientUsageBuffer.ts index a916edc1de..4c7b1e48dc 100644 --- a/open-sse/handlers/chatCore/clientUsageBuffer.ts +++ b/open-sse/handlers/chatCore/clientUsageBuffer.ts @@ -22,6 +22,7 @@ import { addBufferToUsage as defaultAddBuffer, filterUsageForFormat as defaultFilterUsage, estimateUsage as defaultEstimateUsage, + sanitizeProviderUsageForRequest, } from "../../utils/usageTracking.ts"; type ResponseLike = @@ -103,6 +104,14 @@ export function applyClientUsageBuffer( deps: ClientUsageBufferDeps = DEFAULT_DEPS ): void { const { preserveContextBudgetInVisibleUsage = false } = options; + if (translatedResponse?.usage) { + translatedResponse.usage = sanitizeProviderUsageForRequest( + translatedResponse.usage, + body, + clientResponseFormat + ); + } + // Add buffer and filter usage for client (to prevent CLI context errors) if (translatedResponse?.usage && !isEmptyUsage(translatedResponse.usage)) { const buffered = deps.addBufferToUsage(translatedResponse.usage) as Record; diff --git a/open-sse/utils/passthroughTailProcessor.ts b/open-sse/utils/passthroughTailProcessor.ts index 23d04e8a27..ab45fb5474 100644 --- a/open-sse/utils/passthroughTailProcessor.ts +++ b/open-sse/utils/passthroughTailProcessor.ts @@ -31,6 +31,7 @@ export type PassthroughTailProcessorContext = { emitConvertedOutput: (output: string) => void; pushProviderPayload: (payload: unknown) => void; pushClientPayload: (payload: unknown) => void; + sanitizeUsagePayload: (payload: unknown) => boolean; setPassthroughResponsesId: (value: string) => void; setUsage: (value: unknown) => void; addTotalContentLength: (value: number) => void; @@ -284,6 +285,9 @@ export function processBufferedPassthroughLine( } const parsed = parsedPassthroughData as JsonRecord; + if (context.sanitizeUsagePayload(parsed)) { + output = `data: ${JSON.stringify(parsed)}\n\n`; + } const parsedType = typeof parsed.type === "string" ? parsed.type : ""; const isResponses = parsedType.startsWith("response."); const isClaude = context.isClaudeEventPayload(parsed); diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 8bc363017a..121b3950e6 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -9,6 +9,7 @@ import { addBufferToUsage, filterUsageForFormat, normalizeUsage as normalizeTokenUsage, + sanitizeUsagePayloadForRequest, } from "./usageTracking.ts"; import { parseSSELine, @@ -1317,7 +1318,10 @@ export function createSSEStream(options: StreamOptions = {}) { parsed.type.startsWith("content_block") || parsed.type === "ping" || parsed.type === "error"); - + if (sanitizeUsagePayloadForRequest(parsed, body, clientResponseFormat)) { + output = `data: ${JSON.stringify(parsed)}\n\n`; + injectedUsage = true; + } if (isResponsesSSE) { // #6199/#6561 — statefully drop internal commentary-phase output (see // ./responsesCommentaryDrop.ts) and clear the buffered `event:` line @@ -1985,13 +1989,11 @@ export function createSSEStream(options: StreamOptions = {}) { } if (shouldDropResponsesCommentary && dropCommentary(parsed as JsonRecord)) continue; - providerPayloadCollector.push(parsed); - if (parsed && parsed.done) { continue; } - + sanitizeUsagePayloadForRequest(parsed, body, targetFormat); if (parsed.choices?.[0]?.delta?.tool_calls) { lastToolCallChunkTime = now; } @@ -2193,6 +2195,7 @@ export function createSSEStream(options: StreamOptions = {}) { }, pushProviderPayload: (payload: unknown) => providerPayloadCollector.push(payload), pushClientPayload: (payload: unknown) => clientPayloadCollector.push(payload), + sanitizeUsagePayload: (payload: unknown) => sanitizeUsagePayloadForRequest(payload, body, clientResponseFormat), setPassthroughResponsesId: (value: string) => { passthroughResponsesId = value; }, @@ -2241,7 +2244,6 @@ export function createSSEStream(options: StreamOptions = {}) { return; } } - const bufferedLine = buffer.trim(); if (skipPassthroughEvent || /^event:\s*keepalive\b/i.test(bufferedLine)) { skipPassthroughEvent = false; @@ -2254,6 +2256,7 @@ export function createSSEStream(options: StreamOptions = {}) { const bufferedPayload = parseSSELine(bufferedLine); if (bufferedPayload) { providerPayloadCollector.push(bufferedPayload); + if (sanitizeUsagePayloadForRequest(bufferedPayload, body, clientResponseFormat)) output = `data: ${JSON.stringify(bufferedPayload)}\n\n`; if ( shouldInjectClaudeEmptyResponseBeforeCurrentEvent( claudeEmptyResponseLifecycle, @@ -2267,7 +2270,7 @@ export function createSSEStream(options: StreamOptions = {}) { updateClaudeEmptyResponseLifecycle(claudeEmptyResponseLifecycle, bufferedPayload); } clientPayloadCollector.push(bufferedPayload); - + // Normalize numeric IDs for final buffered data: chunk (same as transform path) if (typeof bufferedPayload === "object" && !Array.isArray(bufferedPayload)) { const flushedParsed = bufferedPayload as JsonRecord; const flushedType = diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index 2382cb3f5b..e2abc2a67d 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -152,11 +152,11 @@ export function addBufferToUsage(usage) { result.context_budget_prompt_tokens = result.prompt_tokens + buffer; } - // Calculate or update the context-budget total + // Keep real total_tokens intact and calculate separate context-budget headroom. if (result.total_tokens !== undefined) { 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 (real value — not buffered) + // Calculate a real total if the provider omitted it. result.total_tokens = result.prompt_tokens + result.completion_tokens; result.context_budget_total_tokens = result.total_tokens + buffer; } @@ -282,6 +282,232 @@ export function filterUsageForFormat(usage, targetFormat) { return pickFields(fields); } +// Provider usage is normally authoritative, but compatibility gateways can return +// stale/cumulative cache counters. A token cannot encode less than one UTF-8 byte, +// so a stateless request's input count must remain related to the complete wire +// body. The 2x multiplier plus fixed allowance deliberately tolerates provider +// templates, tokenization differences, and format translation while still catching +// catastrophic values such as 336k tokens for a 115 KB request. +const INPUT_USAGE_BYTE_MULTIPLIER = 2; +const INPUT_USAGE_FIXED_ALLOWANCE = 8192; + +const REMOTE_CONTEXT_REFERENCE_KEYS = new Set([ + "previous_response_id", + "previousResponseId", + "conversation_id", + "conversationId", + "thread_id", + "threadId", + "parent_message_id", + "parentMessageId", + "cached_content", + "cachedContent", + "file_id", + "fileId", + "image_url", + "imageUrl", + "audio_url", + "audioUrl", + "video_url", + "videoUrl", +]); + +function hasValue(value): boolean { + if (value === null || value === undefined || value === false) return false; + if (typeof value === "string") return value.trim().length > 0; + if (Array.isArray(value)) return value.length > 0; + if (typeof value === "object") return Object.keys(value).length > 0; + return true; +} + +function hasRemoteContextReference(value, depth = 0): boolean { + if (!value || typeof value !== "object" || depth > 8) return false; + + if (Array.isArray(value)) { + return value.some((item) => hasRemoteContextReference(item, depth + 1)); + } + + for (const [key, nested] of Object.entries(value)) { + if (REMOTE_CONTEXT_REFERENCE_KEYS.has(key) && hasValue(nested)) { + return true; + } + if (hasRemoteContextReference(nested, depth + 1)) { + return true; + } + } + return false; +} + +function getSerializedBodyBytes(body): number | null { + if (!body || typeof body !== "object" || hasRemoteContextReference(body)) return null; + try { + const serialized = JSON.stringify(body); + if (!serialized) return null; + return Buffer.byteLength(serialized, "utf8"); + } catch { + return null; + } +} + +function tokenNumber(value): number { + return typeof value === "number" && Number.isFinite(value) ? value : 0; +} + +/** + * Return true when a provider-reported input count is plausible for this request. + * `null`/unserializable bodies and server-side context references fail open. + */ +export function isInputTokenCountPlausible(inputTokens, body): boolean { + if (typeof inputTokens !== "number" || !Number.isFinite(inputTokens) || inputTokens < 0) { + return false; + } + + const bodyBytes = getSerializedBodyBytes(body); + if (bodyBytes === null) return true; + const maximum = bodyBytes * INPUT_USAGE_BYTE_MULTIPLIER + INPUT_USAGE_FIXED_ALLOWANCE; + return inputTokens <= maximum; +} + +function resolveUsageFormat(usage, targetFormat) { + if (targetFormat === FORMATS.CLAUDE) return FORMATS.CLAUDE; + if (targetFormat === FORMATS.GEMINI || targetFormat === FORMATS.ANTIGRAVITY) { + return FORMATS.GEMINI; + } + if (targetFormat === FORMATS.OPENAI_RESPONSES || targetFormat === FORMATS.OPENAI_RESPONSE) { + return FORMATS.OPENAI_RESPONSES; + } + if (targetFormat === FORMATS.OPENAI) return FORMATS.OPENAI; + + if (usage?.promptTokenCount !== undefined || usage?.candidatesTokenCount !== undefined) { + return FORMATS.GEMINI; + } + if ( + usage?.cache_read_input_tokens !== undefined || + usage?.cache_creation_input_tokens !== undefined + ) { + return FORMATS.CLAUDE; + } + if (usage?.input_tokens_details !== undefined) return FORMATS.OPENAI_RESPONSES; + return FORMATS.OPENAI; +} + +function getReportedInputTokens(usage, format): number { + if (format === FORMATS.CLAUDE) { + return ( + tokenNumber(usage.input_tokens) + + tokenNumber(usage.cache_read_input_tokens) + + tokenNumber(usage.cache_creation_input_tokens) + ); + } + if (format === FORMATS.GEMINI) { + return tokenNumber(usage.promptTokenCount); + } + if (format === FORMATS.OPENAI_RESPONSES) { + return tokenNumber(usage.input_tokens ?? usage.prompt_tokens); + } + return tokenNumber(usage.prompt_tokens ?? usage.input_tokens); +} + +function clearCachedTokenDetail(value) { + if (!value || typeof value !== "object" || Array.isArray(value)) return value; + const result = { ...value }; + if (result.cached_tokens !== undefined) result.cached_tokens = 0; + return result; +} + +/** + * Replace only physically implausible provider input/cache usage with the local + * request estimate. Valid usage is returned by reference and remains untouched. + */ +export function sanitizeProviderUsageForRequest(usage, body, targetFormat = null) { + if (!usage || typeof usage !== "object" || Array.isArray(usage)) return usage; + + const format = resolveUsageFormat(usage, targetFormat); + const reportedInput = getReportedInputTokens(usage, format); + if (reportedInput <= 0 || isInputTokenCountPlausible(reportedInput, body)) { + return usage; + } + + const estimatedInput = Math.max(1, estimateInputTokens(body)); + const result = { ...usage }; + + if (format === FORMATS.CLAUDE) { + result.input_tokens = estimatedInput; + result.cache_read_input_tokens = 0; + result.cache_creation_input_tokens = 0; + return result; + } + + if (format === FORMATS.GEMINI) { + const output = + tokenNumber(result.candidatesTokenCount) + tokenNumber(result.thoughtsTokenCount); + result.promptTokenCount = estimatedInput; + result.cachedContentTokenCount = 0; + if (result.totalTokenCount !== undefined) { + result.totalTokenCount = estimatedInput + output; + } + return result; + } + + if (format === FORMATS.OPENAI_RESPONSES) { + result.input_tokens = estimatedInput; + result.input_tokens_details = clearCachedTokenDetail(result.input_tokens_details); + result.cache_read_input_tokens = 0; + result.cache_creation_input_tokens = 0; + if (result.total_tokens !== undefined) { + result.total_tokens = estimatedInput + tokenNumber(result.output_tokens); + } + return result; + } + + result.prompt_tokens = estimatedInput; + result.cached_tokens = 0; + result.cache_read_input_tokens = 0; + result.cache_creation_input_tokens = 0; + result.prompt_tokens_details = clearCachedTokenDetail(result.prompt_tokens_details); + if (result.total_tokens !== undefined) { + result.total_tokens = estimatedInput + tokenNumber(result.completion_tokens); + } + return result; +} + +/** + * Sanitize the usage container used by native provider responses/SSE events. + * Returns true only when the payload was changed and must be re-serialized. + */ +export function sanitizeUsagePayloadForRequest(payload, body, targetFormat = null): boolean { + if (!payload || typeof payload !== "object" || Array.isArray(payload)) return false; + + const replaceUsage = (owner, key, format) => { + if (!owner || typeof owner !== "object" || !owner[key]) return false; + const sanitized = sanitizeProviderUsageForRequest(owner[key], body, format); + if (sanitized === owner[key]) return false; + owner[key] = sanitized; + return true; + }; + + if (payload.type === "message_start" && payload.message?.usage) { + return replaceUsage(payload.message, "usage", FORMATS.CLAUDE); + } + if (payload.type === "message_delta" && payload.usage) { + return replaceUsage(payload, "usage", FORMATS.CLAUDE); + } + if (payload.response?.usage) { + return replaceUsage(payload.response, "usage", FORMATS.OPENAI_RESPONSES); + } + if (payload.response?.usageMetadata) { + return replaceUsage(payload.response, "usageMetadata", FORMATS.GEMINI); + } + if (payload.usageMetadata) { + return replaceUsage(payload, "usageMetadata", FORMATS.GEMINI); + } + if (payload.usage) { + const format = payload.type === "message" ? FORMATS.CLAUDE : targetFormat; + return replaceUsage(payload, "usage", format); + } + return false; +} + /** * Normalize usage object - ensure all values are valid numbers */ diff --git a/src/app/api/v1/messages/count_tokens/route.ts b/src/app/api/v1/messages/count_tokens/route.ts index af18840799..80bbc5d77f 100644 --- a/src/app/api/v1/messages/count_tokens/route.ts +++ b/src/app/api/v1/messages/count_tokens/route.ts @@ -8,6 +8,7 @@ import { getModelInfo } from "@/sse/services/model"; import { extractApiKey, getProviderCredentials, isValidApiKey } from "@/sse/services/auth"; import { safeResolveProxy } from "@/sse/handlers/chatHelpers"; import * as log from "@/sse/utils/logger"; +import { isInputTokenCountPlausible } from "@omniroute/open-sse/utils/usageTracking.ts"; /** * Handle CORS preflight @@ -82,7 +83,11 @@ export async function POST(request) { }) ); - if (!counted || !Number.isFinite(counted.input_tokens)) { + if ( + !counted || + !Number.isFinite(counted.input_tokens) || + !isInputTokenCountPlausible(counted.input_tokens, body) + ) { return estimated; } diff --git a/tests/unit/chatcore-client-usage-buffer.test.ts b/tests/unit/chatcore-client-usage-buffer.test.ts index 5ca07f0da0..01d742999e 100644 --- a/tests/unit/chatcore-client-usage-buffer.test.ts +++ b/tests/unit/chatcore-client-usage-buffer.test.ts @@ -10,6 +10,7 @@ const { applyClientUsageBuffer } = await import("../../open-sse/handlers/chatCore/clientUsageBuffer.ts"); const { resolveChatCoreRequestFormat } = await import("../../open-sse/handlers/chatCore/requestFormat.ts"); +const { invalidateBufferTokensCache } = await import("../../open-sse/utils/usageTracking.ts"); function makeDeps(overrides: Record = {}) { const calls = { buffer: [] as unknown[], estimate: [] as unknown[], filter: [] as unknown[] }; @@ -157,3 +158,44 @@ test("without the option the visible usage keeps the real unbuffered #8331 numbe const filtered = calls.filter[0] as Record; assert.equal(filtered.prompt_tokens, 5, "default path must not inflate client-visible metering"); }); + +test("real client-visible usage is not inflated by the context safety buffer", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "2000"; + invalidateBufferTokensCache(); + + try { + const response: Record = { + usage: { prompt_tokens: 69, completion_tokens: 5, total_tokens: 74 }, + }; + applyClientUsageBuffer(response, { messages: [{ role: "user", content: "hello" }] }, "openai"); + + assert.deepEqual(response.usage, { + prompt_tokens: 69, + completion_tokens: 5, + total_tokens: 74, + }); + } finally { + if (saved === undefined) delete process.env.USAGE_TOKEN_BUFFER; + else process.env.USAGE_TOKEN_BUFFER = saved; + invalidateBufferTokensCache(); + } +}); + +test("usage is validated against the provider-bound body with injected context", () => { + const providerBody = { + system: "x".repeat(10_000), + messages: [{ role: "user", content: "hello" }], + }; + const response: Record = { + usage: { prompt_tokens: 15_000, completion_tokens: 5, total_tokens: 15_005 }, + }; + + applyClientUsageBuffer(response, providerBody, "openai"); + + assert.deepEqual(response.usage, { + prompt_tokens: 15_000, + completion_tokens: 5, + total_tokens: 15_005, + }); +}); diff --git a/tests/unit/messages-count-tokens-route.test.ts b/tests/unit/messages-count-tokens-route.test.ts index a066e9cb63..f3a345a433 100644 --- a/tests/unit/messages-count-tokens-route.test.ts +++ b/tests/unit/messages-count-tokens-route.test.ts @@ -216,3 +216,35 @@ test("messages/count_tokens falls back to estimate when real upstream count fail globalThis.fetch = originalFetch; } }); + +test("messages/count_tokens rejects an impossible provider count and uses the local estimate", async () => { + await seedConnection("anthropic", { apiKey: "sk-ant-impossible-count" }); + + const originalFetch = globalThis.fetch; + globalThis.fetch = async () => + new Response(JSON.stringify({ input_tokens: 336409 }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + + try { + const response = await POST( + new Request("http://localhost/api/v1/messages/count_tokens", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "anthropic/claude-opus-4.6", + messages: [{ role: "user", content: "hello world" }], + }), + }) + ); + + assert.equal(response.status, 200); + const body = (await response.json()) as CountTokensResponse; + assert.equal(body.input_tokens, 2); + assert.equal(body.source, "local"); + assert.equal(body.provider, undefined); + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/tests/unit/stream-impossible-input-usage.test.ts b/tests/unit/stream-impossible-input-usage.test.ts new file mode 100644 index 0000000000..1716f6851d --- /dev/null +++ b/tests/unit/stream-impossible-input-usage.test.ts @@ -0,0 +1,266 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-usage-sanity-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const { createSSEStream } = await import("../../open-sse/utils/stream.ts"); +const { sanitizeProviderUsageForRequest } = await import("../../open-sse/utils/usageTracking.ts"); +const { FORMATS } = await import("../../open-sse/translator/formats.ts"); + +const textEncoder = new TextEncoder(); + +async function readTransformed(chunks, options) { + const source = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(textEncoder.encode(chunk)); + } + controller.close(); + }, + }); + + return new Response(source.pipeThrough(createSSEStream(options))).text(); +} + +function parseSsePayloads(text: string): Array> { + return text + .split("\n") + .filter((line) => line.startsWith("data: ") && line.slice(6).trim() !== "[DONE]") + .map((line) => JSON.parse(line.slice(6))); +} + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("native Claude passthrough repairs impossible AgentRouter cache usage before forwarding", async () => { + const body = { + model: "claude-opus-4-8", + system: "You are a coding assistant.", + messages: [{ role: "user", content: "hey" }], + tools: Array.from({ length: 25 }, (_, index) => ({ + name: `tool_${index}`, + description: "A small test tool", + input_schema: { + type: "object", + properties: { value: { type: "string" } }, + }, + })), + }; + let onCompletePayload = null; + + const text = await readTransformed( + [ + `event: message_start\ndata: ${JSON.stringify({ + type: "message_start", + message: { + id: "msg_agentrouter", + type: "message", + role: "assistant", + model: "claude-opus-4-8", + content: [], + stop_reason: null, + stop_sequence: null, + usage: { + input_tokens: 2, + cache_creation_input_tokens: 8186, + cache_read_input_tokens: 328221, + output_tokens: 0, + }, + }, + })}\n\n`, + `event: content_block_start\ndata: ${JSON.stringify({ + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + })}\n\n`, + `event: content_block_delta\ndata: ${JSON.stringify({ + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "Hey!" }, + })}\n\n`, + `event: content_block_stop\ndata: ${JSON.stringify({ + type: "content_block_stop", + index: 0, + })}\n\n`, + `event: message_delta\ndata: ${JSON.stringify({ + type: "message_delta", + delta: { stop_reason: "end_turn", stop_sequence: null }, + usage: { output_tokens: 44 }, + })}\n\n`, + `event: message_stop\ndata: ${JSON.stringify({ type: "message_stop" })}\n\n`, + ], + { + mode: "passthrough", + provider: "agentrouter", + model: "claude-opus-4-8", + body, + clientResponseFormat: FORMATS.CLAUDE, + onComplete: (payload) => { + onCompletePayload = payload; + }, + } + ); + + const payloads = parseSsePayloads(text); + const start = payloads.find((payload) => payload.type === "message_start") as { + message: { usage: Record }; + }; + assert.ok(start, "message_start must still be forwarded"); + + const usage = start.message.usage; + const clientVisibleInput = + (usage.input_tokens || 0) + + (usage.cache_creation_input_tokens || 0) + + (usage.cache_read_input_tokens || 0); + const requestBytes = Buffer.byteLength(JSON.stringify(body), "utf8"); + assert.ok( + clientVisibleInput <= requestBytes * 2 + 8192, + `client-visible input usage must be plausible for ${requestBytes} request bytes, got ${clientVisibleInput}` + ); + assert.equal(usage.cache_creation_input_tokens || 0, 0); + assert.equal(usage.cache_read_input_tokens || 0, 0); + + assert.ok(onCompletePayload, "stream completion callback must run"); + const completedUsage = (onCompletePayload as { usage: Record }).usage; + assert.ok( + completedUsage.prompt_tokens <= requestBytes * 2 + 8192, + "internal usage/logging must use the repaired count too" + ); + assert.equal(completedUsage.completion_tokens, 44); +}); + +test("valid Claude cache usage remains untouched", () => { + const body = { + messages: [{ role: "user", content: "x".repeat(20_000) }], + }; + const usage = { + input_tokens: 20, + cache_creation_input_tokens: 500, + cache_read_input_tokens: 3000, + output_tokens: 12, + }; + + const sanitized = sanitizeProviderUsageForRequest(usage, body, FORMATS.CLAUDE); + assert.equal(sanitized, usage, "plausible provider usage should retain object identity"); + assert.deepEqual(sanitized, usage); +}); + +test("impossible input usage is repaired for OpenAI, Responses, and Gemini shapes", () => { + const body = { messages: [{ role: "user", content: "hello" }] }; + + const openai = sanitizeProviderUsageForRequest( + { + prompt_tokens: 300_000, + completion_tokens: 7, + total_tokens: 300_007, + prompt_tokens_details: { cached_tokens: 299_000 }, + }, + body, + FORMATS.OPENAI + ); + assert.ok(openai.prompt_tokens < 300_000); + assert.equal(openai.prompt_tokens_details.cached_tokens, 0); + assert.equal(openai.total_tokens, openai.prompt_tokens + 7); + + const responses = sanitizeProviderUsageForRequest( + { + input_tokens: 300_000, + output_tokens: 8, + total_tokens: 300_008, + input_tokens_details: { cached_tokens: 299_000 }, + }, + body, + FORMATS.OPENAI_RESPONSES + ); + assert.ok(responses.input_tokens < 300_000); + assert.equal(responses.input_tokens_details.cached_tokens, 0); + assert.equal(responses.total_tokens, responses.input_tokens + 8); + + const gemini = sanitizeProviderUsageForRequest( + { + promptTokenCount: 300_000, + candidatesTokenCount: 9, + thoughtsTokenCount: 3, + cachedContentTokenCount: 299_000, + totalTokenCount: 300_012, + }, + body, + FORMATS.GEMINI + ); + assert.ok(gemini.promptTokenCount < 300_000); + assert.equal(gemini.cachedContentTokenCount, 0); + assert.equal(gemini.totalTokenCount, gemini.promptTokenCount + 12); +}); + +test("server-side context and remote file references bypass the body-byte guard", () => { + const statefulBody = { + previous_response_id: "resp_previous", + input: "continue", + }; + const statefulUsage = { input_tokens: 300_000, output_tokens: 3 }; + assert.equal( + sanitizeProviderUsageForRequest(statefulUsage, statefulBody, FORMATS.OPENAI_RESPONSES), + statefulUsage + ); + + const remoteFileBody = { + messages: [ + { + role: "user", + content: [{ type: "input_file", file_id: "file_large_document" }], + }, + ], + }; + const remoteFileUsage = { prompt_tokens: 300_000, completion_tokens: 3 }; + assert.equal( + sanitizeProviderUsageForRequest(remoteFileUsage, remoteFileBody, FORMATS.OPENAI), + remoteFileUsage + ); +}); + +test("final usage frame without a trailing newline is still sanitized", async () => { + const body = { messages: [{ role: "user", content: "hello" }] }; + const text = await readTransformed( + [ + `data: ${JSON.stringify({ + id: "chatcmpl_tail", + object: "chat.completion.chunk", + choices: [], + usage: { + prompt_tokens: 300_000, + completion_tokens: 4, + total_tokens: 300_004, + prompt_tokens_details: { cached_tokens: 299_000 }, + }, + })}`, + ], + { + mode: "passthrough", + provider: "generic-openai-compatible", + model: "test-model", + body, + clientResponseFormat: FORMATS.OPENAI, + } + ); + + const payloads = parseSsePayloads(text); + const usagePayloads = payloads.filter((payload) => payload.usage); + assert.equal(usagePayloads.length, 1, "usage must not be duplicated during tail flush"); + const providerPayload = payloads.find((payload) => payload.id === "chatcmpl_tail"); + assert.ok(providerPayload, "the provider's final usage frame must be forwarded"); + assert.equal(usagePayloads[0], providerPayload); + const usage = providerPayload.usage as Record; + assert.ok((usage.prompt_tokens as number) < 300_000); + assert.equal((usage.prompt_tokens_details as Record).cached_tokens, 0); + assert.equal( + usage.total_tokens, + (usage.prompt_tokens as number) + (usage.completion_tokens as number) + ); +});