[BUG] Streaming output_tokens estimated as bytes//40 instead of provider completion_tokens #2712

Title: [BUG] Streaming output_tokens falls back to total_bytes // 40 when no usage frame is parsed — over-counts completion_tokens by ~3-6x (provider-dependent)

Description

When a streaming response reaches the end-of-stream accounting and the proxy was unable to parse an output_tokens / completion_tokens value from the SSE stream, it fabricates one by integer-dividing the accumulated wire byte count by 40:

headroom/proxy/handlers/streaming.py:881-889

output_tokens = stream_state["output_tokens"]
output_tokens_source = "provider"
if output_tokens is None:
    output_tokens = stream_state["total_bytes"] // 40
    output_tokens_source = "estimated_bytes"
    logger.warning(
        f"[{request_id}] Could not parse output_tokens from SSE, "
        f"estimating {output_tokens} from {stream_state['total_bytes']} bytes"
    )

total_bytes is accumulated at streaming.py:1422 as stream_state["total_bytes"] += len(chunk) for every raw upstream chunk — i.e. the full SSE wire bytes including data: framing, the JSON envelope, role/finish fields, and any tool-call deltas. It is not a content-byte count, and the 40 divisor is not calibrated to any tokenizer.

The fallback fires whenever stream_state["output_tokens"] is still None at line 881. That happens in at least two situations the surrounding code already knows about:

1. The stream is truncated before the usage event arrives (client disconnect, network drop, connection reset). The tail-flush logic at lines 850-880 exists precisely because the message_start / message_delta usage events can sit unparsed in the residual buffer — and when that flush still comes up empty, execution falls straight into the // 40 estimator.
2. The upstream simply does not emit a usage frame on the stream (e.g. an OpenAI Chat Completions stream where the client did not set stream_options.include_usage, or a provider that never puts completion_tokens on the wire).

Because the bytes-per-token ratio of a raw SSE stream is provider- and chunking-specific, the error has no fixed sign or magnitude. On the providers I tested it over-counts: roughly 3.4x-6.4x the real completion_tokens (GLM ~3.4x, Moonshot/Kimi ~6.4x). The fabricated number then flows into RequestOutcome (line 892 tags output_tokens_source=estimated_bytes, but the output_tokens value itself is populated as if real) and so into downstream cost / savings accounting.

To Reproduce

The fallback is the last line of normal stream accounting, so any stream that omits or loses its usage frame triggers it. Two ways to see it:

A. In the running proxy (provider-agnostic). Point a streaming request at the proxy with a provider/path that does not emit a usage frame, or kill the connection mid-stream. Watch the proxy log:

[<request_id>] Could not parse output_tokens from SSE, estimating <N> from <M> bytes

<N> is M // 40. Compare it to the completion_tokens the same provider reports on a non-streaming call (or on a stream that does carry usage) for identical content — they will not agree, and the gap is provider-specific.

B. Self-contained replay of the estimator (no API keys, SYNTHETIC stream).

git clone https://github.com/headroomlabs-ai/headroom.git
cd headroom
# Python >=3.10 + a Rust toolchain (maturin build backend). Either:
uv sync                          # repo ships uv.lock
# ...or: pip install -e .
python3 /tmp/replay_estimator.py

/tmp/replay_estimator.py (the byte stream below is synthetic — fabricated to show the estimator, not captured from a real provider):

#!/usr/bin/env python3
# Replays the estimator at headroom/proxy/handlers/streaming.py:884 on a
# SYNTHETIC recorded stream. No provider keys needed.
import json

# --- SYNTHETIC Chat-Completions-style stream ------------------------------
# Fine-grained chunks: each carries ~1 token of content but a full JSON envelope,
# which is exactly what inflates total_bytes relative to completion_tokens.
content_chunks = [
    b"data: " + json.dumps({"choices": [{"index": 0, "delta": {"content": f"word{i} "}}]}).encode() + b"\n\n"
    for i in range(120)
]
# The usage frame the provider WOULD emit if the stream completed cleanly:
usage_chunk = b'data: {"choices": [], "usage": {"completion_tokens": 120}}\n\n'

