diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index e265ba80c..4e0020e39 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -31,6 +31,49 @@ from headroom.copilot_auth import apply_copilot_api_auth logger = logging.getLogger("headroom.proxy") +# Upstream controls every SSE event body, so the shape of anything reached +# through `json.loads` is not guaranteed. These normalize a wrong-shaped value +# to a harmless one instead of letting it raise on the streaming path, where +# the raise escapes into `_finalize_stream_response` and tears down a stream +# the client is already reading. They guard *shape*, not value validity. + + +def _sse_dict(value: Any) -> dict[str, Any]: + """Upstream-supplied object, or ``{}`` when it is not one. + + An explicit ``"key": null`` bypasses a ``.get(key, {})`` default — the + default only applies to a *missing* key — so nested objects are + type-checked rather than defaulted. + """ + return value if isinstance(value, dict) else {} + + +def _sse_str(value: Any) -> str: + """Upstream-supplied text, or ``""`` when it is not a string. + + Accumulating a non-string delta would raise ``TypeError`` on ``+``. + """ + return value if isinstance(value, str) else "" + + +def _sse_index(value: Any, default: int | None = None) -> int | None: + """Upstream-supplied block index, or ``default`` when unusable as a key. + + ``index`` keys ``blocks_by_index`` and ``appended_block_keys``; a list or + dict there raises ``TypeError: unhashable type``. ``bool`` is rejected too + so ``true`` cannot silently alias index ``1``. + """ + if isinstance(value, bool) or not isinstance(value, int): + return default + return value + + +# Keys `_parse_sse_to_response` owns on its in-progress block dicts. Upstream +# must not seed them via the non-standard-block copy-through, or a non-string +# would reach the delta accumulators. `type` is set explicitly, not copied. +_BLOCK_SCRATCH_KEYS = frozenset({"type", "_partial_json", "thinking_buffer"}) + + def _parse_completion_tokens_from_sse_chunk(chunk_bytes: bytes) -> int | None: """Extract `usage.completion_tokens` from a single SSE chunk if present. @@ -180,6 +223,9 @@ class StreamingMixin: except json.JSONDecodeError: continue + if not isinstance(data, dict): + continue + usage = {} if provider == "anthropic": @@ -188,8 +234,7 @@ class StreamingMixin: event_type = data.get("type", "") if event_type == "message_start": - msg = data.get("message", {}) - msg_usage = msg.get("usage", {}) + msg_usage = _sse_dict(_sse_dict(data.get("message")).get("usage")) if msg_usage: usage["input_tokens"] = msg_usage.get("input_tokens", 0) usage["cache_read_input_tokens"] = msg_usage.get( @@ -205,24 +250,24 @@ class StreamingMixin: usage["cache_creation_ephemeral_1h_input_tokens"] = cache_write_1h elif event_type == "message_delta": - delta_usage = data.get("usage", {}) + delta_usage = _sse_dict(data.get("usage")) if delta_usage: usage["output_tokens"] = delta_usage.get("output_tokens", 0) elif provider == "openai": # OpenAI sends usage in final chunk (when stream_options.include_usage=true) - chunk_usage = data.get("usage") + chunk_usage = _sse_dict(data.get("usage")) if chunk_usage: usage["input_tokens"] = chunk_usage.get("prompt_tokens", 0) usage["output_tokens"] = chunk_usage.get("completion_tokens", 0) # OpenAI has cached tokens in prompt_tokens_details - details = chunk_usage.get("prompt_tokens_details") or {} + details = _sse_dict(chunk_usage.get("prompt_tokens_details")) usage["cache_read_input_tokens"] = details.get("cached_tokens", 0) elif provider == "gemini": # Gemini sends usageMetadata in each streaming chunk # Format: {"usageMetadata": {"promptTokenCount": N, "candidatesTokenCount": M}} - usage_meta = data.get("usageMetadata") + usage_meta = _sse_dict(data.get("usageMetadata")) if usage_meta: usage["input_tokens"] = usage_meta.get("promptTokenCount", 0) usage["output_tokens"] = gemini_output_tokens(usage_meta) @@ -276,11 +321,18 @@ class StreamingMixin: except json.JSONDecodeError: continue + # The upstream controls the event body. A valid-JSON non-object + # (a bare array or string, e.g. from an overloaded frontend's + # error page framed as SSE) makes `.get` raise AttributeError, + # which the JSONDecodeError guard above does not catch. Skip it + # like any other unusable event. + if not isinstance(data, dict): + continue + if provider == "anthropic": event_type = data.get("type", "") if event_type == "message_start": - msg = data.get("message", {}) - msg_usage = msg.get("usage", {}) + msg_usage = _sse_dict(_sse_dict(data.get("message")).get("usage")) if msg_usage: usage_found["input_tokens"] = msg_usage.get("input_tokens", 0) usage_found["cache_read_input_tokens"] = msg_usage.get( @@ -300,7 +352,7 @@ class StreamingMixin: f"cache_write={usage_found.get('cache_creation_input_tokens')}" ) elif event_type == "message_delta": - delta_usage = data.get("usage", {}) + delta_usage = _sse_dict(data.get("usage")) if delta_usage: usage_found["output_tokens"] = delta_usage.get("output_tokens", 0) @@ -340,7 +392,7 @@ class StreamingMixin: ) elif provider == "gemini": - usage_meta = data.get("usageMetadata") + usage_meta = _sse_dict(data.get("usageMetadata")) if usage_meta: usage_found["input_tokens"] = usage_meta.get("promptTokenCount", 0) usage_found["output_tokens"] = gemini_output_tokens(usage_meta) @@ -403,29 +455,36 @@ class StreamingMixin: except json.JSONDecodeError: continue + # Same upstream-controlled shape guard as + # ``_parse_sse_usage_from_buffer``: a valid-JSON non-object event + # would raise AttributeError below, and here that escapes into + # ``_finalize_stream_response`` and tears down a stream the client + # was already reading. + if not isinstance(data, dict): + continue + event_type = data.get("type", "") if event_type == "message_start": - msg = data.get("message", {}) + msg = _sse_dict(data.get("message")) response["id"] = msg.get("id") response["model"] = msg.get("model") response["role"] = msg.get("role", "assistant") response["stop_reason"] = msg.get("stop_reason") if "stop_details" in msg: response["stop_details"] = msg["stop_details"] - if msg.get("usage"): - response["usage"].update(msg["usage"]) + response["usage"].update(_sse_dict(msg.get("usage"))) elif event_type == "content_block_start": - block = data.get("content_block", {}) - block_index = data.get("index", len(response["content"])) + block = _sse_dict(data.get("content_block")) + block_index = _sse_index(data.get("index"), len(response["content"])) btype = block.get("type") current_block = { "type": btype, "index": block_index, } if btype == "text": - current_block["text"] = block.get("text", "") + current_block["text"] = _sse_str(block.get("text")) elif btype == "tool_use": current_block["id"] = block.get("id") current_block["name"] = block.get("name") @@ -434,7 +493,7 @@ class StreamingMixin: # Thinking block — accumulate text via # `thinking_delta`; signature arrives via # `signature_delta` (single value, not accumulated). - current_block["thinking_buffer"] = block.get("thinking", "") + current_block["thinking_buffer"] = _sse_str(block.get("thinking")) if "signature" in block: current_block["signature"] = block["signature"] elif btype == "redacted_thinking": @@ -451,31 +510,33 @@ class StreamingMixin: # {type, index}. Mirrors the sibling reconstructor # `_reconstruct_anthropic_response`, which does `dict(block)`. for _k, _v in block.items(): - if _k != "type": + if _k not in _BLOCK_SCRATCH_KEYS: current_block[_k] = _v blocks_by_index[block_index] = current_block elif event_type == "content_block_delta": # Resolve the target block by index (preferred) or fall # back to current_block for legacy linear streams. - idx = data.get("index") + idx = _sse_index(data.get("index")) target = (blocks_by_index.get(idx) if idx is not None else None) or current_block if target is not None: - delta = data.get("delta", {}) + # `"delta": null` bypasses a `.get(..., {})` default the same + # way `"message": null` does on message_start. + delta = _sse_dict(data.get("delta")) dtype = delta.get("type") if dtype == "text_delta": - target["text"] = target.get("text", "") + delta.get("text", "") + target["text"] = _sse_str(target.get("text")) + _sse_str(delta.get("text")) elif dtype == "input_json_delta": # Accumulate partial JSON for tool input. - partial = delta.get("partial_json", "") - target["_partial_json"] = target.get("_partial_json", "") + partial + partial = _sse_str(delta.get("partial_json")) + target["_partial_json"] = _sse_str(target.get("_partial_json")) + partial elif dtype == "thinking_delta": # Accumulate thinking text into the dedicated # buffer so it never collides with `text` on # text blocks (separate field per guide §2.7). - target["thinking_buffer"] = target.get("thinking_buffer", "") + delta.get( - "thinking", "" - ) + target["thinking_buffer"] = _sse_str( + target.get("thinking_buffer") + ) + _sse_str(delta.get("thinking")) elif dtype == "signature_delta": # Single value, not accumulated. Last-write # wins per Anthropic spec. @@ -486,13 +547,18 @@ class StreamingMixin: # list so multi-citation blocks reconstruct # correctly. Per guide §2.5: each delta carries # one full citation object under `citation`. - citations = target.setdefault("citations", []) + # A non-standard block start can copy a non-list + # `citations` through, which `.append` would reject. + citations = target.get("citations") + if not isinstance(citations, list): + citations = [] + target["citations"] = citations citation = delta.get("citation") if citation is not None: citations.append(citation) elif event_type == "content_block_stop": - idx = data.get("index") + idx = _sse_index(data.get("index")) target = (blocks_by_index.get(idx) if idx is not None else None) or current_block if target is not None: # Parse accumulated JSON into `input` for any block that @@ -526,13 +592,15 @@ class StreamingMixin: current_block = None elif event_type == "message_delta": - delta = data.get("delta", {}) + delta = _sse_dict(data.get("delta")) if "stop_reason" in delta: response["stop_reason"] = delta["stop_reason"] if "stop_details" in delta: response["stop_details"] = delta["stop_details"] - if data.get("usage"): - response["usage"].update(data["usage"]) + # Type-checked, not truthiness-checked: `dict.update` on a + # non-mapping raises. The message_start twin above already + # guarded this; this sibling did not. + response["usage"].update(_sse_dict(data.get("usage"))) return response if response.get("content") else None diff --git a/tests/test_streaming_sse_malformed_events.py b/tests/test_streaming_sse_malformed_events.py new file mode 100644 index 000000000..7cdf28ddc --- /dev/null +++ b/tests/test_streaming_sse_malformed_events.py @@ -0,0 +1,611 @@ +"""Malformed-upstream tolerance for SSE event parsing. + +Three parsers read upstream-controlled SSE bodies: ``_parse_sse_usage``, +``_parse_sse_usage_from_buffer`` and ``_parse_sse_to_response``. A valid-JSON +event of unexpected *shape* must be skipped rather than raise. In +``_parse_sse_to_response`` the raise escapes into ``_finalize_stream_response`` +and tears down a stream the client is already reading. + +Scope: these cover **frame shape** — the type of a container the parser reaches +into. Value validity (a token count that is a string, ``Infinity``, or a +>4300-digit integer literal) is a separate defect class and is not covered here. + +The happy-path fixtures are transcribed from real ``api.anthropic.com`` +``/v1/messages`` streams (``claude-haiku-4-5``), not hand-invented. Details a +stand-in gets wrong, which these pin: + +* ``message_start.message`` sends ``stop_reason``, ``stop_sequence`` and + ``stop_details`` as **null**. A guard that rejected nulls indiscriminately + would break every ordinary request. +* ``message_delta.usage`` repeats the *full* usage block, not just + ``output_tokens``. +* ``usage.cache_creation`` is always present as a nested object — it is what + ``_extract_anthropic_cache_ttl_metrics`` reads. +* A ``tool_use`` block carries a non-standard ``caller`` field, which is what + exercises the copy-through branch in ``content_block_start``. +* Extended thinking emits ``thinking_delta`` then ``signature_delta`` on block + 0, then a separate text block at index 1. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from headroom.proxy.handlers.streaming import StreamingMixin + + +class _Handler(StreamingMixin): + """Bare mixin host — the parsers touch no other handler state.""" + + +# `None` is itself a value under test, so it cannot double as "not supplied". +_UNSET = object() + +_WRONG_TYPES = [None, "str", 42, 3.5, True, ["x"], [], {}] + + +def _frame(event: str, payload: Any) -> str: + """One SSE frame in Anthropic's wire shape: an ``event:`` line then ``data:``.""" + return f"event: {event}\ndata: {json.dumps(payload)}\n\n" + + +# --------------------------------------------------------------------------- +# Fixtures transcribed from real api.anthropic.com streams. +# --------------------------------------------------------------------------- + +_REAL_USAGE = { + "input_tokens": 12, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "not_available", +} + +_REAL_MESSAGE = { + "id": "msg_011CdQD1YDscyCRWpGCcrwkA", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5-20251001", + "content": [], + # Real Anthropic sends all three as null on message_start. + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": _REAL_USAGE, +} + +_MESSAGE_DELTA = _frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None}, + # The real message_delta repeats the whole usage block. + "usage": { + "input_tokens": 12, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 4, + }, + }, +) + + +def _text_stream(*, message: Any = _UNSET, extra: str = "", mid: str = "") -> str: + """The real single-text-block stream, optionally with frames injected. + + ``extra`` lands right after ``message_start``; ``mid`` lands *inside* the + open content block, after ``content_block_start``. The distinction matters: + a ``content_block_delta`` injected before any block has opened resolves to + ``target is None`` and the parser skips its whole body, so a delta-shaped + fault placed in ``extra`` would never reach the code it is meant to test. + """ + msg = _REAL_MESSAGE if message is _UNSET else message + return ( + _frame("message_start", {"type": "message_start", "message": msg}) + + extra + + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + + _frame("ping", {"type": "ping"}) + + mid + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hi"}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}) + + _MESSAGE_DELTA + + _frame("message_stop", {"type": "message_stop"}) + ) + + +def _tool_use_stream() -> str: + """The real tool_use stream — note the non-standard ``caller`` field.""" + partials = ["", '{"ci', "ty", '": "Par', 'is"}'] + return ( + _frame("message_start", {"type": "message_start", "message": _REAL_MESSAGE}) + + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_01Kpi77sYFPtGoQbz6YMCSok", + "name": "get_weather", + "input": {}, + "caller": {"type": "direct"}, + }, + }, + ) + + "".join( + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": p}, + }, + ) + for p in partials + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}) + + _frame("message_stop", {"type": "message_stop"}) + ) + + +def _thinking_stream() -> str: + """The real extended-thinking stream: thinking block 0, then text block 1.""" + return ( + _frame("message_start", {"type": "message_start", "message": _REAL_MESSAGE}) + + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": "17"}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": " * 23 = 391"}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "signature_delta", "signature": "EtQCCpMBCBAYAipA"}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}) + + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "text_delta", "text": "391"}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 1}) + + _MESSAGE_DELTA + + _frame("message_stop", {"type": "message_stop"}) + ) + + +def _usage(payload: str, provider: str = "anthropic") -> dict[str, int] | None: + """Run the buffered usage parser over a whole stream.""" + return _Handler()._parse_sse_usage_from_buffer( + {"sse_buffer": bytearray(payload.encode())}, provider + ) + + +def _usage_chunk(payload: str, provider: str = "anthropic") -> dict[str, int] | None: + """Run the single-chunk usage parser (the third, previously unguarded one).""" + return _Handler()._parse_sse_usage(payload.encode(), provider) + + +def _response(payload: str, provider: str = "anthropic") -> dict[str, Any] | None: + return _Handler()._parse_sse_to_response(payload, provider) + + +# --------------------------------------------------------------------------- +# Happy path — the real streams must parse, nulls and ping frames included. +# --------------------------------------------------------------------------- + + +def test_real_text_stream_usage_is_parsed() -> None: + usage = _usage(_text_stream()) + assert usage is not None + assert usage["input_tokens"] == 12 + assert usage["cache_read_input_tokens"] == 0 + # message_delta's output_tokens wins over message_start's provisional 1. + assert usage["output_tokens"] == 4 + # Sourced from the nested usage.cache_creation object real traffic always sends. + assert usage["cache_creation_ephemeral_5m_input_tokens"] == 0 + assert usage["cache_creation_ephemeral_1h_input_tokens"] == 0 + + +def test_real_text_stream_reconstructs_the_response() -> None: + resp = _response(_text_stream()) + assert resp is not None + assert resp["id"] == "msg_011CdQD1YDscyCRWpGCcrwkA" + assert resp["model"] == "claude-haiku-4-5-20251001" + assert resp["role"] == "assistant" + assert resp["stop_reason"] == "end_turn" + assert [b["text"] for b in resp["content"]] == ["hi"] + assert resp["usage"]["output_tokens"] == 4 + + +def test_real_tool_use_stream_reconstructs_tool_input() -> None: + resp = _response(_tool_use_stream()) + assert resp is not None + block = resp["content"][0] + assert block["type"] == "tool_use" + assert block["name"] == "get_weather" + # The five partial_json deltas must accumulate and parse. + assert block["input"] == {"city": "Paris"} + assert "_partial_json" not in block + # Pins current behaviour, not desired behaviour: the copy-through branch + # runs only for *non-standard* block types, so `tool_use` keeps just + # id/name/input and the real `caller` field is dropped on reconstruction. + # Out of scope here (this is fidelity, not a shape crash); see the + # `server_tool_use` cases below for the copy-through path itself. + assert "caller" not in block + + +def test_real_thinking_stream_reconstructs_both_blocks() -> None: + resp = _response(_thinking_stream()) + assert resp is not None + thinking, text = resp["content"] + assert thinking["type"] == "thinking" + assert thinking["thinking"] == "17 * 23 = 391" + assert thinking["signature"] == "EtQCCpMBCBAYAipA" + assert "thinking_buffer" not in thinking + assert text["text"] == "391" + + +def test_legitimate_nulls_are_not_treated_as_malformed() -> None: + # stop_reason/stop_sequence/stop_details are null on every real + # message_start. The guard must not mistake them for a malformed frame. + head = ( + _frame("message_start", {"type": "message_start", "message": _REAL_MESSAGE}) + + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + partial = _response(head) + assert partial is not None + assert partial["id"] == "msg_011CdQD1YDscyCRWpGCcrwkA" + assert partial["stop_reason"] is None + assert partial["stop_details"] is None + + +# --------------------------------------------------------------------------- +# Shape sweep. Each entry injects a wrong-typed value at one position the +# parser reaches into, inside an otherwise-real stream. The frame must be +# skipped, never fatal, and the surrounding good frames must still parse. +# --------------------------------------------------------------------------- + + +def _at_message(bad: Any) -> str: + return _text_stream(message=bad) + + +def _at_message_usage(bad: Any) -> str: + return _text_stream(message={**_REAL_MESSAGE, "usage": bad}) + + +def _at_message_stop_details(bad: Any) -> str: + return _text_stream(message={**_REAL_MESSAGE, "stop_details": bad}) + + +def _at_content_block(bad: Any) -> str: + return _text_stream( + extra=_frame( + "content_block_start", + {"type": "content_block_start", "index": 5, "content_block": bad}, + ) + ) + + +def _at_block_start_index(bad: Any) -> str: + return _text_stream( + extra=_frame( + "content_block_start", + { + "type": "content_block_start", + "index": bad, + "content_block": {"type": "text", "text": ""}, + }, + ) + ) + + +def _at_block_delta_index(bad: Any) -> str: + return _text_stream( + extra=_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": bad, + "delta": {"type": "text_delta", "text": "X"}, + }, + ) + ) + + +def _at_block_stop_index(bad: Any) -> str: + return _text_stream( + extra=_frame("content_block_stop", {"type": "content_block_stop", "index": bad}) + ) + + +def _at_content_block_delta(bad: Any) -> str: + return _text_stream( + mid=_frame("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": bad}) + ) + + +def _at_text_delta_text(bad: Any) -> str: + return _text_stream( + mid=_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": bad}, + }, + ) + ) + + +def _at_partial_json(bad: Any) -> str: + return _text_stream( + mid=_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": bad}, + }, + ) + ) + + +def _at_thinking_delta(bad: Any) -> str: + return _text_stream( + mid=_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "thinking_delta", "thinking": bad}, + }, + ) + ) + + +def _at_message_delta(bad: Any) -> str: + return _text_stream(extra=_frame("message_delta", {"type": "message_delta", "delta": bad})) + + +def _at_message_delta_usage(bad: Any) -> str: + return _text_stream( + extra=_frame("message_delta", {"type": "message_delta", "delta": {}, "usage": bad}) + ) + + +def _at_top_level_event(bad: Any) -> str: + return _text_stream(extra=f"event: garbage\ndata: {json.dumps(bad)}\n\n") + + +def _at_citations_copy_through(bad: Any) -> str: + """Upstream seeds a non-list `citations` on a non-standard block, then a + citations_delta tries to append to it.""" + return _text_stream( + extra=_frame( + "content_block_start", + { + "type": "content_block_start", + "index": 9, + "content_block": {"type": "server_tool_use", "citations": bad}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 9, + "delta": {"type": "citations_delta", "citation": {"cited_text": "x"}}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 9}) + ) + + +def _at_scratch_key_injection(bad: Any) -> str: + """Upstream tries to seed the parser's own accumulator scratch keys.""" + return _text_stream( + extra=_frame( + "content_block_start", + { + "type": "content_block_start", + "index": 8, + "content_block": { + "type": "server_tool_use", + "_partial_json": bad, + "thinking_buffer": bad, + }, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 8, + "delta": {"type": "input_json_delta", "partial_json": '{"a":1}'}, + }, + ) + + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 8, + "delta": {"type": "thinking_delta", "thinking": "t"}, + }, + ) + + _frame("content_block_stop", {"type": "content_block_stop", "index": 8}) + ) + + +_INJECTIONS = [ + ("message", _at_message), + ("message.usage", _at_message_usage), + # No guard needed at this position — the parser copies stop_details + # through without inspecting it. Kept so the sweep's contract ("no + # position raises") stays exhaustive if that ever changes. + ("message.stop_details", _at_message_stop_details), + ("content_block", _at_content_block), + ("content_block_start.index", _at_block_start_index), + ("content_block_delta.index", _at_block_delta_index), + ("content_block_stop.index", _at_block_stop_index), + ("content_block_delta.delta", _at_content_block_delta), + ("text_delta.text", _at_text_delta_text), + ("input_json_delta.partial_json", _at_partial_json), + ("thinking_delta.thinking", _at_thinking_delta), + ("message_delta.delta", _at_message_delta), + ("message_delta.usage", _at_message_delta_usage), + ("top-level event", _at_top_level_event), + ("citations copy-through", _at_citations_copy_through), + ("scratch-key injection", _at_scratch_key_injection), +] + + +@pytest.mark.parametrize("position,build", _INJECTIONS, ids=[p for p, _ in _INJECTIONS]) +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_response_parser_survives_wrong_shape(position: str, build: Any, bad: Any) -> None: + resp = _response(build(bad)) + assert resp is not None, f"{position}={bad!r} lost the whole response" + # The good delta's text still lands. Substring rather than equality: when + # `bad` happens to be a valid string it is a legitimate text_delta and + # correctly accumulates alongside "hi". + texts = [b.get("text") or "" for b in resp["content"] if b.get("type") == "text"] + assert any("hi" in t for t in texts), f"{position}={bad!r} destroyed the good content block" + + +@pytest.mark.parametrize("position,build", _INJECTIONS, ids=[p for p, _ in _INJECTIONS]) +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_usage_parser_survives_wrong_shape(position: str, build: Any, bad: Any) -> None: + usage = _usage(build(bad)) + assert usage is not None, f"{position}={bad!r} lost all usage" + # message_start's input_tokens survives whatever was injected around it. + if position not in {"message", "message.usage"}: + assert usage["input_tokens"] == 12, f"{position}={bad!r} corrupted input_tokens" + + +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_message_start_wrong_shape_still_yields_later_usage(bad: Any) -> None: + # A malformed message_start must not prevent message_delta's usage landing. + usage = _usage(_text_stream(message=bad)) + assert usage is not None + assert usage["output_tokens"] == 4 + + +# --------------------------------------------------------------------------- +# The two usage parsers, across every provider branch. +# --------------------------------------------------------------------------- + +_USAGE_POSITIONS = [ + ("anthropic", lambda b: {"type": "message_start", "message": b}), + ("anthropic", lambda b: {"type": "message_start", "message": {"usage": b}}), + ("anthropic", lambda b: {"type": "message_delta", "usage": b}), + ("openai", lambda b: {"usage": b}), + ("openai", lambda b: {"response": b}), + ("openai", lambda b: {"usage": {"prompt_tokens_details": b}}), + ("gemini", lambda b: {"usageMetadata": b}), +] + + +@pytest.mark.parametrize("provider,build", _USAGE_POSITIONS) +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_buffered_usage_parser_survives_wrong_shape(provider: str, build: Any, bad: Any) -> None: + # Must not raise; returning None (nothing usable) is the correct outcome. + _usage(_frame("x", build(bad)), provider) + + +@pytest.mark.parametrize("provider,build", _USAGE_POSITIONS) +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_chunk_usage_parser_survives_wrong_shape(provider: str, build: Any, bad: Any) -> None: + _usage_chunk(_frame("x", build(bad)), provider) + + +@pytest.mark.parametrize("provider", ["anthropic", "openai", "gemini"]) +@pytest.mark.parametrize("bad", _WRONG_TYPES, ids=repr) +def test_both_usage_parsers_skip_non_object_events(provider: str, bad: Any) -> None: + payload = f"data: {json.dumps(bad)}\n\n" + assert not _usage(payload, provider) + assert not _usage_chunk(payload, provider) + + +def test_gemini_usage_still_parses(provider: str = "gemini") -> None: + # Pins that the gemini guard did not break the branch it hardened. + usage = _usage( + _frame("x", {"usageMetadata": {"promptTokenCount": 7, "candidatesTokenCount": 3}}), provider + ) + assert usage is not None + assert usage["input_tokens"] == 7 + assert usage["output_tokens"] == 3 + + +def test_openai_usage_still_parses() -> None: + usage = _usage(_frame("x", {"usage": {"prompt_tokens": 9, "completion_tokens": 2}}), "openai") + assert usage is not None + assert usage["input_tokens"] == 9 + assert usage["output_tokens"] == 2 + + +def test_unparseable_json_is_still_skipped() -> None: + # Pre-existing behaviour: the JSONDecodeError guard already covered this. + usage = _usage(_text_stream(extra="event: broken\ndata: {not json\n\n")) + assert usage is not None + assert usage["input_tokens"] == 12