diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index b3fc68547..3b864f9dd 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -565,8 +565,17 @@ class StreamingMixin: } events.append(f"event: message_start\ndata: {json.dumps(msg_start)}\n\n".encode()) - # Content blocks - for idx, block in enumerate(response.get("content", [])): + # Content blocks. `content` is provider/reconstruction-controlled, so a + # present-but-null value or a non-list would crash `enumerate`, and a + # non-dict element would crash `block.get(...)`. Guard both, matching the + # sibling `_record_ccr_feedback_from_response` below. This is reached from + # a call site (anthropic.py buffered CCR path) that only catches + # ValueError, so an unguarded TypeError/AttributeError would 500 the + # streamed request. + content = response.get("content") + for idx, block in enumerate(content if isinstance(content, list) else []): + if not isinstance(block, dict): + continue # content_block_start if block.get("type") == "text": block_start = { @@ -680,10 +689,13 @@ class StreamingMixin: msg_delta_payload["stop_reason"] = response["stop_reason"] if "stop_details" in response: msg_delta_payload["stop_details"] = response["stop_details"] + usage = response.get("usage") + if not isinstance(usage, dict): + usage = {} msg_delta = { "type": "message_delta", "delta": msg_delta_payload, - "usage": {"output_tokens": response.get("usage", {}).get("output_tokens", 0)}, + "usage": {"output_tokens": usage.get("output_tokens", 0)}, } events.append(f"event: message_delta\ndata: {json.dumps(msg_delta)}\n\n".encode()) diff --git a/tests/test_sse_thinking_blocks.py b/tests/test_sse_thinking_blocks.py index 0b32881f8..236b43c23 100644 --- a/tests/test_sse_thinking_blocks.py +++ b/tests/test_sse_thinking_blocks.py @@ -324,6 +324,40 @@ def test_response_to_sse_emits_unknown_content_block_verbatim() -> None: assert not any(ev["type"] == "content_block_delta" for ev in events) +def test_response_to_sse_tolerates_malformed_content_and_usage() -> None: + # `_response_to_sse` runs on provider/reconstruction-controlled JSON and is + # reached from a call site (anthropic.py buffered CCR path) that only catches + # ValueError, so a present-but-null `content`/`usage` or a non-dict block must + # not raise a TypeError/AttributeError that would 500 the streamed request. + # The sibling `_record_ccr_feedback_from_response` guards `content` the same + # way. Each of these once crashed the unguarded loop. + parser = _Parser() + + for response in ( + {"content": None, "usage": {"output_tokens": 5}}, + {"content": "not-a-list"}, + {"content": [None, {"type": "text", "text": "hi"}]}, + {"content": [], "usage": None}, + ): + sse_text = b"".join(parser._response_to_sse(response, "anthropic")).decode("utf-8") + # Always a well-formed envelope, regardless of the malformed body. + assert "event: message_start" in sse_text + assert "event: message_stop" in sse_text + + # The one valid block alongside a null element is still rendered. + sse_text = b"".join( + parser._response_to_sse({"content": [None, {"type": "text", "text": "hi"}]}, "anthropic") + ).decode("utf-8") + events = _sse_events(sse_text) + text_deltas = [ + ev + for ev in events + if ev["type"] == "content_block_delta" and ev["delta"].get("type") == "text_delta" + ] + assert len(text_deltas) == 1 + assert text_deltas[0]["delta"]["text"] == "hi" + + def test_response_to_sse_emits_server_tool_use_without_delta() -> None: parser = _Parser() server_tool_use = {