diff --git a/changelog.d/fixes/13686-estimated-usage-guard.md b/changelog.d/fixes/13686-estimated-usage-guard.md new file mode 100644 index 0000000000..cb8cfbd2e0 --- /dev/null +++ b/changelog.d/fixes/13686-estimated-usage-guard.md @@ -0,0 +1 @@ +- **fix(usage):** mark locally estimated token usage in the call log (`_omniroute.usageEstimated` on the logged response) so operators can tell estimated counts and costs from provider-reported ones — covers OmniRoute's own estimate for streams without upstream usage and web executors that report `estimated: true`; billing, API-key budgets, quota-share and client payloads are unchanged ([#13686](https://github.com/diegosouzapw/OmniRoute/pull/13686)) — thanks @maxmad64bis diff --git a/open-sse/handlers/chatCore/attemptLogging.ts b/open-sse/handlers/chatCore/attemptLogging.ts index f1ddcc1678..b434645595 100644 --- a/open-sse/handlers/chatCore/attemptLogging.ts +++ b/open-sse/handlers/chatCore/attemptLogging.ts @@ -20,6 +20,7 @@ import type { VideoBridgeLogRedactionEntry } from "@/lib/guardrails/videoBridge" import { FORMATS } from "../../translator/formats.ts"; import { takeEarlyKeepaliveBytes } from "../../utils/earlyKeepaliveByteBuffer.ts"; import { sanitizeErrorMessage } from "../../utils/error.ts"; +import { isEstimatedUsage } from "../../utils/usageTracking.ts"; import { cloneBoundedChatLogPayload, truncateForLog } from "./logTruncation.ts"; import { attachLogMeta } from "./cacheUsageMeta.ts"; @@ -493,6 +494,9 @@ export function persistAttemptLogs(args: PersistAttemptLogsArgs, ctx: PersistAtt } : null, claudePromptCacheUsage: claudeCacheUsageMeta, + // Operators can tell estimated token counts (and the cost derived from them) + // apart from provider-reported ones. Log-only: billing is unchanged. + usageEstimated: isEstimatedUsage(tokens) ? true : null, }) ), error: error || null, diff --git a/open-sse/handlers/chatCore/quotaShareConsumption.ts b/open-sse/handlers/chatCore/quotaShareConsumption.ts index cb50ea2e22..70de289966 100644 --- a/open-sse/handlers/chatCore/quotaShareConsumption.ts +++ b/open-sse/handlers/chatCore/quotaShareConsumption.ts @@ -21,9 +21,8 @@ export async function scheduleQuotaShareConsumption(args: { }): Promise { if (!args.apiKeyId || !args.connectionId) return; try { - const { scheduleRecordConsumption, buildConsumptionCost } = await import( - "@/lib/quota/spendRecorder" - ); + const { scheduleRecordConsumption, buildConsumptionCost } = + await import("@/lib/quota/spendRecorder"); scheduleRecordConsumption( { apiKeyId: args.apiKeyId, diff --git a/open-sse/handlers/usageExtractor.ts b/open-sse/handlers/usageExtractor.ts index 789b1019ca..e0076112a9 100644 --- a/open-sse/handlers/usageExtractor.ts +++ b/open-sse/handlers/usageExtractor.ts @@ -2,6 +2,8 @@ * Extract usage from non-streaming response body * Handles different provider response formats */ +import { carryEstimatedUsageMarker } from "../utils/usageTracking.ts"; + export function extractUsageFromResponse(responseBody, provider) { if (!responseBody || typeof responseBody !== "object") return null; const providerId = typeof provider === "string" ? provider.toLowerCase() : ""; @@ -23,7 +25,7 @@ export function extractUsageFromResponse(responseBody, provider) { responseBody.usage.prompt_tokens_details?.cache_write_tokens ?? responseBody.usage.input_tokens_details?.cache_write_tokens ?? responseBody.usage.cache_write_tokens; - return { + const openAiUsage = { prompt_tokens: responseBody.usage.prompt_tokens || 0, completion_tokens: responseBody.usage.completion_tokens || 0, // DeepSeek native API uses flat prompt_cache_hit_tokens (NOT @@ -60,6 +62,7 @@ export function extractUsageFromResponse(responseBody, provider) { ? { cost_in_usd_ticks: responseBody.usage.cost_in_usd_ticks } : {}), }; + return carryEstimatedUsageMarker(responseBody.usage, openAiUsage); } // Claude format diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index d71b858cc8..16339a95af 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -642,6 +642,33 @@ export function normalizeUsage(usage: UsageLike | null | undefined) { return normalized; } +// Internal marker for usage that was estimated locally (a web/cookie executor with no +// upstream metering). A NON-enumerable symbol: JSON.stringify, object spread and +// filterUsageForFormat never copy it, so it cannot reach a client payload or change any +// usage field, cost or budget — it only lets the call-log sink tell estimated usage apart +// after extraction rebuilt the object without the provider's `estimated` flag. +const ESTIMATED_USAGE_MARKER = Symbol.for("omniroute.usage.estimated"); + +export function carryEstimatedUsageMarker(source: unknown, rebuilt: T): T { + const estimated = + !!source && typeof source === "object" && (source as UsageLike).estimated === true; + if (estimated && rebuilt && typeof rebuilt === "object") { + Object.defineProperty(rebuilt, ESTIMATED_USAGE_MARKER, { value: true, enumerable: false }); + } + return rebuilt; +} + +/** + * True when token usage was estimated locally instead of reported by the provider: either + * the usage still carries `estimated: true` (OmniRoute's own estimateUsage fallback) or + * extraction kept the internal marker. Observability only — billing does not read it. + */ +export function isEstimatedUsage(usage: unknown): boolean { + if (!usage || typeof usage !== "object") return false; + if ((usage as UsageLike).estimated === true) return true; + return Reflect.get(usage, ESTIMATED_USAGE_MARKER) === true; +} + /** * Check if usage has valid token data * Valid = has at least one token field with value > 0 @@ -786,7 +813,7 @@ export function extractUsage(chunk: UsagePayloadLike | null | undefined) { typeof chunk.usage === "object" && (chunk.usage.prompt_tokens !== undefined || chunk.usage.input_tokens !== undefined) ) { - return normalizeUsage({ + const normalized = normalizeUsage({ prompt_tokens: chunk.usage.prompt_tokens ?? chunk.usage.input_tokens ?? 0, completion_tokens: chunk.usage.completion_tokens ?? chunk.usage.output_tokens ?? 0, cached_tokens: @@ -804,6 +831,7 @@ export function extractUsage(chunk: UsagePayloadLike | null | undefined) { // xAI's exact provider-reported cost (port of decolua/9router#2453, capability A). cost_in_usd_ticks: chunk.usage.cost_in_usd_ticks, }); + return carryEstimatedUsageMarker(chunk.usage, normalized); } // Gemini format (Antigravity) diff --git a/tests/unit/estimated-usage-billing-guard.test.ts b/tests/unit/estimated-usage-billing-guard.test.ts new file mode 100644 index 0000000000..193843772c --- /dev/null +++ b/tests/unit/estimated-usage-billing-guard.test.ts @@ -0,0 +1,199 @@ +// Estimated token usage: billing stays exactly as it is, and the call log records that the +// counts were estimated. Drives the real handleChatCore (non-streaming and streaming) with a +// fetch stub, then reads the persisted call log and the API-key spend ledger. +import { after, before, 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-estimated-usage-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const callLogs = await import("../../src/lib/usage/callLogs.ts"); +const { getDailyTotal } = await import("../../src/domain/costRules.ts"); +const { handleChatCore } = await import("../../open-sse/handlers/chatCore.ts"); +const { extractUsage, filterUsageForFormat, isEstimatedUsage } = + await import("../../open-sse/utils/usageTracking.ts"); +const { extractUsageFromResponse } = await import("../../open-sse/handlers/usageExtractor.ts"); + +const originalFetch = globalThis.fetch; +const silentLog = { debug() {}, info() {}, warn() {}, error() {} }; +const MODEL = "gpt-4o-mini"; + +before(() => { + core.resetDbInstance(); +}); + +after(async () => { + globalThis.fetch = originalFetch; + await callLogs.closeCallLogSaves(5_000); + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); +}); + +const USAGE = { prompt_tokens: 1200, completion_tokens: 800, total_tokens: 2000 }; + +function jsonCompletion(usage: Record): Response { + return new Response( + JSON.stringify({ + id: "chatcmpl-estimated", + object: "chat.completion", + model: MODEL, + choices: [{ index: 0, message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage, + }), + { status: 200, headers: { "content-type": "application/json" } } + ); +} + +function sseCompletion(events: unknown[]): Response { + const body = events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("") + "data: [DONE]\n\n"; + return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } }); +} + +const textChunk = (content: string, finish: string | null = null) => ({ + id: "chatcmpl-estimated", + object: "chat.completion.chunk", + model: MODEL, + choices: [{ index: 0, delta: { content }, finish_reason: finish }], +}); + +async function runChat(apiKeyId: string, stream: boolean, response: () => Response) { + globalThis.fetch = (async () => response()) as typeof fetch; + const body = { model: MODEL, stream, messages: [{ role: "user", content: "hello there" }] }; + const result = (await handleChatCore({ + body, + modelInfo: { provider: "openai", model: MODEL, extendedContext: false }, + credentials: { apiKey: "sk-test-estimated" }, + clientRawRequest: { + endpoint: "/v1/chat/completions", + body, + headers: new Headers({ accept: stream ? "text/event-stream" : "application/json" }), + }, + apiKeyInfo: { id: apiKeyId, name: apiKeyId }, + userAgent: "unit-test", + isCombo: false, + log: silentLog, + } as unknown as Parameters[0])) as { response?: Response }; + const clientText = result.response ? await result.response.text() : ""; + return clientText; +} + +async function persistedLog(apiKeyId: string) { + const deadline = Date.now() + 15_000; + for (;;) { + await callLogs.waitForCallLogSaves(5_000); + const rows = (await callLogs.getCallLogs({})) as Array<{ id: string; apiKeyId: string }>; + const row = rows.find((r) => r.apiKeyId === apiKeyId); + if (row) return callLogs.getCallLogById(row.id); + if (Date.now() > deadline) throw new Error(`no call log for ${apiKeyId}`); + await new Promise((r) => setTimeout(r, 50)); + } +} + +async function spend(apiKeyId: string): Promise { + const deadline = Date.now() + 5_000; + let total = getDailyTotal(apiKeyId); + while (total === 0 && Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 50)); + total = getDailyTotal(apiKeyId); + } + return total; +} + +function usageEstimatedMeta(entry: unknown): unknown { + const responseBody = (entry as { responseBody?: { _omniroute?: Record } }) + ?.responseBody; + return responseBody?._omniroute?.usageEstimated; +} + +test("extraction keeps an internal estimated marker that never serializes or spreads", () => { + const estimatedChunk = { choices: [], usage: { ...USAGE, estimated: true } }; + const reportedChunk = { choices: [], usage: { ...USAGE } }; + const estimated = extractUsage(estimatedChunk); + const reported = extractUsage(reportedChunk); + assert.equal(isEstimatedUsage(estimated), true); + assert.equal(isEstimatedUsage(reported), false); + assert.deepStrictEqual(estimated, reported, "token fields are untouched"); + assert.equal(JSON.stringify(estimated), JSON.stringify(reported)); + assert.equal(isEstimatedUsage({ ...estimated }), false, "spread copies never carry it"); + assert.equal(isEstimatedUsage(filterUsageForFormat(estimated, "openai")), false); + + const fromResponse = extractUsageFromResponse({ usage: { ...USAGE, estimated: true } }, "x"); + assert.equal(isEstimatedUsage(fromResponse), true); + assert.equal(JSON.stringify(fromResponse).includes("estimated"), false); + assert.equal(isEstimatedUsage(extractUsageFromResponse({ usage: { ...USAGE } }, "x")), false); +}); + +test("non-streaming estimated usage is still billed and is marked in the call log", async () => { + const clientText = await runChat("key-json-estimated", false, () => + jsonCompletion({ ...USAGE, estimated: true }) + ); + assert.ok((await spend("key-json-estimated")) > 0, "API-key spend still records the cost"); + const entry = await persistedLog("key-json-estimated"); + assert.equal(entry?.tokens?.in, USAGE.prompt_tokens); + assert.equal(entry?.tokens?.out, USAGE.completion_tokens); + assert.equal(usageEstimatedMeta(entry), true); + assert.doesNotMatch(clientText, /usageEstimated/); +}); + +test("non-streaming provider-reported usage carries no estimated marker", async () => { + await runChat("key-json-reported", false, () => jsonCompletion({ ...USAGE })); + assert.ok((await spend("key-json-reported")) > 0); + const entry = await persistedLog("key-json-reported"); + assert.equal(usageEstimatedMeta(entry), undefined); +}); + +test("a stream without upstream usage is billed on the estimate and marked in the call log", async () => { + const clientText = await runChat("key-sse-silent", true, () => + sseCompletion([textChunk("hello from the model"), textChunk("", "stop")]) + ); + assert.match(clientText, /hello from the model/); + assert.ok((await spend("key-sse-silent")) > 0, "API-key spend still records the estimate"); + const entry = await persistedLog("key-sse-silent"); + assert.ok((entry?.tokens?.out ?? 0) > 0); + assert.equal(usageEstimatedMeta(entry), true); + assert.doesNotMatch(clientText, /usageEstimated/); +}); + +test("a stream whose executor reports estimated usage is billed and marked in the call log", async () => { + const clientText = await runChat("key-sse-executor", true, () => + sseCompletion([ + textChunk("hello from the model"), + textChunk("", "stop"), + { + id: "chatcmpl-estimated", + object: "chat.completion.chunk", + model: MODEL, + choices: [], + usage: { ...USAGE, estimated: true }, + }, + ]) + ); + assert.ok((await spend("key-sse-executor")) > 0); + const entry = await persistedLog("key-sse-executor"); + assert.equal(entry?.tokens?.in, USAGE.prompt_tokens); + assert.equal(usageEstimatedMeta(entry), true); + assert.doesNotMatch(clientText, /usageEstimated/); +}); + +test("a stream with provider-reported usage carries no estimated marker", async () => { + await runChat("key-sse-reported", true, () => + sseCompletion([ + textChunk("hello from the model"), + textChunk("", "stop"), + { + id: "chatcmpl-estimated", + object: "chat.completion.chunk", + model: MODEL, + choices: [], + usage: { ...USAGE }, + }, + ]) + ); + const entry = await persistedLog("key-sse-reported"); + assert.equal(entry?.tokens?.in, USAGE.prompt_tokens); + assert.equal(usageEstimatedMeta(entry), undefined); +});