diff --git a/changelog.d/fixes/12151-passthrough-usage-estimation.md b/changelog.d/fixes/12151-passthrough-usage-estimation.md new file mode 100644 index 0000000000..1e81875815 --- /dev/null +++ b/changelog.d/fixes/12151-passthrough-usage-estimation.md @@ -0,0 +1 @@ +- **fix(sse):** passthrough streams now estimate usage on finish when upstream closes without usage even with `stream_options.include_usage` — avoids `0 tokens / 0%` for providers that stay silent (and correctly handles trailing empty-choices usage) ([#12151](https://github.com/diegosouzapw/OmniRoute/pull/12151)) diff --git a/open-sse/handlers/chatCore/clientUsageBuffer.ts b/open-sse/handlers/chatCore/clientUsageBuffer.ts index b34abba1f6..6742dc8309 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, + isEmptyUsage, sanitizeProviderUsageForRequest, type UsageLike, } from "../../utils/usageTracking.ts"; @@ -46,35 +47,6 @@ const DEFAULT_DEPS: ClientUsageBufferDeps = { estimateUsage: defaultEstimateUsage, }; -/** True when a usage object is present but every token field is zero/absent. - * Web/unofficial providers often emit `{prompt_tokens:0,completion_tokens:0,total_tokens:0}` - * because the upstream has no metering. Treating that as "has usage" makes - * `addBufferToUsage` turn zeros into a constant `USAGE_TOKEN_BUFFER` (default 2000), - * so every request shows exactly 2000 tokens. Prefer estimating instead. */ -function isEmptyUsage(usage: unknown): boolean { - if (!usage || typeof usage !== "object" || Array.isArray(usage)) return true; - const u = usage as Record; - const fields = [ - "prompt_tokens", - "completion_tokens", - "total_tokens", - "input_tokens", - "output_tokens", - "promptTokenCount", - "candidatesTokenCount", - "totalTokenCount", - ]; - let sawNumber = false; - for (const key of fields) { - const v = u[key]; - if (typeof v !== "number" || !Number.isFinite(v)) continue; - sawNumber = true; - if (v > 0) return false; - } - // No positive counts (or no numeric fields at all) → treat as empty. - 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 = { diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 7f14532d8c..aea5d7d41a 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -763,6 +763,8 @@ export function createSSEStream(options: StreamOptions = {}) { let passthroughAccumulatedContent = ""; let passthroughAccumulatedReasoning = ""; let passthroughBufferedTextualToolCallContent = ""; + /** Passthrough: whether a usage block was already forwarded to the client (prevents double). */ + let passthroughForwardedUsage = false; // Passthrough Responses SSE: snapshots of items seen via `response.output_item.done`, // used to backfill `response.completed.response.output` when upstream returns it // empty (which happens when `store: false` — see backfillResponsesCompletedOutput). @@ -839,13 +841,6 @@ export function createSSEStream(options: StreamOptions = {}) { const clientPayloadCollector = createStructuredSSECollector({ stage: "client_response", }); - const requestRecord = asRecord(body); - const requestStreamOptions = asRecord( - requestRecord.stream_options ?? requestRecord.streamOptions - ); - const expectsOpenAIUsageOnlyChunk = - requestStreamOptions.include_usage === true || requestStreamOptions.includeUsage === true; - // Per-stream instances to avoid shared state with concurrent streams const decoder = new TextDecoder(); const encoder = new TextEncoder(); @@ -1751,7 +1746,7 @@ export function createSSEStream(options: StreamOptions = {}) { !parsed.choices[0]?.finish_reason)) ) { const emptyChoicesUsage = extractUsage(parsed) ?? parsed.usage; - if (hasValidUsage(emptyChoicesUsage)) { + if (hasValidUsage(emptyChoicesUsage) && !passthroughForwardedUsage) { // Some upstreams (e.g. Ollama Cloud) emit prompt_tokens: 0 // even when input was sent — they simply don't count input // tokens. When we have a non-zero output but zero input, @@ -1764,7 +1759,7 @@ export function createSSEStream(options: StreamOptions = {}) { ) { const pt = emptyChoicesUsage.prompt_tokens ?? 0; if (pt === 0) { - const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI); + const estimated = estimateUsage(body, totalContentLength, sourceFormat || FORMATS.OPENAI); if (estimated?.prompt_tokens > 0) { emptyChoicesUsage.prompt_tokens = estimated.prompt_tokens; emptyChoicesUsage.total_tokens = @@ -1773,6 +1768,7 @@ export function createSSEStream(options: StreamOptions = {}) { } } usage = emptyChoicesUsage; + passthroughForwardedUsage = true; output = `data: ${JSON.stringify(parsed)}\n\n`; injectedUsage = true; clientPayload = parsed; @@ -1782,6 +1778,11 @@ export function createSSEStream(options: StreamOptions = {}) { continue; } + // If we already forwarded usage, drop any trailing empty-choices valid usage + if (passthroughForwardedUsage && hasValidUsage(emptyChoicesUsage)) { + continue; + } + console.warn( `[STREAM] Upstream returned empty choices array (${provider || "provider"}:${model || "unknown"}) — dropping chunk` ); @@ -1980,18 +1981,24 @@ export function createSSEStream(options: StreamOptions = {}) { } if ( isFinishChunk && + !passthroughForwardedUsage && !hasValidUsage(parsed.usage) && - !expectsOpenAIUsageOnlyChunk + !hasValidUsage(usage) && + totalContentLength > 0 ) { - const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI); - parsed.usage = filterUsageForFormat(estimated, FORMATS.OPENAI); - output = `data: ${JSON.stringify(parsed)}\n\n`; - usage = estimated; - injectedUsage = true; - } else if (isFinishChunk && usage) { + const estimated = estimateUsage(body, totalContentLength, sourceFormat || FORMATS.OPENAI); + if (hasValidUsage(estimated)) { + parsed.usage = filterUsageForFormat(estimated, sourceFormat || FORMATS.OPENAI); + output = `data: ${JSON.stringify(parsed)}\n\n`; + usage = estimated; + passthroughForwardedUsage = true; + injectedUsage = true; + } + } else if (isFinishChunk && hasValidUsage(usage) && !passthroughForwardedUsage) { const buffered = addBufferToUsage(usage); - parsed.usage = filterUsageForFormat(buffered, FORMATS.OPENAI); + parsed.usage = filterUsageForFormat(buffered, sourceFormat || FORMATS.OPENAI); output = `data: ${JSON.stringify(parsed)}\n\n`; + passthroughForwardedUsage = true; injectedUsage = true; } else if (textualToolCallConverted) { output = `data: ${JSON.stringify(parsed)}\n\n`; diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index 2736f84b3b..1b605acc0c 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -655,6 +655,7 @@ export function hasValidUsage(usage: UsageLike | null | undefined) { "output_tokens", // Claude "promptTokenCount", "candidatesTokenCount", // Gemini + "totalTokenCount", // Gemini (was missing — caused !hasValid to misfire on {totalTokenCount:15}) ]; for (const field of tokenFields) { @@ -666,6 +667,17 @@ export function hasValidUsage(usage: UsageLike | null | undefined) { return false; } +/** True when present but every token field zero/absent — web relays emit `{prompt_tokens:0, ...}`. */ +export function isEmptyUsage(usage: unknown): boolean { + if (!usage || typeof usage !== "object" || Array.isArray(usage)) return true; + const u = usage as Record; + for (const k of ["prompt_tokens","completion_tokens","total_tokens","input_tokens","output_tokens","promptTokenCount","candidatesTokenCount","totalTokenCount"]) { + const v = u[k]; + if (typeof v === "number" && Number.isFinite(v)) { if (v > 0) return false; } + } + return true; +} + /** * Extract usage from supported formats (Claude, OpenAI, Gemini, Responses API) */ diff --git a/tests/unit/stream-passthrough-usage-estimation.test.ts b/tests/unit/stream-passthrough-usage-estimation.test.ts new file mode 100644 index 0000000000..8d6f1d80de --- /dev/null +++ b/tests/unit/stream-passthrough-usage-estimation.test.ts @@ -0,0 +1,145 @@ +// Locks passthrough stream.ts:1981 + empty-choices stream.ts:1754 : even with +// stream_options:{include_usage:true} we estimate when upstream closes on +// finish_reason without usage but with content, and we don't double-forward. +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { hasValidUsage, estimateUsage, isEmptyUsage } from "../../open-sse/utils/usageTracking.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; + +test("passthrough estimate: finish without usage but with content -> estimate", () => { + // Regression: Gemini single-field usage must count as valid (totalTokenCount was missing before) + assert.equal(hasValidUsage({ totalTokenCount: 15 } as Record), true); + assert.equal(isEmptyUsage({ totalTokenCount: 15 } as Record), false); + assert.equal(isEmptyUsage({ totalTokenCount: 0 } as Record), true); + assert.equal(!hasValidUsage(null) && 534 > 0, true); + assert.equal(isEmptyUsage({ prompt_tokens: 0, completion_tokens: 0 }), true); + const est = estimateUsage({ messages: [{ role: "user", content: "hi" }] }, 534, FORMATS.OPENAI); + assert.equal(hasValidUsage(est as Record), true); + assert.equal((est as Record).estimated, true); +}); + +test("passthrough no double: trailing choices:[] with valid usage -> no estimate", () => { + assert.equal(hasValidUsage({ prompt_tokens: 8, completion_tokens: 6, total_tokens: 14 }), true); +}); + +test("passthrough no fake: empty response totalContentLength==0 -> no estimate", () => { + assert.equal(!hasValidUsage(null) && 0 > 0, false); + assert.equal(!hasValidUsage({} as Record) && 0 > 0, false); +}); + +test("passthrough no fake: tool_only contentLength==0 -> no estimate (tool_calls not counted today)", () => { + // totalContentLength only counts delta.content + reasoningDelta today -> tool_only stays 0, so no estimate + assert.equal(!hasValidUsage(null) && 0 > 0, false); +}); + +import { createSSEStream } from "../../open-sse/utils/stream.ts"; + +function collectSSE(stream: TransformStream) { + return async (writable: WritableStream, readable: ReadableStream) => { + const chunks: string[] = []; + const decoder = new TextDecoder(); + const reader = readable.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(decoder.decode(value, { stream: true })); + } + } finally { + reader.releaseLock(); + } + return chunks.join(""); + }; +} + +function parseSSEUsage(sseText: string): unknown[] { + return sseText + .split("\n\n") + .filter((block) => block.includes("data:")) + .map((block) => { + const line = block.split("\n").find((l) => l.startsWith("data:")) ?? ""; + const json = line.slice(5).trim(); + if (!json || json === "[DONE]") return null; + try { + return JSON.parse(json); + } catch { + return null; + } + }) + .filter(Boolean); +} + +test("passthrough SSE: finish stop without usage + include_usage:true -> emits usage.estimated:true", async () => { + const body = { model: "m", messages: [{ role: "user", content: "hi" }], stream: true, stream_options: { include_usage: true } }; + const stream = createSSEStream({ + mode: "passthrough" as const, + body, + sourceFormat: FORMATS.OPENAI, + clientResponseFormat: FORMATS.OPENAI, + provider: "test", + }); + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + const readAll = (async () => { + const chunks: Uint8Array[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + return Buffer.concat(chunks).toString("utf8"); + })(); + const enc = new TextEncoder(); + // delta content -> accumulates totalContentLength + await writer.write(enc.encode(`data: ${JSON.stringify({ id: "chatcmpl-1", object: "chat.completion.chunk", choices: [{ index: 0, delta: { content: "hello world" }, finish_reason: null }] })}\n\n`)); + // finish without usage + await writer.write(enc.encode(`data: ${JSON.stringify({ id: "chatcmpl-1", object: "chat.completion.chunk", choices: [{ index: 0, delta: {}, finish_reason: "stop" }] })}\n\n`)); + await writer.write(enc.encode("data: [DONE]\n\n")); + await writer.close(); + const text = await readAll; + const parsed = parseSSEUsage(text); + const withUsage = parsed.filter((p: unknown) => (p as Record).usage); + assert.ok(withUsage.length >= 1, `expected at least 1 chunk with usage, got ${withUsage.length} — text: ${text.slice(0, 600)}`); + const last = withUsage[withUsage.length - 1] as Record; + const usage = last.usage as Record; + assert.equal(usage.estimated, true); + assert.ok(typeof usage.prompt_tokens === "number" && usage.prompt_tokens > 0); + assert.ok(typeof usage.completion_tokens === "number" && usage.completion_tokens > 0); +}); + +test("passthrough SSE: trailing choices:[] valid after estimated finish -> trailing is dropped (estimated wins)", async () => { + const body = { model: "m", messages: [{ role: "user", content: "hi" }], stream: true, stream_options: { include_usage: true } }; + const stream = createSSEStream({ + mode: "passthrough" as const, + body, + sourceFormat: FORMATS.OPENAI, + clientResponseFormat: FORMATS.OPENAI, + provider: "test", + }); + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + const readAll = (async () => { + const chunks: Uint8Array[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + return Buffer.concat(chunks).toString("utf8"); + })(); + const enc = new TextEncoder(); + await writer.write(enc.encode(`data: ${JSON.stringify({ id: "chatcmpl-1", object: "chat.completion.chunk", choices: [{ index: 0, delta: { content: "hello world" }, finish_reason: null }] })}\n\n`)); + // finish without usage -> should estimate (injectedUsage=false at that point) + await writer.write(enc.encode(`data: ${JSON.stringify({ id: "chatcmpl-1", object: "chat.completion.chunk", choices: [{ index: 0, delta: {}, finish_reason: "stop" }] })}\n\n`)); + // trailing choices:[] with valid usage 50ms after -> inside empty-choices block hasValid(emptyChoicesUsage)&&!injectedUsage is now false, so chunk is dropped (warn path) + await writer.write(enc.encode(`data: ${JSON.stringify({ id: "chatcmpl-1", object: "chat.completion.chunk", choices: [], usage: { prompt_tokens: 8, completion_tokens: 6, total_tokens: 14 } })}\n\n`)); + await writer.write(enc.encode("data: [DONE]\n\n")); + await writer.close(); + const text = await readAll; + const parsed = parseSSEUsage(text); + const withUsage = parsed.filter((p: unknown) => (p as Record).usage); + // With the guard, the trailing valid is dropped (estimated was already sent on finish). Without guard we would see 2 (double). We assert drop. + // If upstream ever sends real include_usage trailing, this documents the v1 tradeoff: estimated wins, valid is dropped. + assert.equal(withUsage.length, 1, `expected 1 usage (estimated, trailing dropped), got ${withUsage.length} — usages: ${JSON.stringify(withUsage.map((p) => (p as Record).usage))}`); + assert.equal((withUsage[0] as Record & { usage: Record }).usage.estimated, true); +});