diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 060c18aaa..0b875d923 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -588,6 +588,38 @@ class StreamingMixin: from headroom.proxy.cost import _summarize_transforms total_latency = (time.time() - start_time) * 1000 + + # Per-chunk SSE parsing only flushes events terminated by ``\n\n``. + # When upstream truncates mid-event (client disconnect, network + # drop, connection reset), the message_start (cache_read / + # cache_creation) or message_delta (output_tokens) usage events + # can sit in the residual buffer and never be parsed — surfacing + # as cache_read=cache_write=0 in PERF logs and poisoning the + # downstream freeze heuristic for the next request. Append the + # terminator so the buffer parser drains whatever's there. The + # per-event try/except in the parser swallows incomplete JSON, + # so this is safe even when the truncation cut mid-payload. + sse_buffer = stream_state.get("sse_buffer") + if isinstance(sse_buffer, bytearray) and len(sse_buffer) > 0: + sse_buffer.extend(b"\n\n") + late_usage = self._parse_sse_usage_from_buffer(stream_state, provider) or {} + for key in ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "cache_creation_ephemeral_5m_input_tokens", + "cache_creation_ephemeral_1h_input_tokens", + ): + if key not in late_usage: + continue + current = stream_state.get(key) + # Only fill in unset (None) or default-zero slots so a + # real cache_read=0 from earlier in the stream isn't + # clobbered by a later partial event. + if current is None or current == 0: + stream_state[key] = late_usage[key] + output_tokens = stream_state["output_tokens"] if output_tokens is None: output_tokens = stream_state["total_bytes"] // 40 diff --git a/tests/test_proxy_streaming_request_logger.py b/tests/test_proxy_streaming_request_logger.py index aad6d50c0..035d9570c 100644 --- a/tests/test_proxy_streaming_request_logger.py +++ b/tests/test_proxy_streaming_request_logger.py @@ -228,6 +228,57 @@ async def test_finalize_openai_responses_stream_uses_provider_usage_for_dashboar assert cost_kwargs["uncached_tokens"] == 186_600 +@pytest.mark.asyncio +async def test_finalize_stream_response_recovers_usage_from_truncated_buffer() -> None: + """When upstream truncates mid-event (no trailing \\n\\n), the per-chunk + parser leaves the message_start usage event sitting in sse_buffer and + PERF logs cache_read=cache_write=0 — which then poisons the freeze + heuristic on the next request. The finalizer must flush the residual + buffer so the real cache_read / cache_creation tokens still land in + the log even on aborted streams. + """ + proxy = _build_proxy_with_real_logger(log_full_messages=False) + + partial_message_start = ( + b"event: message_start\n" + b'data: {"type":"message_start","message":{"id":"msg_x",' + b'"type":"message","role":"assistant","model":"claude-sonnet-4-6",' + b'"content":[],"stop_reason":null,"usage":{' + b'"input_tokens":1234,"cache_read_input_tokens":50000,' + b'"cache_creation_input_tokens":2500,"output_tokens":1}}}' + ) + + state = { + "output_tokens": None, + "total_bytes": len(partial_message_start), + "ttfb_ms": 35.0, + "input_tokens": 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, + "sse_buffer": bytearray(partial_message_start), + } + + await proxy._finalize_stream_response( + body={"messages": [{"role": "user", "content": "hi"}]}, + provider="anthropic", + model="claude-sonnet-4-6", + request_id="req-stream-truncated", + original_tokens=2000, + optimized_tokens=1800, + tokens_saved=200, + transforms_applied=[], + optimization_latency=5.0, + stream_state=state, + start_time=0.0, + ) + + assert state["input_tokens"] == 1234 + assert state["cache_read_input_tokens"] == 50000 + assert state["cache_creation_input_tokens"] == 2500 + + @pytest.mark.asyncio async def test_finalize_stream_response_no_op_when_logger_disabled(): proxy = _build_proxy_with_real_logger(log_full_messages=False)