mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-09-19 21:32:20 +03:00
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:
1
changelog.d/fixes/12151-passthrough-usage-estimation.md
Normal file
1
changelog.d/fixes/12151-passthrough-usage-estimation.md
Normal 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))
|
||||
@@ -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> = {
|
||||
|
||||
@@ -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`;
|
||||
|
||||
@@ -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)
|
||||
*/
|
||||
|
||||
145
tests/unit/stream-passthrough-usage-estimation.test.ts
Normal file
145
tests/unit/stream-passthrough-usage-estimation.test.ts
Normal 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);
|
||||
});
|
||||
Reference in New Issue
Block a user