def account(frames, *, usage_was_parsed: bool):
    total_bytes = 0
    for raw in frames:
        total_bytes += len(raw)                       # streaming.py:1422
    output_tokens = 120 if usage_was_parsed else None  # streaming.py:1483-1484
    source = "provider"
    if output_tokens is None:
        output_tokens = total_bytes // 40             # streaming.py:884 (the line under test)
        source = "estimated_bytes"
    return total_bytes, output_tokens, source

tb_a, ot_a, src_a = account(content_chunks + [usage_chunk], usage_was_parsed=True)
tb_b, ot_b, src_b = account(content_chunks,          usage_was_parsed=False)  # truncated / no usage
truth = 120

print(f"truth (provider completion_tokens)      = {truth}")
print(f"A usage-present: total_bytes={tb_a} output_tokens={ot_a} source={src_a}")
print(f"B usage-missing: total_bytes={tb_b} output_tokens={ot_b} source={src_b}")
print(f"B bias vs truth  = {ot_b / truth:.1f}x  (SYNTHETIC stream; ratio is set by chunk")
print("                    granularity, not a calibrated constant)")

Sample output on my machine:

truth (provider completion_tokens)      = 120
A usage-present: total_bytes=14957 output_tokens=120 source=provider
B usage-missing: total_bytes=13457 output_tokens=336 source=estimated_bytes
B bias vs truth  = 2.8x  (SYNTHETIC stream; ratio is set by chunk
                    granularity, not a calibrated constant)

Case A is correct (the provider value is used). Case B is the bug: with the exact same content, dropping the usage frame makes the proxy log 336 output tokens instead of 120, and that 336 is what gets recorded. The 2.8x in this synthetic run is arbitrary — it is whatever the framing produces — which is the core problem. On live upstreams I measured the same fallback over-counting by about 3.4x-6.4x depending on provider (GLM ~3.4x, Moonshot/Kimi ~6.4x).

Expected Behavior

output_tokens should be the provider's reported completion_tokens when available. When no usage frame could be parsed, the proxy should not substitute a fabricated, uncalibrated total_bytes // 40 into the cost / savings pipeline — it should record the value as unknown / null (and let the output_tokens_source tag already present at line 892 gate it out of any spend math), or at minimum make the estimate conservative and clearly marked so it cannot be mistaken for a real token count.

Actual Behavior

On any stream that arrives at line 881 without a parsed usage frame, the proxy records total_bytes // 40 as output_tokens with only a logger.warning and the estimated_bytes tag. Because total_bytes is raw SSE wire bytes (framing + JSON envelope, accumulated per chunk at line 1422), the result diverges from the real completion_tokens by a provider-dependent factor — ~3.4x-6.4x over across the two providers I tested.

Code Sample

See the replay_estimator.py snippet in To Reproduce — it exercises the exact line (output_tokens = stream_state["total_bytes"] // 40).

Error Output

No exception is raised; the failure is a silent accounting error. The only signal is the warning log line at streaming.py:886-889:

[<request_id>] Could not parse output_tokens from SSE, estimating <N> from <M> bytes

…and an output_tokens_source=estimated_bytes tag on the RequestOutcome. Neither blocks the value from being consumed as if it were a real token count.

Environment

• Headroom version: headroom-ai 0.33.0 (current main, headroom/proxy/handlers/streaming.py lines 881-889 / 1422)
• Python version: 3.10+
• OS: macOS / Linux
• LLM Provider: any; the fallback is provider-agnostic and fires whenever the stream omits or truncates the usage frame. Over-count magnitude is provider-specific (observed ~3.4x on GLM, ~6.4x on Moonshot/Kimi).

Additional Context

• The tail-flush block immediately above (streaming.py:850-880) shows the code already knows the usage event can be lost to truncation; the // 40 fallback is what catches the cases that flush still does not recover.
• The output_tokens_source tag is already there, so a minimal fix is to stop populating output_tokens with the fabricated estimate and instead emit None / 0 with the estimated_bytes tag, then have cost / savings consumers skip records whose source is not provider. That keeps telemetry (stream happened, byte volume seen) without letting an uncalibrated number into spend accounting.
• The over-count direction and size depend on provider wire format (chunk granularity, JSON verbosity, CJK vs. latin content), so I would avoid replacing 40 with a different single constant — it would just move the bias. Either carry the real completion_tokens through, or mark the value unknown.