From d91254a0f2419dbd77e6b8f3b6718a029b725082 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 8 May 2026 17:23:10 -0700 Subject: [PATCH] fix(proxy): recover SSE usage from truncated final event MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-chunk SSE parser only flushes events terminated by `\n\n`. When upstream truncates mid-event (client disconnect, network drop, RemoteProtocolError), the message_start (cache_read / cache_creation) or message_delta (output_tokens) usage events sit in the residual sse_buffer and never get parsed — the finalizer then logs cache_read=cache_write=0, which the freeze heuristic on the next request reads as "no provider cache, reprocess everything," producing a different prefix and a real cache miss on the *next* turn. Append `\n\n` to the residual buffer at end-of-stream so the existing parser drains the partial event. Only fills None / 0 slots so a real cache_read=0 from earlier in the stream isn't clobbered. --- headroom/proxy/handlers/streaming.py | 32 ++++++++++++ tests/test_proxy_streaming_request_logger.py | 51 ++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index f3b4f4e45..c70404e7a 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -561,6 +561,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 bc36dd6e3..c1f769fc2 100644 --- a/tests/test_proxy_streaming_request_logger.py +++ b/tests/test_proxy_streaming_request_logger.py @@ -153,6 +153,57 @@ async def test_finalize_stream_response_handles_zero_original_tokens(): assert entries[0]["savings_percent"] == 0 +@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)