From 1c9585d42ee4167def21cf5e823cdb7ee329c33f Mon Sep 17 00:00:00 2001 From: Chester Date: Thu, 16 Jul 2026 02:17:24 +0800 Subject: [PATCH] fix(stats): tag streamed output token source (#2214) ## Description Preserve the existing SSE output-token fallback while making its provenance visible to request logs and downstream statistics. Closes #2213 ## 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 - Tag provider-reported streaming output tokens with `output_tokens_source=provider`. - Tag the existing `total_bytes // 40` fallback with `output_tokens_source=estimated_bytes`. - Copy incoming tags before adding provenance so caller-owned dictionaries are not mutated. - Add focused coverage for both source values and the unchanged fallback estimate. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uv run --with pytest --with pytest-asyncio --with fastapi --with httpx --with numpy pytest \ tests/test_proxy_streaming_request_logger.py \ tests/test_request_outcome.py \ tests/test_proxy_handler_helpers.py -q 77 passed $ uv run --with ruff ruff check headroom/proxy/handlers/streaming.py tests/test_proxy_streaming_request_logger.py All checks passed! $ uv run --with ruff ruff format --check headroom/proxy/handlers/streaming.py tests/test_proxy_streaming_request_logger.py 2 files already formatted ``` ## Real Behavior Proof - Environment: Python 3.13 with the real request logger and synthetic stream state - Exact command / steps: run the focused test set above - Observed result: parsed usage records `provider`; a 200-byte no-usage stream still records 5 output tokens and tags it `estimated_bytes` - Not tested: live provider stream, full repository suite ## Review Readiness - [x] I have performed a self-review - [ ] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes This PR does not change the fallback formula or token totals. Documentation and changelog changes are not needed for the new internal outcome tag. Co-authored-by: JerrettDavis --- headroom/proxy/handlers/streaming.py | 7 +++- tests/test_proxy_streaming_request_logger.py | 35 ++++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) 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