diff --git a/headroom/proxy/body_forwarding.py b/headroom/proxy/body_forwarding.py index f93c68af9..4648fce6c 100644 --- a/headroom/proxy/body_forwarding.py +++ b/headroom/proxy/body_forwarding.py @@ -32,6 +32,15 @@ class OutboundBody: content: bytes source: OutboundBodySource + #: True when byte-faithful passthrough won over a mutated body, so every + #: edit the handler made to ``body`` was discarded before the wire. Callers + #: that gated a downstream decision on their own mutation (for example + #: flipping ``stream`` to False to buffer a reply) MUST consult this — the + #: request upstream actually sees is the client's original one. + dropped_mutations: bool = False + #: The ``BodyMutationTracker`` reasons discarded alongside it, when the + #: caller supplied them. Empty when the caller passed no reason list. + dropped_mutation_reasons: tuple[str, ...] = () def get_python_forwarder_mode() -> PythonForwarderMode: @@ -104,11 +113,23 @@ def select_outbound_body( original_body_bytes: bytes | None, body_mutated: bool, forwarder_mode: PythonForwarderMode | None = None, + mutation_reasons: list[str] | None = None, ) -> OutboundBody: - """Select the exact bytes to forward upstream.""" + """Select the exact bytes to forward upstream. + + ``mutation_reasons`` is optional and only used for reporting: when the + signed-thinking passthrough overrides a mutated body, the discarded reasons + are echoed back on the result so the call site can log what never made it + upstream instead of silently claiming the edit landed. + """ mode = forwarder_mode if forwarder_mode is not None else get_python_forwarder_mode() if original_body_bytes is not None and has_signed_thinking_blocks(body): - return OutboundBody(content=original_body_bytes, source="passthrough") + return OutboundBody( + content=original_body_bytes, + source="passthrough", + dropped_mutations=body_mutated, + dropped_mutation_reasons=tuple(mutation_reasons or ()) if body_mutated else (), + ) if mode == "legacy_json_kwarg": content = json.dumps(body, separators=(", ", ": "), ensure_ascii=True).encode("utf-8") @@ -125,12 +146,38 @@ def prepare_outbound_body_bytes( original_body_bytes: bytes | None, body_mutated: bool, forwarder_mode: PythonForwarderMode | None = None, + mutation_reasons: list[str] | None = None, ) -> tuple[bytes, OutboundBodySource]: - """Compatibility tuple wrapper around :func:`select_outbound_body`.""" + """Compatibility tuple wrapper around :func:`select_outbound_body`. + + Keeps the two-value shape its existing callers unpack. Call + :func:`select_outbound_body` directly when you need the dropped-mutation + reporting. + """ outbound = select_outbound_body( body=body, original_body_bytes=original_body_bytes, body_mutated=body_mutated, forwarder_mode=forwarder_mode, + mutation_reasons=mutation_reasons, ) return outbound.content, outbound.source + + +def outbound_body_is_client_bytes( + *, + body: dict[str, Any], + original_body_bytes: bytes | None, +) -> bool: + """Return whether the wire body will be the client's original bytes. + + A handler that changes ``body`` to steer its own upstream call — the + ``stream`` flip that buys a buffered reply is the load-bearing case — has to + know that the signed-thinking passthrough will throw that change away and + send the client's bytes verbatim. Asking before acting is cheaper than + discovering it from a reply in the wrong wire format. + + Mirrors the first branch of :func:`select_outbound_body`; the forwarder mode + is deliberately not consulted because that branch overrides it too. + """ + return original_body_bytes is not None and has_signed_thinking_blocks(body) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 5426d0cf5..c08da0fa1 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -73,6 +73,20 @@ def _strip_index_from_content_blocks(content: Any) -> None: _strip_index_from_content_blocks(block.get("content")) +def _looks_like_sse_response(response: httpx.Response) -> bool: + """Return whether an upstream reply is a Server-Sent Events stream. + + Trusts the declared content-type first and falls back to sniffing the + leading bytes for an SSE field, because a gateway in front of Anthropic may + relay the stream under a vaguer type. + """ + content_type = (response.headers.get("content-type") or "").lower() + if "text/event-stream" in content_type: + return True + head = response.content[:64].lstrip() + return head.startswith(b"event:") or head.startswith(b"data:") + + class AnthropicHandlerMixin: """Mixin providing Anthropic API handler methods for HeadroomProxy.""" @@ -1054,6 +1068,11 @@ class AnthropicHandlerMixin: response_headers = dict(cached.response_headers) response_headers.pop("content-encoding", None) response_headers.pop("content-length", None) + # Drop the stored content-type too. Starlette lets an + # explicit header win over ``media_type``, so keeping the + # producing request's type would let a cache entry hand this + # caller a wire format it never asked for (#2952). + response_headers.pop("content-type", None) # Unit 4: release the pre-upstream semaphore on cache # hit — no upstream call will happen. @@ -3111,13 +3130,37 @@ class AnthropicHandlerMixin: ccr_response_handler_enabled = bool( self.ccr_response_handler and getattr(ccr_handler_config, "enabled", True) ) - buffered_stream_ccr = bool( + # A body carrying signed thinking blocks leaves as the client's + # original bytes (see ``select_outbound_body``), which throws + # away every edit made here — including the ``stream`` flip + # below. Taking the buffered path anyway asks upstream for a + # stream:true reply and then tries to read it as buffered JSON: + # the parse fails, SSE resynthesis is skipped, and the client + # gets a 200 with no usable body (#2952). The retrieve tool is + # itself an injected (and equally discarded) mutation on these + # turns, so the plain streaming path is the coherent choice. + from headroom.proxy.body_forwarding import outbound_body_is_client_bytes + + outbound_locked_to_client_bytes = outbound_body_is_client_bytes( + body=body, + original_body_bytes=original_body_bytes, + ) + wants_buffered_stream_ccr = bool( stream and ccr_response_handler_enabled and self._has_headroom_retrieve_tool( tools if tools is not None else body.get("tools") ) ) + buffered_stream_ccr = ( + wants_buffered_stream_ccr and not outbound_locked_to_client_bytes + ) + if wants_buffered_stream_ccr and outbound_locked_to_client_bytes: + logger.info( + f"[{request_id}] CCR: signed thinking blocks force byte-faithful " + "passthrough, so a stream:false flip could not reach upstream; " + "using the plain streaming path instead of buffered retrieval" + ) if buffered_stream_ccr: if body.get("stream") is not False: body["stream"] = False @@ -3400,9 +3443,23 @@ class AnthropicHandlerMixin: try: resp_json = response.json() except (json.JSONDecodeError, ValueError) as e: - logger.debug( - f"[{request_id}] Failed to parse response JSON for CCR handling: {e}" - ) + # DEBUG is right for the buffered non-stream path, where + # an unparseable body is just "no CCR handling". On the + # buffered-stream path it means the reply came back in a + # wire format we did not ask for, and every downstream + # step (retrieval, SSE resynthesis, usage accounting) + # silently no-ops — that has to be visible (#2952). + if buffered_stream_ccr: + logger.warning( + f"[{request_id}] CCR: buffered stream:false request got a " + f"non-JSON {response.status_code} reply " + f"(content-type={response.headers.get('content-type')!r}): {e}" + ) + else: + logger.debug( + f"[{request_id}] Failed to parse response JSON for CCR " + f"handling: {e}" + ) # CCR Response Handling: Handle headroom_retrieve tool calls automatically if ( @@ -3730,7 +3787,13 @@ class AnthropicHandlerMixin: # Cache response under the SAME key it was looked up by: # cache_lookup_messages is the raw pre-mutation snapshot, not # the live (compressed/hooked) `messages` (#327). - if self.cache and response.status_code == 200: + # ``resp_json`` is None when the reply did not parse as + # JSON — an SSE stream, most often. Caching those bytes + # poisons the entry for every later caller that shares + # the key: the cache key has no ``stream`` component, so + # a buffered request would be answered with a stream it + # cannot read (#2952). + if self.cache and response.status_code == 200 and resp_json is not None: await self.cache.set( cache_lookup_messages, model, @@ -3893,6 +3956,44 @@ class AnthropicHandlerMixin: f"[{request_id}] Security response scan error: {sec_err}" ) + if ( + buffered_stream_ccr + and response.status_code == 200 + and not resp_json + and _looks_like_sse_response(response) + ): + # Upstream streamed instead of buffering, so there is + # nothing to resynthesize — but the client asked for a + # stream and this already is one. Relay it verbatim + # rather than falling through to a plain Response the + # _BufferedCCRResponse wrapper can only turn into a bare + # error event (#2952). + logger.warning( + f"[{request_id}] CCR: relaying the upstream SSE reply verbatim; " + "server-side retrieval was skipped for this turn" + ) + relay_headers = { + k: v + for k, v in response_headers.items() + if k.lower() + not in ( + "content-encoding", + "content-length", + "transfer-encoding", + "content-type", + ) + } + relayed_sse = response.content + + async def _upstream_sse_relay(): + yield relayed_sse + + return StreamingResponse( + _upstream_sse_relay(), + media_type="text/event-stream", + headers=relay_headers, + ) + if buffered_stream_ccr and response.status_code == 200 and resp_json: sse_headers = { k: v diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 1092c58a5..b9cce8a23 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -327,12 +327,18 @@ def log_outbound_request( mutation_reasons: list[str], request_id: str | None, source: str, + dropped_mutation_reasons: tuple[str, ...] | list[str] | None = None, ) -> None: """Structured log line for every outbound forwarder call. Per realignment build constraints: every cache-affecting decision is logged. Never includes ``Authorization``/``x-api-key`` content or full body bytes. + + ``dropped_mutation_reasons`` records edits that byte-faithful passthrough + discarded before the wire. That is a WARNING, not a detail: the line above + reports the transforms Headroom *decided* on, and without this the operator + reads savings and injections that the upstream never saw. """ logger.info( "event=outbound_request forwarder=%s method=%s path=%s body_bytes=%d " @@ -346,6 +352,16 @@ def log_outbound_request( source, request_id or "", ) + if dropped_mutation_reasons: + logger.warning( + "event=outbound_body_mutations_dropped forwarder=%s source=%s " + "dropped_mutation_reasons=%s request_id=%s (signed thinking blocks force " + "byte-faithful passthrough, so these body edits did NOT reach upstream)", + forwarder, + source, + ",".join(dropped_mutation_reasons), + request_id or "", + ) def count_cache_breakpoints( diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 1c9fae65e..bd2dc3acb 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -2128,16 +2128,18 @@ class HeadroomProxy( construct their body from scratch, so canonical serialization is correct and original bytes do not exist). """ - from headroom.proxy.body_forwarding import prepare_outbound_body_bytes + from headroom.proxy.body_forwarding import select_outbound_body from headroom.proxy.helpers import log_outbound_request last_error = None reasons = list(mutation_reasons or []) - outbound_bytes, source = prepare_outbound_body_bytes( + outbound = select_outbound_body( body=body, original_body_bytes=original_body_bytes, body_mutated=body_mutated, + mutation_reasons=reasons, ) + outbound_bytes, source = outbound.content, outbound.source outbound_headers = {**headers, "content-type": "application/json"} log_outbound_request( @@ -2149,6 +2151,7 @@ class HeadroomProxy( mutation_reasons=reasons, request_id=request_id, source=source, + dropped_mutation_reasons=outbound.dropped_mutation_reasons, ) post_kwargs: dict = {"content": outbound_bytes, "headers": outbound_headers} diff --git a/tests/test_ccr_buffered_stream_signed_thinking.py b/tests/test_ccr_buffered_stream_signed_thinking.py new file mode 100644 index 000000000..b90a599d2 --- /dev/null +++ b/tests/test_ccr_buffered_stream_signed_thinking.py @@ -0,0 +1,219 @@ +"""Buffered-CCR streaming vs. byte-faithful passthrough (issue #2952). + +The buffered-CCR path is the one place the Anthropic handler changes the +request *for its own benefit*: it flips ``stream`` to False so the reply comes +back as one JSON document it can inspect for ``headroom_retrieve`` calls, then +resynthesizes SSE for the client. + +That only works if the flip reaches the wire. When conversation history carries +a signed ``thinking`` block, ``select_outbound_body`` forwards the client's +original bytes instead — ``"stream": true`` and all — so upstream streams, the +JSON parse fails, resynthesis is skipped, and the client is left with a 200 and +nothing it can read. These tests pin the three defenses: don't take the path, +survive the reply if we somehow do, and never cache a body in the wrong format. +""" + +from __future__ import annotations + +import asyncio +import json +from datetime import datetime + +import pytest + +fastapi = pytest.importorskip("fastapi") +httpx = pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.models import CacheEntry # noqa: E402 +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +RETRIEVE_TOOL = { + "name": "headroom_retrieve", + "description": "Retrieve original content", + "input_schema": {"type": "object", "properties": {}}, +} + +SIGNED_THINKING_TURN = { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "private reasoning", + "signature": "sig-abc123", + }, + {"type": "text", "text": "Answered."}, + ], +} + +SSE_BODY = ( + b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' +) + + +def _config() -> ProxyConfig: + return ProxyConfig( + optimize=False, + cache_enabled=True, + rate_limit_enabled=False, + memory_enabled=False, + ) + + +def _body(*, with_thinking: bool) -> dict: + messages: list[dict] = [{"role": "user", "content": "hi"}] + if with_thinking: + messages.append(SIGNED_THINKING_TURN) + messages.append({"role": "user", "content": "continue"}) + return { + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "stream": True, + "tools": [RETRIEVE_TOOL], + "messages": messages, + } + + +def _headers() -> dict[str, str]: + return {"Authorization": "Bearer test-key", "x-api-key": "test-key"} + + +@pytest.mark.parametrize( + ("with_thinking", "expect_plain_streaming"), + [(True, True), (False, False)], +) +def test_signed_thinking_history_skips_the_buffered_ccr_path( + with_thinking: bool, expect_plain_streaming: bool +) -> None: + """The buffered path is only chosen when the stream:false flip can land.""" + calls: dict[str, object] = {} + + async def fake_stream_response(url, headers, body, *args, **kwargs): # noqa: ANN001 + calls["stream_body"] = body + return fastapi.responses.StreamingResponse(iter([SSE_BODY]), media_type="text/event-stream") + + async def fake_retry(method, url, headers, req_body, *args, **kwargs): # noqa: ANN001 + calls["buffered_body"] = json.loads(json.dumps(req_body)) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + headers={"content-type": "application/json"}, + ) + + app = create_app(_config()) + with TestClient(app) as client: + client.app.state.proxy._stream_response = fake_stream_response + client.app.state.proxy._retry_request = fake_retry + resp = client.post( + "/v1/messages", json=_body(with_thinking=with_thinking), headers=_headers() + ) + + assert resp.status_code == 200, resp.text + if expect_plain_streaming: + # Passthrough is locked in, so we must not pretend we can buffer. + assert "stream_body" in calls, "expected the plain streaming path" + assert "buffered_body" not in calls + # The turn still leaves as a streaming request, matching the bytes + # that passthrough will actually forward. + assert calls["stream_body"]["stream"] is True + else: + assert "buffered_body" in calls, "expected the buffered CCR path" + assert calls["buffered_body"]["stream"] is False + + +@pytest.mark.parametrize("upstream_delay", [0.0, 1.2], ids=["prompt", "past-keepalive"]) +def test_buffered_ccr_relays_an_unexpected_sse_reply_and_does_not_cache_it( + upstream_delay: float, +) -> None: + """A 200 SSE reply on the buffered path reaches the client as a stream. + + The delay matters: ``_BufferedCCRResponse`` commits SSE response headers + after a 1 s keepalive, and past that point it can only forward a result + that exposes a ``body_iterator``. A plain ``Response`` there degrades to a + bare ``event: error`` — which is what a real (multi-second) Anthropic turn + hit in #2952. + """ + + async def fake_retry(method, url, headers, req_body, *args, **kwargs): # noqa: ANN001 + if upstream_delay: + await asyncio.sleep(upstream_delay) + return httpx.Response(200, content=SSE_BODY, headers={"content-type": "text/event-stream"}) + + app = create_app(_config()) + with TestClient(app) as client: + proxy = client.app.state.proxy + proxy._retry_request = fake_retry + resp = client.post("/v1/messages", json=_body(with_thinking=False), headers=_headers()) + + assert resp.status_code == 200, resp.text + assert resp.headers["content-type"].startswith("text/event-stream") + assert b"message_start" in resp.content + # Caching SSE bytes under a key with no `stream` component is what + # served a stream to a buffered caller in the first place. + assert proxy.cache is not None + assert len(proxy.cache._cache) == 0, "an unparseable body must never be cached" + + +def test_cache_hit_never_replays_a_foreign_content_type() -> None: + """A cache entry cannot hand a caller a wire format it did not ask for.""" + body = { + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "stream": False, + "messages": [{"role": "user", "content": "hi"}], + } + payload = json.dumps( + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-20250514", + "content": [{"type": "text", "text": "cached"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ).encode() + + async def fail_retry(*args, **kwargs): # noqa: ANN001, ANN002, ANN003 + raise AssertionError("upstream must not be called on a cache hit") + + app = create_app(_config()) + with TestClient(app) as client: + proxy = client.app.state.proxy + proxy._retry_request = fail_retry + key = proxy.cache._compute_key( + body["messages"], + body["model"], + system=None, + tools=None, + tool_choice=None, + temperature=None, + top_p=None, + top_k=None, + max_tokens=64, + stop=None, + thinking=None, + output_config=None, + ) + proxy.cache._cache[key] = CacheEntry( + response_body=payload, + response_headers={"content-type": "text/event-stream"}, + created_at=datetime.now(), + ttl_seconds=3600, + ) + + resp = client.post("/v1/messages", json=body, headers=_headers()) + + assert resp.status_code == 200, resp.text + assert resp.headers["content-type"].startswith("application/json") + assert resp.json()["content"][0]["text"] == "cached" diff --git a/tests/test_proxy_byte_faithful_forwarding.py b/tests/test_proxy_byte_faithful_forwarding.py index 7fc2dbd69..eacfdc0bd 100644 --- a/tests/test_proxy_byte_faithful_forwarding.py +++ b/tests/test_proxy_byte_faithful_forwarding.py @@ -34,6 +34,7 @@ from headroom.proxy.body_forwarding import ( BodyMutationTracker, OutboundBody, get_python_forwarder_mode, + outbound_body_is_client_bytes, prepare_outbound_body_bytes, select_outbound_body, serialize_body_canonical, @@ -252,7 +253,113 @@ def test_signed_thinking_history_overrides_legacy_encoder() -> None: forwarder_mode="legacy_json_kwarg", ) - assert outbound == OutboundBody(content=original, source="passthrough") + assert outbound == OutboundBody(content=original, source="passthrough", dropped_mutations=True) + + +def test_signed_thinking_passthrough_reports_the_mutations_it_discarded() -> None: + """Passthrough silently winning over a mutated body is what hid #2952.""" + body = { + "stream": False, + "messages": [ + { + "role": "assistant", + "content": [{"type": "thinking", "signature": "sig123"}], + } + ], + } + original = json.dumps({**body, "stream": True}).encode("utf-8") + + outbound = select_outbound_body( + body=body, + original_body_bytes=original, + body_mutated=True, + forwarder_mode="byte_faithful", + mutation_reasons=["ccr_streaming_retrieve_buffered_non_stream"], + ) + + assert outbound.source == "passthrough" + assert outbound.dropped_mutations is True + assert outbound.dropped_mutation_reasons == ("ccr_streaming_retrieve_buffered_non_stream",) + + +def test_signed_thinking_passthrough_reports_nothing_when_body_unmutated() -> None: + body = { + "messages": [ + { + "role": "assistant", + "content": [{"type": "thinking", "signature": "sig123"}], + } + ] + } + original = json.dumps(body).encode("utf-8") + + outbound = select_outbound_body( + body=body, + original_body_bytes=original, + body_mutated=False, + forwarder_mode="byte_faithful", + mutation_reasons=["irrelevant"], + ) + + assert outbound.source == "passthrough" + assert outbound.dropped_mutations is False + assert outbound.dropped_mutation_reasons == () + + +def test_canonical_path_reports_no_dropped_mutations() -> None: + body = {"messages": [{"role": "user", "content": "hi"}]} + + outbound = select_outbound_body( + body=body, + original_body_bytes=b'{"messages": []}', + body_mutated=True, + forwarder_mode="byte_faithful", + mutation_reasons=["compression"], + ) + + assert outbound.source == "canonical" + assert outbound.dropped_mutations is False + assert outbound.dropped_mutation_reasons == () + + +@pytest.mark.parametrize( + ("original_body_bytes", "expected"), + [(b'{"messages": []}', True), (None, False)], +) +def test_outbound_body_is_client_bytes_matches_selection( + original_body_bytes: bytes | None, expected: bool +) -> None: + """Handlers gate on this before mutating a body for their own upstream call.""" + body = { + "messages": [ + { + "role": "assistant", + "content": [{"type": "thinking", "signature": "sig123"}], + } + ] + } + + assert ( + outbound_body_is_client_bytes(body=body, original_body_bytes=original_body_bytes) + is expected + ) + outbound = select_outbound_body( + body=body, + original_body_bytes=original_body_bytes, + body_mutated=True, + forwarder_mode="byte_faithful", + ) + assert (outbound.source == "passthrough") is expected + + +def test_outbound_body_is_client_bytes_false_without_thinking_blocks() -> None: + assert ( + outbound_body_is_client_bytes( + body={"messages": [{"role": "user", "content": "hi"}]}, + original_body_bytes=b'{"messages": []}', + ) + is False + ) def test_prepare_outbound_no_original_bytes_uses_canonical() -> None: