diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 2c60548a0..fc0cacd30 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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:,} " diff --git a/tests/test_backend_streaming_cache_metrics.py b/tests/test_backend_streaming_cache_metrics.py new file mode 100644 index 000000000..f9a31d107 --- /dev/null +++ b/tests/test_backend_streaming_cache_metrics.py @@ -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\d+)\s+cache_write=(?P\d+)\s+cache_hit_pct=(?P\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." + )