akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): record cache reads/writes on backend-routed streaming (#327)

Browse source 
Two regressions surfaced as "Cache write: 0" in `headroom perf` and the
dashboard for every backend-routed streaming request (e.g. SvenMeyer's
DROID CLI > headroom > Azure GPT-5.5 setup):

* `_stream_openai_via_backend` parsed only `completion_tokens` and never
  read `prompt_tokens_details.cached_tokens` from the upstream usage
  frame. It also emitted no PERF log line at all, so `headroom perf`
  couldn't even count the request to report numbers. Now buffers SSE
  bytes, drains via `_parse_sse_usage_from_buffer(provider="openai")`,
  infers writes via `_infer_openai_cache_write_tokens` (only when the
  upstream actually reported usage — mirrors `_extract_responses_usage`),
  threads cache values into `record_request`, `cost_tracker.record_tokens`,
  the RequestLog, and a real PERF log line.

* `_stream_response_bedrock` hardcoded `cache_read=0 cache_write=0
  cache_hit_pct=0` in its PERF line regardless of what `message_start.usage`
  reported. Extended `stream_state` with `cache_read_input_tokens` and
  `cache_creation_input_tokens` (plus 5m/1h TTL buckets), captures them
  from `message_start`, threads through `record_request(cached=...)`,
  `cost_tracker.record_tokens(...)`, and `RequestLog(cache_hit=...)`.

Tests: four new tests in `test_backend_streaming_cache_metrics.py` cover
both paths plus a source-level regression guard against the hardcoded
zero string reappearing.
This commit is contained in:
chopratejas 2026-05-14 11:01:57 -07:00
parent bcf5517259
commit 3e97a0cb73
2 changed files with 490 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

151
headroom/proxy/handlers/streaming.py
View file

@ -1303,11 +1303,17 @@ class StreamingMixin:
start_time = time.time()
# Mutable state for the generator
# Mutable state for the generator. Cache fields mirror the
# native ``_finalize_stream_response`` shape so the PERF log
# values match between paths (issue #327).
stream_state: dict[str, Any] = {
"input_tokens": 0,
"output_tokens": 0,
"ttfb_ms": None,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_creation_ephemeral_5m_input_tokens": 0,
"cache_creation_ephemeral_1h_input_tokens": 0,
}
async def generate():
@ -1332,6 +1338,15 @@ class StreamingMixin:
usage = msg.get("usage", {})
if "input_tokens" in usage:
stream_state["input_tokens"] = usage["input_tokens"]
stream_state["cache_read_input_tokens"] = usage.get(
"cache_read_input_tokens", 0
)
stream_state["cache_creation_input_tokens"] = usage.get(
"cache_creation_input_tokens", 0
)
cw_5m, cw_1h = self._extract_anthropic_cache_ttl_metrics(usage)
stream_state["cache_creation_ephemeral_5m_input_tokens"] = cw_5m
stream_state["cache_creation_ephemeral_1h_input_tokens"] = cw_1h
# Track output tokens from message_delta
if event.event_type == "message_delta":
@ -1355,6 +1370,15 @@ class StreamingMixin:
# Record metrics
total_latency = (time.time() - start_time) * 1000
output_tokens = stream_state["output_tokens"]
cache_read_tokens = stream_state["cache_read_input_tokens"]
cache_write_tokens = stream_state["cache_creation_input_tokens"]
cache_write_5m_tokens = stream_state["cache_creation_ephemeral_5m_input_tokens"]
cache_write_1h_tokens = stream_state["cache_creation_ephemeral_1h_input_tokens"]
cache_hit_pct = (
round(cache_read_tokens / (cache_read_tokens + cache_write_tokens) * 100)
if (cache_read_tokens + cache_write_tokens) > 0
else 0
)
_backend_name = (
self.anthropic_backend.name if self.anthropic_backend else "anthropic"
@ -1371,7 +1395,7 @@ class StreamingMixin:
output_tokens=output_tokens,
tokens_saved=tokens_saved,
latency_ms=total_latency,
cached=False,
cached=cache_read_tokens > 0,
overhead_ms=optimization_latency,
ttfb_ms=stream_state["ttfb_ms"] or 0,
pipeline_timing=pipeline_timing,
@ -1379,7 +1403,15 @@ class StreamingMixin:
)
if self.cost_tracker:
self.cost_tracker.record_tokens(model, tokens_saved, optimized_tokens)
self.cost_tracker.record_tokens(
model,
tokens_saved,
optimized_tokens,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
cache_write_5m_tokens=cache_write_5m_tokens,
cache_write_1h_tokens=cache_write_1h_tokens,
)
# Log request
if self.logger:
@ -1399,7 +1431,7 @@ class StreamingMixin:
optimization_latency_ms=optimization_latency,
total_latency_ms=total_latency,
tags=tags,
cache_hit=False,
cache_hit=cache_read_tokens > 0,
transforms_applied=transforms_applied,
request_messages=body.get("messages")
if self.config.log_full_messages
@ -1417,7 +1449,8 @@ class StreamingMixin:
f"model={model} msgs={num_msgs} "
f"tok_before={original_tokens} tok_after={optimized_tokens} "
f"tok_saved={tokens_saved} "
f"cache_read=0 cache_write=0 cache_hit_pct=0 "
f"cache_read={cache_read_tokens} cache_write={cache_write_tokens} "
f"cache_hit_pct={cache_hit_pct} "
f"opt_ms={optimization_latency:.0f} "
f"transforms={_summarize_transforms(transforms_applied)}"
)
@ -1445,25 +1478,51 @@ class StreamingMixin:
) -> StreamingResponse:
"""Stream OpenAI chat completion response from backend.
Routes stream:true requests through the backend's stream_openai_message(),
yielding SSE events to the client. Tracks the final
`usage.completion_tokens` online (LiteLLM emits this only when the
request included ``stream_options.include_usage=true``) using
:func:`_parse_completion_tokens_from_sse_chunk`, so memory stays
O(1) regardless of stream length.
Routes stream:true requests through the backend's
``stream_openai_message()``, yielding SSE events to the client.
Buffers chunk bytes into ``stream_state["sse_buffer"]`` and
incrementally drains complete events via
:meth:`_parse_sse_usage_from_buffer` so the final usage frame
(LiteLLM/OpenAI emits this only when the request included
``stream_options.include_usage=true``) yields ``prompt_tokens``,
``completion_tokens``, and
``prompt_tokens_details.cached_tokens``. OpenAI exposes no
separate cache-write counter, so the write portion is inferred
via :func:`_infer_openai_cache_write_tokens`. Memory stays O(1)
because the buffer-parser consumes whole events as they arrive.
"""
from fastapi.responses import StreamingResponse
from headroom.proxy.cost import _summarize_transforms
from headroom.proxy.handlers.openai import _infer_openai_cache_write_tokens
assert self.anthropic_backend is not None
async def generate():
output_tokens = 0
stream_state: dict[str, Any] = {
"sse_buffer": bytearray(),
"input_tokens": None,
"output_tokens": None,
"cache_read_input_tokens": None,
}
def _absorb(usage: dict[str, int] | None) -> None:
if not usage:
return
for key in ("input_tokens", "output_tokens", "cache_read_input_tokens"):
if key in usage and not stream_state.get(key):
stream_state[key] = usage[key]
try:
async for sse_chunk in self.anthropic_backend.stream_openai_message(body, headers):
chunk_bytes = sse_chunk.encode() if isinstance(sse_chunk, str) else sse_chunk
stream_state["sse_buffer"].extend(chunk_bytes)
_absorb(self._parse_sse_usage_from_buffer(stream_state, "openai"))
# Per-chunk fallback for upstreams that emit only
# ``completion_tokens`` and not a full usage frame.
parsed = _parse_completion_tokens_from_sse_chunk(chunk_bytes)
if parsed is not None:
output_tokens = parsed
if parsed is not None and not stream_state["output_tokens"]:
stream_state["output_tokens"] = parsed
yield chunk_bytes
except Exception as e:
logger.error(f"[{request_id}] Backend streaming error: {e}")
@ -1477,6 +1536,39 @@ class StreamingMixin:
yield f"data: {json.dumps(error_data)}\n\n".encode()
yield b"data: [DONE]\n\n"
finally:
# Late-flush: if upstream truncated the stream mid-event,
# the buffer parser hasn't seen the closing ``\n\n`` yet.
# Mirror _finalize_stream_response: append the terminator
# and drain anything still parseable.
buf = stream_state["sse_buffer"]
if len(buf) > 0:
buf.extend(b"\n\n")
_absorb(self._parse_sse_usage_from_buffer(stream_state, "openai"))
# Mirror the non-streaming sibling (``_extract_responses_usage``
# in handlers/openai.py): only infer cache metrics when
# upstream actually reported a usage frame. Otherwise the
# proxy-side ``optimized_tokens`` would masquerade as a
# cache write — wrong, and indistinguishable from a real
# hit-rate-zero call in the dashboard.
upstream_input = stream_state["input_tokens"]
output_tokens = stream_state["output_tokens"] or 0
cache_read_tokens = stream_state["cache_read_input_tokens"] or 0
if upstream_input is None:
input_tokens = 0
cache_write_tokens = 0
uncached_input_tokens = 0
cache_hit_pct = 0
else:
input_tokens = upstream_input
cache_write_tokens = _infer_openai_cache_write_tokens(
input_tokens, cache_read_tokens
)
uncached_input_tokens = max(input_tokens - cache_read_tokens, 0)
cache_hit_pct = (
round(cache_read_tokens / input_tokens * 100) if input_tokens > 0 else 0
)
total_latency = (time.time() - start_time) * 1000
# Active-compression denominator for backend-routed
# streaming. No per-message live-zone tracking is wired
@ -1493,13 +1585,23 @@ class StreamingMixin:
output_tokens=output_tokens,
tokens_saved=tokens_saved,
latency_ms=total_latency,
cached=False,
cached=cache_read_tokens > 0,
overhead_ms=optimization_latency,
pipeline_timing=pipeline_timing,
waste_signals=waste_signals,
attempted_input_tokens=attempted_input_tokens,
)
if self.cost_tracker:
self.cost_tracker.record_tokens(
model,
tokens_saved,
optimized_tokens,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
uncached_tokens=uncached_input_tokens,
)
# Mirror the Anthropic-stream path: log to RequestLogger so
# /stats.recent_requests and /transformations/feed see this
# request. Without it the OpenAI-via-backend path is invisible.
@ -1522,7 +1624,7 @@ class StreamingMixin:
optimization_latency_ms=optimization_latency,
total_latency_ms=total_latency,
tags=tags or {},
cache_hit=False,
cache_hit=cache_read_tokens > 0,
transforms_applied=transforms_applied,
waste_signals=waste_signals,
request_messages=body.get("messages")
@ -1531,6 +1633,23 @@ class StreamingMixin:
)
)
# Structured perf log line for `headroom perf` analysis.
# Missing this line is why issue #327 reported
# ``Cache write: 0`` on Azure-GPT/Codex backend-routed
# traffic: the entire request was invisible to the perf
# parser, not zero — but the symptom is identical.
num_msgs = len(body.get("messages", []))
logger.info(
f"[{request_id}] PERF "
f"model={model} msgs={num_msgs} "
f"tok_before={original_tokens} tok_after={optimized_tokens} "
f"tok_saved={tokens_saved} "
f"cache_read={cache_read_tokens} cache_write={cache_write_tokens} "
f"cache_hit_pct={cache_hit_pct} "
f"opt_ms={optimization_latency:.0f} "
f"transforms={_summarize_transforms(transforms_applied)}"
)
if tokens_saved > 0:
logger.info(
f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} "

