diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 5162cdd88..325b3e502 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -845,13 +845,18 @@ class StreamingMixin: stream_state[key] = late_usage[key] 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" ) + outcome_tags = dict(tags or {}) + outcome_tags["output_tokens_source"] = output_tokens_source + provider_input_tokens = stream_state.get("input_tokens") effective_optimized_tokens = optimized_tokens effective_original_tokens = original_tokens @@ -940,7 +945,7 @@ class StreamingMixin: transforms_applied=transforms_applied, total_latency_ms=total_latency, overhead_ms=optimization_latency, - tags=tags, + tags=outcome_tags, client=client, log_full_messages=getattr(self.config, "log_full_messages", False), cache_read_tokens=cache_read_tokens, diff --git a/tests/test_proxy_streaming_request_logger.py b/tests/test_proxy_streaming_request_logger.py index 1aed8bf38..f39aa4b87 100644 --- a/tests/test_proxy_streaming_request_logger.py +++ b/tests/test_proxy_streaming_request_logger.py @@ -82,6 +82,7 @@ def test_parse_openai_responses_completed_usage_from_sse_buffer(): @pytest.mark.asyncio async def test_finalize_stream_response_logs_request_for_feed(): proxy = _build_proxy_with_real_logger(log_full_messages=False) + request_tags = {"stack": "wrap_claude"} await proxy._finalize_stream_response( body={"messages": [{"role": "user", "content": "hi"}]}, @@ -95,7 +96,7 @@ async def test_finalize_stream_response_logs_request_for_feed(): optimization_latency=12.0, stream_state=_stream_state(), start_time=0.0, - tags={"stack": "wrap_claude"}, + tags=request_tags, ) entries = proxy.logger.get_recent(10) @@ -109,10 +110,40 @@ async def test_finalize_stream_response_logs_request_for_feed(): assert entry["tokens_saved"] == 400 assert entry["savings_percent"] == pytest.approx(40.0) assert entry["transforms_applied"] == ["smart_crusher"] - assert entry["tags"] == {"stack": "wrap_claude"} + assert entry["tags"] == { + "stack": "wrap_claude", + "output_tokens_source": "provider", + } + assert request_tags == {"stack": "wrap_claude"} assert entry["cache_hit"] is False +@pytest.mark.asyncio +async def test_finalize_stream_response_marks_estimated_output_tokens() -> None: + proxy = _build_proxy_with_real_logger(log_full_messages=False) + state = _stream_state() + state["output_tokens"] = None + state["total_bytes"] = 200 + + await proxy._finalize_stream_response( + body={"messages": [{"role": "user", "content": "hi"}]}, + provider="anthropic", + model="claude-sonnet-4-6", + request_id="req-stream-estimated", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + optimization_latency=1.0, + stream_state=state, + start_time=0.0, + ) + + entry = proxy.logger.get_recent(1)[0] + assert entry["output_tokens"] == 5 + assert entry["tags"]["output_tokens_source"] == "estimated_bytes" + + @pytest.mark.asyncio async def test_finalize_stream_response_logs_original_and_compressed_messages(): """With log_full_messages enabled, both sides of the compression are