fix(sse): estimate usage in passthrough stream even with include_usage (#12151)

A passthrough stream could end with no usage even though the client asked for it via stream_options: {include_usage: true}, so providers that do not meter always showed 0 tokens. The fix estimates usage at the finish marker when the upstream stays silent (flagged estimated: true) and drops any duplicate trailing usage chunk so the client never sees two.

open-sse/utils/stream.ts:1982,1749 · open-sse/utils/usageTracking.ts:651,664

Six cases: the predicate (finish without usage but with content, trailing valid, empty response, tool-only) plus two SSE harness cases through createSSEStream passthrough.

Note on base: this branch forked 442 commits back and carried a base-red marker for #12109, which is now closed — the release tip has no open base-red issue. It merged cleanly against the current tip regardless.

Verified in a combined batch worktree: 174/174 focused tests across all 11 PRs of this batch (this PR's stream-passthrough-usage-estimation suite included), typecheck:core clean, check:cycles and check:docs-counts green.

Thanks @maxmad64bis.
This commit is contained in:
Dizzle
2026-09-01 16:50:15 +02:00
committed by GitHub
parent 2f33f2c20d
commit 5e6c9a92dc
5 changed files with 183 additions and 46 deletions

View File

@@ -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))

View File

@@ -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<string, unknown>;
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<string, string> = {

View File

@@ -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`;

View File

@@ -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<string, unknown>;
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)
*/

View File

@@ -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<string, unknown>), true);
assert.equal(isEmptyUsage({ totalTokenCount: 15 } as Record<string, unknown>), false);
assert.equal(isEmptyUsage({ totalTokenCount: 0 } as Record<string, unknown>), 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<string, unknown>), true);
assert.equal((est as Record<string, unknown>).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<string, unknown>) && 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<Uint8Array, Uint8Array>) {
return async (writable: WritableStream<Uint8Array>, readable: ReadableStream<Uint8Array>) => {
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<string, unknown>).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<string, unknown>;
const usage = last.usage as Record<string, unknown>;
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<string, unknown>).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<string, unknown>).usage))}`);
assert.equal((withUsage[0] as Record<string, unknown> & { usage: Record<string, unknown> }).usage.estimated, true);
});