From 8369b7ff37317290b444569fc1cfbe6490cca053 Mon Sep 17 00:00:00 2001 From: Navid Kamali Date: Sat, 25 Jul 2026 23:47:04 -0700 Subject: [PATCH] fix(proxy/streaming): tolerate malformed upstream SSE events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Three parsers read upstream-controlled SSE bodies with `json.loads` and then reach into the result without checking its shape: `_parse_sse_usage`, `_parse_sse_usage_from_buffer` and `_parse_sse_to_response`. A valid-JSON value of the wrong shape passes the `json.JSONDecodeError` guard and then raises. Three distinct ways: - a non-object event (`["x"]`, `"str"`, `42`) makes `.get` raise `AttributeError` - an explicit `"key": null` bypasses a `.get(key, {})` default, because the default applies only to a *missing* key, not a present-and-null one - an unhashable `index` (a list or dict) raises `TypeError: unhashable type` when used to key the block map In `_parse_sse_to_response` the raise escapes into `_finalize_stream_response` and tears down a stream the client is already reading. Several call sites reach it from a bare `finally`, so there is no handler above it at all. **This is a robustness invariant, not a bug report about any specific upstream.** I am not claiming Anthropic emits these shapes; I have no evidence of that. The Anthropic handler serves any Anthropic-shaped upstream (`--backend anthropic|bedrock|openrouter|anyllm|litellm-`, `ANTHROPIC_TARGET_API_URL`), so Bedrock, OpenRouter, LiteLLM, vLLM and Vertex all flow through these parsers. A parser on a proxy's critical streaming path should not raise on a well-formed-JSON frame of unexpected shape, whatever the source. The guard costs one `isinstance`; being wrong costs the user's session. **Scope: frame shape only.** Value validity — a token count that is a string, `Infinity`, or a >4300-digit integer literal — is a separate defect class that reaches the same parsers and is deliberately left for a follow-up. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - **headroom/proxy/handlers/streaming.py** - Add three shared normalizers — `_sse_dict`, `_sse_str`, `_sse_index` — so every upstream-derived container is shape-checked with one idiom instead of a mix of inline `isinstance` and `.get(..., {})` defaults. - `_sse_index` rejects `bool` as well as non-`int`, so `true` cannot silently alias block index `1`, and unhashable values cannot reach the block map. - Guard the positions the earlier pass missed: `content_block_delta.delta`, `message_delta.delta`, `message_delta.usage` (`dict.update` on a non-mapping raises — the `message_start` twin was already guarded, this sibling was not), the `index` on all three `content_block_*` events, and the string accumulators for `text` / `partial_json` / `thinking`. - Guard `_parse_sse_usage`, which had no shape checks at all — not on the top-level event, nor on the anthropic, openai or gemini branches. - Guard the gemini `usageMetadata` branch in `_parse_sse_usage_from_buffer`, which used a truthiness check where its two sibling branches used `isinstance`. - Stop the non-standard-block copy-through from seeding the parser's own scratch keys (`_partial_json`, `thinking_buffer`) via `_BLOCK_SCRATCH_KEYS`, and tolerate a non-list `citations` that a copy-through already placed. - **tests/test_streaming_sse_malformed_events.py** — new file, 408 cases. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`, `ruff format`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added The sweep is enumerated, not randomized: the failure surface is finite, so the cases can be exact and no `hypothesis` dependency is needed. This matches the existing precedent in `tests/test_toin_observation_only.py`. Happy-path fixtures are transcribed from real `api.anthropic.com` streams rather than hand-written, which corrected four wrong assumptions in the previous fixtures: `message_delta.usage` repeats the *full* usage block (not just `output_tokens`); `message_start.message` carries `stop_details: null` alongside the other two nulls; `usage.cache_creation` is always present as a nested object and is what `_extract_anthropic_cache_ttl_metrics` reads; and a `tool_use` block carries a non-standard `caller` field. ### Test Output ```text $ pytest tests/test_streaming_sse_malformed_events.py 408 passed # same file against origin/main's parsers 191 failed, 217 passed # regression sweep: all 22 test files touching SSE parsing 710 passed, 5 skipped $ mypy headroom --ignore-missing-imports Success: no issues found in 509 source files ``` ## Real Behavior Proof - **Environment:** macOS 15 (Darwin 25.4.0, arm64), Python 3.13.7, Rust 1.95.0 per `rust-toolchain.toml`, `uv sync --extra dev`. - **Real upstream capture:** three live streaming `POST /v1/messages` calls to `api.anthropic.com` (`claude-haiku-4-5`) — a text stream, a `tool_use` stream and an extended-thinking stream — captured as raw wire bytes and used to build the happy-path fixtures. This is what surfaced the four fixture errors above. - **Differential:** the 408-case file was run against `origin/main`'s parsers via `git show origin/main:...` (not `git stash`, which silently no-ops once the fix is committed and produces a false pass). 191 cases fail there and pass here. - **Exhaustive shape fuzz:** 700 generated cases crossing every upstream-derived container position in the three parsers with `None`, `"str"`, `42`, `3.5`, `True`, `False`, `["x"]`, `[]`, `{}`, `{"k":"v"}`. Zero raises after the change; before it, the positions listed under Changes Made all raise. - **Per-position legitimacy check:** each injection position in the sweep was confirmed to produce at least one genuine failure against `origin/main`, so no case passes vacuously. This caught four delta-shaped tests that were injecting their fault before any content block had opened, where the parser skips the body and the guard is never reached. - **Not tested:** - **No evidence these shapes occur in the wild.** The fault shapes are derived from the parsers' own unguarded access positions, not observed in real upstream traffic. I searched the tracker for this crash signature and both function names and found no report. The claim is survivability, not incidence. - The real API cannot be made to emit malformed frames on demand, so the fault injection is synthetic; only the happy-path fixtures come from live traffic. - No live third-party gateway (Bedrock/OpenRouter/LiteLLM/Vertex) exercised, despite those being the motivating case. - No WebSocket path, no Windows, not run against a live Claude Code session. - Value-validity faults are out of scope and still raise; see Additional Notes. ## Additional Notes Two findings left deliberately unfixed, both pre-existing and both outside the frame-shape scope of this change: 1. **Value validity.** `_extract_anthropic_cache_ttl_metrics` calls `int()` on an upstream value with no `try`; `_usage_int` does not catch `OverflowError`, and `json.loads` accepts bare `Infinity`; `except json.JSONDecodeError` is too narrow for a >4300-digit integer literal (plain `ValueError`) or deep nesting (`RecursionError`); and the anthropic and gemini branches export non-`int` token values out of a `dict[str, int]`-annotated function. 2. **Reconstruction fidelity.** `content_block_start` copies through unknown fields only for non-standard block types, so the real `caller` field on a `tool_use` block is dropped on reconstruction. `test_real_tool_use_stream_reconstructs_tool_input` pins the current behaviour rather than asserting the desired one. ## Review Readiness - [x] I have performed a self-review - [x] 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 - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation — n/a, no user-facing surface changed - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` No new dependencies added. --- headroom/proxy/handlers/streaming.py | 130 +++- tests/test_streaming_sse_malformed_events.py | 611 +++++++++++++++++++ 2 files changed, 710 insertions(+), 31 deletions(-) create mode 100644 tests/test_streaming_sse_malformed_events.py diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 3b864f9dd..22b689117 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -30,6 +30,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. @@ -179,6 +222,9 @@ class StreamingMixin: except json.JSONDecodeError: continue + if not isinstance(data, dict): + continue + usage = {} if provider == "anthropic": @@ -187,8 +233,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( @@ -204,24 +249,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"] = usage_meta.get("candidatesTokenCount", 0) @@ -275,11 +320,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( @@ -299,7 +351,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) @@ -339,7 +391,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"] = usage_meta.get("candidatesTokenCount", 0) @@ -402,29 +454,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") @@ -433,7 +492,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": @@ -450,31 +509,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. @@ -485,13 +546,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 @@ -525,13 +591,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