355
tests/test_backend_streaming_cache_metrics.py Normal file
View file

@ -0,0 +1,355 @@
"""Cache-metric coverage for backend-routed streaming.
Two regressions in main as of 2026-05-14 (issue #327):
* ``StreamingMixin._stream_openai_via_backend`` (Azure/LiteLLM/AnyLLM OpenAI
streaming) never inspects ``usage.prompt_tokens_details.cached_tokens`` from
the upstream SSE chunks. Cache reads/writes are absent from
``cost_tracker.record_tokens``, ``SavingsTracker.record_request``, the
``RequestLog``, *and* the ``PERF`` log line — the latter is missing entirely
for this path, so ``headroom perf`` shows 0 cache writes for every
Azure-GPT/Codex backend-routed request.
* ``StreamingMixin._stream_response_bedrock`` (Bedrock-native streaming) hard-
codes ``cache_read=0 cache_write=0 cache_hit_pct=0`` in its PERF log line
regardless of what ``message_start.usage`` reported.
Both surface to the user as "Cache write: 0 tokens" in ``headroom perf``.
"""
from __future__ import annotations
import json
import logging
import re
from collections.abc import AsyncIterator
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
fastapi = pytest.importorskip("fastapi")
httpx = pytest.importorskip("httpx")
from fastapi.testclient import TestClient # noqa: E402
from headroom.backends.base import StreamEvent # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
PERF_RE = re.compile(
r"\bcache_read=(?P<cr>\d+)\s+cache_write=(?P<cw>\d+)\s+cache_hit_pct=(?P<chp>\d+)"
)
def _find_perf_record(records: list[logging.LogRecord]) -> tuple[int, int, int]:
"""Find the structured PERF log line and return (cache_read, cache_write, hit_pct)."""
for record in records:
msg = record.getMessage()
if " PERF " not in msg:
continue
m = PERF_RE.search(msg)
if m:
return int(m["cr"]), int(m["cw"]), int(m["chp"])
raise AssertionError(
"No PERF log line with cache_read/cache_write/cache_hit_pct found. "
f"Captured {len(records)} records.\n" + "\n".join(r.getMessage() for r in records[-15:])
)
class _ListHandler(logging.Handler):
"""Tiny direct handler that survives the proxy disabling propagation.
``caplog`` attaches to root; ``headroom.proxy.helpers._setup_file_logging``
flips ``logging.getLogger("headroom").propagate = False`` once a proxy
instance is constructed in the test, after which root-attached handlers
stop receiving headroom-namespaced records. Attaching directly to
``headroom.proxy`` sidesteps that.
"""
def __init__(self) -> None:
super().__init__(level=logging.INFO)
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None: # noqa: D401
self.records.append(record)
def _attach_proxy_log_capture():
handler = _ListHandler()
target = logging.getLogger("headroom.proxy")
target.addHandler(handler)
prior_level = target.level
target.setLevel(logging.INFO)
return handler, target, prior_level
def _detach_proxy_log_capture(handler, target, prior_level) -> None:
target.removeHandler(handler)
target.setLevel(prior_level)
def _make_openai_backend(chunks: list[str]) -> MagicMock:
"""Build a mock backend that yields OpenAI-format SSE chunks."""
async def fake_stream(body: dict, headers: dict) -> AsyncIterator[str]:
for chunk in chunks:
yield chunk
mock = MagicMock()
mock.name = "anyllm-openai"
mock.stream_openai_message = fake_stream
return mock
def _make_bedrock_backend(events: list[StreamEvent]) -> MagicMock:
"""Build a mock backend that yields Anthropic StreamEvent objects."""
async def fake_stream(body: dict, headers: dict) -> AsyncIterator[StreamEvent]:
for evt in events:
yield evt
mock = MagicMock()
mock.name = "bedrock"
mock.stream_message = fake_stream
mock.map_model_id = MagicMock(return_value="claude-3-5-sonnet-20241022")
mock.supports_model = MagicMock(return_value=True)
return mock
# =============================================================================
# Bug A — _stream_openai_via_backend (Azure/LiteLLM/AnyLLM OpenAI streaming)
# =============================================================================
def test_openai_backend_streaming_emits_perf_with_cache_read_and_inferred_write() -> None:
"""OpenAI backend streaming must surface cache reads + inferred writes.
Real upstream (OpenAI Chat Completions w/ ``stream_options.include_usage=true``,
or Azure GPT-5.5 through LiteLLM) emits a final chunk carrying::
usage: {
prompt_tokens: 1000,
completion_tokens: 50,
prompt_tokens_details: { cached_tokens: 700 }
}
OpenAI never reports a separate write counter, so we infer it as
``max(prompt_tokens - cached_tokens, 0)`` (see
``_infer_openai_cache_write_tokens``). The PERF log line consumed by
``headroom perf`` must report both.
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="openai",
)
chunks = [
'data: {"id":"c1","object":"chat.completion.chunk",'
'"choices":[{"index":0,"delta":{"role":"assistant","content":"hi"}}]}\n\n',
'data: {"id":"c1","object":"chat.completion.chunk",'
'"choices":[{"index":0,"delta":{"content":" there"},"finish_reason":"stop"}]}\n\n',
'data: {"id":"c1","object":"chat.completion.chunk","choices":[],'
'"usage":{"prompt_tokens":1000,"completion_tokens":50,"total_tokens":1050,'
'"prompt_tokens_details":{"cached_tokens":700}}}\n\n',
"data: [DONE]\n\n",
]
backend = _make_openai_backend(chunks)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/chat/completions",
json={
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "hi"}],
"stream": True,
"stream_options": {"include_usage": True},
},
headers={"Authorization": "Bearer test-key"},
)
assert resp.status_code == 200, resp.text[:200]
body = resp.text
assert "[DONE]" in body, body[:300]
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert cr == 700, f"expected cache_read=700, got {cr}"
assert cw == 300, f"expected inferred cache_write=300 (=1000-700), got {cw}"
assert chp == 70, f"expected cache_hit_pct=70, got {chp}"
def test_openai_backend_streaming_perf_zeros_when_upstream_omits_usage() -> None:
"""When the upstream omits a usage chunk, cache values must be zero — not absent.
Without ``stream_options.include_usage=true`` (or when upstream drops the
final usage chunk) the PERF line still has to emit so ``headroom perf``
counts the request.
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="openai",
)
chunks = [
'data: {"id":"c1","object":"chat.completion.chunk",'
'"choices":[{"index":0,"delta":{"role":"assistant","content":"hi"}}]}\n\n',
"data: [DONE]\n\n",
]
backend = _make_openai_backend(chunks)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/chat/completions",
json={
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "hi"}],
"stream": True,
},
headers={"Authorization": "Bearer test-key"},
)
assert resp.status_code == 200
assert "[DONE]" in resp.text
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert (cr, cw, chp) == (0, 0, 0)
# =============================================================================
# Bug B — _stream_response_bedrock (Bedrock-native Anthropic streaming)
# =============================================================================
def _sse_data(event_type: str, data: dict[str, Any]) -> str:
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
def test_bedrock_streaming_emits_perf_with_message_start_cache_usage() -> None:
"""Bedrock streaming must surface cache_read + cache_write from message_start.
Anthropic streaming reports cache usage on ``message_start.message.usage``
(cache_read_input_tokens + cache_creation_input_tokens). The Bedrock streamer
currently captures only ``input_tokens`` and ``output_tokens`` from the same
event and hardcodes ``cache_read=0 cache_write=0`` into the PERF log.
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="anthropic",
)
message_start = {
"type": "message_start",
"message": {
"id": "msg_1",
"model": "claude-3-5-sonnet-20241022",
"role": "assistant",
"type": "message",
"content": [],
"usage": {
"input_tokens": 1000,
"cache_read_input_tokens": 500,
"cache_creation_input_tokens": 200,
},
},
}
block_start = {
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
}
block_delta = {
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "hi"},
}
block_stop = {"type": "content_block_stop", "index": 0}
message_delta = {
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
}
message_stop = {"type": "message_stop"}
events = [
StreamEvent(
event_type=e["type"],
data=e,
raw_sse=_sse_data(e["type"], e),
)
for e in [
message_start,
block_start,
block_delta,
block_stop,
message_delta,
message_stop,
]
]
backend = _make_bedrock_backend(events)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/messages",
json={
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 64,
"stream": True,
},
headers={
"x-api-key": "sk-ant-test",
"anthropic-version": "2023-06-01",
},
)
assert resp.status_code == 200, resp.text[:200]
assert "message_stop" in resp.text
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert cr == 500, f"expected cache_read=500, got {cr}"
assert cw == 200, f"expected cache_write=200, got {cw}"
# round(500 / (500 + 200) * 100) = round(71.43) = 71
assert chp == 71, f"expected cache_hit_pct=71, got {chp}"
# =============================================================================
# Regression guard
# =============================================================================
def test_streaming_perf_log_has_no_hardcoded_cache_zeros() -> None:
"""Catch any future re-introduction of ``cache_read=0 cache_write=0`` literal."""
from pathlib import Path
src = Path(__file__).resolve().parents[1] / "headroom" / "proxy" / "handlers" / "streaming.py"
text = src.read_text()
assert "cache_read=0 cache_write=0" not in text, (
"streaming.py contains a hardcoded `cache_read=0 cache_write=0` PERF log fragment. "
"Wire the real cache_read/cache_write values into the PERF line instead."
)