fix(proxy): count output tokens from the stream's text, not its wire size (#3163)
## Description From a user's proxy log (Copilot Chat, 0.36.x), on every streamed turn: ``` WARNING Could not parse output_tokens from SSE, estimating 8 from 334 bytes ``` When an upstream sends no usage chunk, output tokens were estimated as `total_bytes // 40` over the **raw SSE wire** — every `data:` prefix, JSON envelope, `role` / `finish_reason` / `id` / `model` field and blank-line framing included. The divisor is a fudge for "bytes per token *including framing overhead*", so its error tracks **how chattily the answer was chunked** rather than how long the answer was. The same text split into more deltas scores higher purely for being split. GitHub's Copilot CAPI is one of the upstreams that omits the usage chunk, so this was every Copilot turn's output number — and output tokens feed both the output-shaping savings estimate and the cost model. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - New pure module `headroom/proxy/stream_output_tokens.py`. The stream's own text is already in the buffer at the estimation site (`_finalize_stream_response` receives `full_sse_data`), so extract it and count that instead of the wire. - Handles all three forwarded surfaces: OpenAI chat `choices[].delta.content`, OpenAI responses `*.delta`, Anthropic `content_block_delta`. - Counts **reasoning deltas and tool-call arguments** too — the provider bills those as output, so omitting them would under-count exactly the most expensive turns. - `bytes // 40` survives only as the last resort for a stream whose text could not be recovered. That is the upstream-error path, which reaches the finalizer with no stream text and has no generated text to count — so it keeps its previous behavior exactly. - The log line named the wrong basis (it always said "from N bytes"), so it now reports which rung produced the number. - Parsing is I/O-free and hardened against malformed input — it runs on the response path, where an exception would break a turn that had already succeeded. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ pytest tests/test_stream_output_tokens.py -q 21 passed in 0.23s $ pytest tests/ -q -k stream 502 passed, 19 skipped $ pytest tests/ -q # this branch 6 failed, 11394 passed, 587 skipped in 426.13s All 6 also fail on clean origin/main, same machine — pre-existing, not regressions: test_graceful_shutdown.py::test_run_server_installs_cancelled_error_filter test_learn/test_integration.py::TestCodexIntegration::test_full_pipeline test_release_workflows.py::test_no_native_tls_in_wheel_build_tree test_providers/test_deepseek.py::... (3 litellm pricing tests) $ ruff check headroom/ All checks passed! $ mypy headroom/proxy/stream_output_tokens.py Success ``` Coverage includes: per-surface extraction; reasoning/tool-argument deltas; multi-line `data:` fields (per the SSE spec); 10 malformed-input shapes that must yield `""` rather than raise; and the two properties that motivated the change — - **chunk-invariance**: the same text split one-delta vs per-character now yields the same count, where the wire estimator disagreed wildly; - **a short answer is never recorded as zero** (integer division would report 0 tokens for `"OK"`). ## Real Behavior Proof - **Environment:** macOS, Python 3.12.13, branch on `origin/main` @ `a3821378`. - **Exact command / steps:** ran the estimator over a synthetic OpenAI chat stream matching the reported shape. - **Observed result:** for a 144-byte stream carrying `"Hello there, this is the answer."` (31 chars), the old path yields `144 // 40 = 3` tokens; the new path extracts the text and yields `8`, tagged `estimated_text`. Chunking the same text per-character leaves the new count unchanged while the wire count changes substantially. - **Not tested:** no live Copilot CAPI stream was captured; SSE fixtures are synthetic. The count remains an approximation (`chars // 4`) — this makes the estimate track the answer instead of the framing, it does not make it exact. Where the provider does send usage, that value is still preferred and untouched. ## Runtime Rollout Safety - **Rollout-managed feature(s):** none. - **Minimum rollout channel:** n/a. - **Stable/default behavior changed:** only for streams with **no** provider usage chunk — reported output tokens become larger and more accurate. Provider-reported usage is preferred exactly as before. - **Kill switch / disable path:** n/a. `output_tokens_source` is already recorded on the outcome tags, so provider vs estimated vs byte-fallback stays distinguishable downstream. - **Unsafe override required:** none. - **Qualification impact:** output-shaping savings and cost estimates for affected upstreams shift to a better-grounded number. - **Rollback path:** revert the commit. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines Co-authored-by: Tejas Chopra <tejas@Tejass-MacBook-Pro.local> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra
@@ -27,6 +27,7 @@ if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
||||
from headroom.copilot_auth import apply_copilot_api_auth
|
||||
from headroom.proxy.stream_output_tokens import estimate_output_tokens
|
||||
|
||||
logger = logging.getLogger("headroom.proxy")
|
||||
|
||||
@@ -952,11 +953,27 @@ class StreamingMixin:
|
||||
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"
|
||||
# No usage chunk from the upstream. Count the stream's OWN TEXT
|
||||
# rather than dividing the raw wire by a constant: `total_bytes`
|
||||
# includes every `data:` prefix, JSON envelope and framing newline,
|
||||
# so its error tracked how chattily the answer was chunked instead
|
||||
# of how long the answer was. Falls back to the byte heuristic only
|
||||
# when no text could be recovered at all.
|
||||
output_tokens, output_tokens_source = estimate_output_tokens(
|
||||
sse_text=full_sse_data,
|
||||
total_bytes=stream_state["total_bytes"],
|
||||
)
|
||||
# Name the actual basis. The old message always said "from N bytes"
|
||||
# even though that is now only true for the fallback rung, and an
|
||||
# operator reading it needs to know which estimate they are looking
|
||||
# at before trusting the number.
|
||||
basis = (
|
||||
"counted from stream text"
|
||||
if output_tokens_source == "estimated_text"
|
||||
else f"estimated from {stream_state['total_bytes']} raw SSE bytes"
|
||||
)
|
||||
logger.warning(
|
||||
f"[{request_id}] Could not parse output_tokens from SSE, "
|
||||
f"estimating {output_tokens} from {stream_state['total_bytes']} bytes"
|
||||
f"[{request_id}] No usage chunk in SSE; output_tokens={output_tokens} ({basis})"
|
||||
)
|
||||
|
||||
outcome_tags = dict(tags or {})
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Recover output-token counts from a finished SSE stream.
|
||||
|
||||
When an upstream omits a usage chunk, the proxy has to estimate how many output
|
||||
tokens the turn produced. The estimate was ``total_bytes // 40`` over the RAW
|
||||
SSE WIRE — every ``data:`` prefix, every JSON envelope, every ``role``/
|
||||
``finish_reason``/``id``/``model`` field, blank-line framing included. On a
|
||||
Copilot chat turn that produced a short answer the log read:
|
||||
|
||||
Could not parse output_tokens from SSE, estimating 8 from 334 bytes
|
||||
|
||||
334 bytes of wire is mostly envelope; the generated text inside it was a
|
||||
fraction of that. The divisor 40 is a fudge for "bytes per token INCLUDING
|
||||
framing overhead", so its error scales with how chatty the framing is rather
|
||||
than with the answer — a stream split into many small deltas is punished for
|
||||
the split, and one delivered in a few large chunks is not.
|
||||
|
||||
The stream's own text is right there in the buffer, so extract it and count
|
||||
that instead. ``bytes // 40`` survives only as the last resort for a stream
|
||||
whose text could not be recovered at all.
|
||||
|
||||
Pure and I/O-free so the parsing is testable without a proxy or a network.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
# Bytes per token when nothing better is available. Applied to the raw wire, so
|
||||
# it must absorb SSE framing overhead as well as the text — which is exactly why
|
||||
# it is a poor estimator and a last resort.
|
||||
WIRE_BYTES_PER_TOKEN = 40
|
||||
|
||||
# Characters per token for extracted text. ~4 is the usual English/code
|
||||
# approximation and is applied to generated text ONLY, with no framing in it.
|
||||
TEXT_CHARS_PER_TOKEN = 4
|
||||
|
||||
|
||||
def _iter_sse_payloads(sse: str) -> Iterator[Any]:
|
||||
"""Yield each ``data:`` payload in an SSE stream as a parsed object.
|
||||
|
||||
Tolerates the two things real streams do that a naive split does not:
|
||||
an event whose ``data:`` is spread over multiple lines, and ``[DONE]``.
|
||||
"""
|
||||
for block in sse.split("\n\n"):
|
||||
lines = [ln for ln in block.split("\n") if ln.startswith("data:")]
|
||||
if not lines:
|
||||
continue
|
||||
# Multi-line data: fields concatenate, per the SSE spec.
|
||||
raw = "".join(ln[5:].lstrip() for ln in lines).strip()
|
||||
if not raw or raw == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
yield json.loads(raw)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
|
||||
|
||||
def extract_stream_text(sse: str) -> str:
|
||||
"""Return the assistant text carried by a completed SSE stream.
|
||||
|
||||
Handles the three surfaces this proxy forwards:
|
||||
|
||||
* OpenAI chat completions — ``choices[].delta.content``
|
||||
* OpenAI responses — ``response.output_text.delta`` / ``delta``
|
||||
* Anthropic messages — ``content_block_delta.delta.text``
|
||||
|
||||
Reasoning/thinking deltas are counted too: the provider bills them as
|
||||
output tokens, so omitting them would under-count exactly the turns where
|
||||
output is most expensive.
|
||||
"""
|
||||
if not sse:
|
||||
return ""
|
||||
|
||||
parts: list[str] = []
|
||||
for obj in _iter_sse_payloads(sse):
|
||||
if not isinstance(obj, dict):
|
||||
continue
|
||||
|
||||
# --- OpenAI chat completions -------------------------------------- #
|
||||
choices = obj.get("choices")
|
||||
if isinstance(choices, list):
|
||||
for choice in choices:
|
||||
if not isinstance(choice, dict):
|
||||
continue
|
||||
delta = choice.get("delta")
|
||||
if not isinstance(delta, dict):
|
||||
continue
|
||||
for key in ("content", "reasoning_content", "refusal"):
|
||||
value = delta.get(key)
|
||||
if isinstance(value, str):
|
||||
parts.append(value)
|
||||
# Tool-call arguments stream as text and are billed as output.
|
||||
tool_calls = delta.get("tool_calls")
|
||||
if isinstance(tool_calls, list):
|
||||
for call in tool_calls:
|
||||
fn = call.get("function") if isinstance(call, dict) else None
|
||||
args = fn.get("arguments") if isinstance(fn, dict) else None
|
||||
if isinstance(args, str):
|
||||
parts.append(args)
|
||||
continue
|
||||
|
||||
obj_type = obj.get("type")
|
||||
|
||||
# --- Anthropic messages ------------------------------------------- #
|
||||
if obj_type == "content_block_delta":
|
||||
delta = obj.get("delta")
|
||||
if isinstance(delta, dict):
|
||||
for key in ("text", "thinking", "partial_json"):
|
||||
value = delta.get(key)
|
||||
if isinstance(value, str):
|
||||
parts.append(value)
|
||||
continue
|
||||
|
||||
# --- OpenAI responses --------------------------------------------- #
|
||||
if isinstance(obj_type, str) and obj_type.endswith(".delta"):
|
||||
value = obj.get("delta")
|
||||
if isinstance(value, str):
|
||||
parts.append(value)
|
||||
continue
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def estimate_output_tokens(*, sse_text: str, total_bytes: int) -> tuple[int, str]:
|
||||
"""Return ``(tokens, source)`` for a stream with no provider usage chunk.
|
||||
|
||||
``source`` names which rung of the ladder produced the number so the caller
|
||||
can log it honestly rather than implying the provider reported it:
|
||||
|
||||
* ``estimated_text`` — counted from the generated text (good)
|
||||
* ``estimated_bytes`` — the raw-wire fallback (poor, last resort)
|
||||
"""
|
||||
text = extract_stream_text(sse_text)
|
||||
if text:
|
||||
# At least one token for any non-empty answer: integer division would
|
||||
# report 0 for a 1-3 character reply ("OK", "42"), and a turn that
|
||||
# produced output must never be recorded as having produced none.
|
||||
return max(1, len(text) // TEXT_CHARS_PER_TOKEN), "estimated_text"
|
||||
return max(0, total_bytes) // WIRE_BYTES_PER_TOKEN, "estimated_bytes"
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Output tokens must be counted from the stream's text, not its wire size.
|
||||
|
||||
When an upstream sends no usage chunk, the proxy estimated output tokens as
|
||||
``total_bytes // 40`` over the RAW SSE WIRE — ``data:`` prefixes, JSON
|
||||
envelopes, ``role``/``finish_reason``/``id``/``model`` fields and blank-line
|
||||
framing all included. From a field log (Copilot Chat, 0.36.x):
|
||||
|
||||
Could not parse output_tokens from SSE, estimating 8 from 334 bytes
|
||||
|
||||
The divisor is a fudge for "bytes per token including framing", so the error
|
||||
tracked how chattily the answer was chunked rather than how long it was: the
|
||||
same answer split into more deltas scores higher purely for being split.
|
||||
|
||||
GitHub's Copilot CAPI is one of the upstreams that omits the usage chunk, so
|
||||
this was every Copilot turn's output number — and output tokens feed the
|
||||
output-shaping savings estimate and the cost model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.stream_output_tokens import (
|
||||
TEXT_CHARS_PER_TOKEN,
|
||||
WIRE_BYTES_PER_TOKEN,
|
||||
estimate_output_tokens,
|
||||
extract_stream_text,
|
||||
)
|
||||
|
||||
|
||||
def _sse(*objs: dict, done: bool = True) -> str:
|
||||
out = "".join(f"data: {json.dumps(o)}\n\n" for o in objs)
|
||||
return out + ("data: [DONE]\n\n" if done else "")
|
||||
|
||||
|
||||
def _chat_delta(text: str) -> dict:
|
||||
return {"choices": [{"index": 0, "delta": {"content": text}}]}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Extraction, per surface
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_openai_chat_deltas() -> None:
|
||||
sse = _sse(
|
||||
{"choices": [{"delta": {"role": "assistant"}}]},
|
||||
_chat_delta("Hello "),
|
||||
_chat_delta("world"),
|
||||
{"choices": [{"delta": {}, "finish_reason": "stop"}]},
|
||||
)
|
||||
assert extract_stream_text(sse) == "Hello world"
|
||||
|
||||
|
||||
def test_anthropic_content_block_deltas() -> None:
|
||||
sse = _sse(
|
||||
{"type": "message_start", "message": {"id": "msg_1"}},
|
||||
{"type": "content_block_delta", "delta": {"type": "text_delta", "text": "abc"}},
|
||||
{"type": "content_block_delta", "delta": {"type": "text_delta", "text": "def"}},
|
||||
{"type": "message_stop"},
|
||||
done=False,
|
||||
)
|
||||
assert extract_stream_text(sse) == "abcdef"
|
||||
|
||||
|
||||
def test_openai_responses_deltas() -> None:
|
||||
sse = _sse(
|
||||
{"type": "response.output_text.delta", "delta": "part one "},
|
||||
{"type": "response.output_text.delta", "delta": "part two"},
|
||||
)
|
||||
assert extract_stream_text(sse) == "part one part two"
|
||||
|
||||
|
||||
def test_reasoning_and_tool_arguments_are_billed_output_too() -> None:
|
||||
"""Omitting these under-counts exactly the most expensive turns."""
|
||||
sse = _sse(
|
||||
{"choices": [{"delta": {"reasoning_content": "thinking hard"}}]},
|
||||
{"choices": [{"delta": {"tool_calls": [{"function": {"arguments": '{"path":"a.py"}'}}]}}]},
|
||||
)
|
||||
text = extract_stream_text(sse)
|
||||
assert "thinking hard" in text
|
||||
assert '{"path":"a.py"}' in text
|
||||
|
||||
|
||||
def test_anthropic_thinking_and_partial_json() -> None:
|
||||
sse = _sse(
|
||||
{"type": "content_block_delta", "delta": {"thinking": "plan"}},
|
||||
{"type": "content_block_delta", "delta": {"partial_json": '{"a":1}'}},
|
||||
done=False,
|
||||
)
|
||||
assert extract_stream_text(sse) == 'plan{"a":1}'
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Malformed input must never raise — this runs on the response path
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.parametrize(
|
||||
"sse",
|
||||
[
|
||||
"",
|
||||
"data: not-json\n\n",
|
||||
"data: [DONE]\n\n",
|
||||
"garbage without data prefix\n\n",
|
||||
'data: {"choices": "not-a-list"}\n\n',
|
||||
'data: {"choices": [null]}\n\n',
|
||||
'data: {"choices": [{"delta": null}]}\n\n',
|
||||
'data: {"choices": [{"delta": {"content": 42}}]}\n\n',
|
||||
'data: {"type": "content_block_delta", "delta": "not-a-dict"}\n\n',
|
||||
"data:\n\n",
|
||||
],
|
||||
)
|
||||
def test_malformed_streams_yield_empty_not_an_exception(sse: str) -> None:
|
||||
assert extract_stream_text(sse) == ""
|
||||
|
||||
|
||||
def test_multi_line_data_fields_concatenate() -> None:
|
||||
"""Per the SSE spec, and real streams do it."""
|
||||
payload = json.dumps(_chat_delta("joined"))
|
||||
half = len(payload) // 2
|
||||
sse = f"data: {payload[:half]}\ndata: {payload[half:]}\n\n"
|
||||
assert extract_stream_text(sse) == "joined"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# The estimate itself
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_text_beats_the_wire_heuristic_on_the_reported_shape() -> None:
|
||||
"""A chunky stream: framing dominates the wire, so bytes//40 misreads it."""
|
||||
sse = _sse(*[_chat_delta(w) for w in ("The ", "quick ", "brown ", "fox ", "jumps")])
|
||||
total_bytes = len(sse.encode())
|
||||
|
||||
tokens, source = estimate_output_tokens(sse_text=sse, total_bytes=total_bytes)
|
||||
|
||||
assert source == "estimated_text"
|
||||
# "The quick brown fox jumps" is 25 chars -> 6 tokens.
|
||||
assert tokens == len("The quick brown fox jumps") // TEXT_CHARS_PER_TOKEN
|
||||
# The old estimator scored this stream far higher purely for its framing.
|
||||
assert total_bytes // WIRE_BYTES_PER_TOKEN > tokens
|
||||
|
||||
|
||||
def test_chunking_no_longer_changes_the_answer() -> None:
|
||||
"""Same text, different delta split — the count must not move."""
|
||||
text = "identical content across both streams"
|
||||
one = _sse(_chat_delta(text))
|
||||
many = _sse(*[_chat_delta(c) for c in text])
|
||||
|
||||
a, _ = estimate_output_tokens(sse_text=one, total_bytes=len(one.encode()))
|
||||
b, _ = estimate_output_tokens(sse_text=many, total_bytes=len(many.encode()))
|
||||
|
||||
assert a == b
|
||||
# And the wire-based estimator would have disagreed wildly.
|
||||
assert len(one.encode()) // WIRE_BYTES_PER_TOKEN != len(many.encode()) // WIRE_BYTES_PER_TOKEN
|
||||
|
||||
|
||||
def test_a_short_answer_is_never_recorded_as_zero() -> None:
|
||||
sse = _sse(_chat_delta("OK"))
|
||||
tokens, source = estimate_output_tokens(sse_text=sse, total_bytes=len(sse.encode()))
|
||||
assert source == "estimated_text"
|
||||
assert tokens == 1
|
||||
|
||||
|
||||
def test_falls_back_to_bytes_when_no_text_is_recoverable() -> None:
|
||||
"""The upstream-error path reaches here with no stream text at all."""
|
||||
tokens, source = estimate_output_tokens(sse_text="", total_bytes=800)
|
||||
assert source == "estimated_bytes"
|
||||
assert tokens == 800 // WIRE_BYTES_PER_TOKEN
|
||||
|
||||
|
||||
def test_negative_or_zero_bytes_are_safe() -> None:
|
||||
assert estimate_output_tokens(sse_text="", total_bytes=0) == (0, "estimated_bytes")
|
||||
assert estimate_output_tokens(sse_text="", total_bytes=-5) == (0, "estimated_bytes")
|
||||
Reference in New Issue
Block a user