diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index f0946ed2d..2ac5bc7db 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -318,6 +318,44 @@ class AnthropicHandlerMixin: return True return False + @staticmethod + def _outgoing_body_has_redeemable_marker(body: Any) -> bool: + """Does the body about to be sent carry a marker retrieval could expand? + + ``headroom_retrieve`` exists only to expand a ``<>`` marker, so + a request carrying none cannot benefit from the buffered path (#3071). + + Ownership is verified rather than shape-matched: the marker shape is not + unique to Headroom, and adopting another context tool's hash would send + the model to an endpoint that is guaranteed to miss (#2836). A hash that + survives ``verify_ownership`` is redeemable right now. + + Errors are swallowed deliberately and answered ``True``. This gates a + wire-format decision, and the safe direction on an unexpected message + shape is the long-standing buffered behavior, not a silent change. + """ + if not isinstance(body, dict): + return True + messages = body.get("messages") + if not isinstance(messages, list) or not messages: + return False + try: + from headroom.ccr.tool_injection import CCRToolInjector + + probe = CCRToolInjector( + provider="anthropic", + inject_tool=False, + inject_system_instructions=False, + ) + probe.scan_for_markers(messages) + if not probe.detected_hashes: + return False + probe.verify_ownership() + return bool(probe.detected_hashes) + except Exception: # pragma: no cover - defensive + logger.debug("CCR: marker probe failed; keeping the buffered path", exc_info=True) + return True + @staticmethod def _extract_anthropic_cache_ttl_metrics(usage: dict[str, Any] | None) -> tuple[int, int]: """Extract observed Anthropic cache-write TTL bucket usage. @@ -3368,13 +3406,40 @@ class AnthropicHandlerMixin: body=body, original_body_bytes=original_body_bytes, ) - wants_buffered_stream_ccr = bool( + # ``headroom_retrieve`` stays resident for the session lifetime so + # the tools array is byte-stable and the prompt cache survives, so + # its mere presence is a poor reason to buffer. Once a session went + # sticky, *every* later streaming turn took the buffered path, and + # buffering replaces incremental delivery with one write at the + # end: time-to-last-byte is roughly unchanged, but time-to-first- + # byte becomes the entire generation. In the traffic reported in + # #3071 that was 8s on average and up to 100s, on 234 requests in + # a single day. + # + # The tool can only expand a marker that is in the outgoing body + # and redeemable now, so a turn carrying none cannot benefit from + # server-side retrieval and should keep streaming. This reads + # ``body`` rather than the earlier scan of ``optimized_messages`` + # because memory hooks, pre-send extensions and the CCR/tool-search + # repairs can all replace the message list after that scan; the + # only list that matters is the one about to go on the wire. + retrieve_tool_is_offered = ( stream and ccr_response_handler_enabled and self._has_headroom_retrieve_tool( tools if tools is not None else body.get("tools") ) ) + buffered_retrieval_can_help = ( + retrieve_tool_is_offered and self._outgoing_body_has_redeemable_marker(body) + ) + wants_buffered_stream_ccr = bool(buffered_retrieval_can_help) + if retrieve_tool_is_offered and not buffered_retrieval_can_help: + logger.info( + f"[{request_id}] CCR: headroom_retrieve is resident but this " + "request carries no redeemable marker, so server-side " + "retrieval cannot fire; keeping the streaming path (#3071)" + ) buffered_stream_ccr = ( wants_buffered_stream_ccr and not outbound_locked_to_client_bytes ) diff --git a/tests/test_ccr_buffered_stream_signed_thinking.py b/tests/test_ccr_buffered_stream_signed_thinking.py index b90a599d2..586fe6d27 100644 --- a/tests/test_ccr_buffered_stream_signed_thinking.py +++ b/tests/test_ccr_buffered_stream_signed_thinking.py @@ -62,8 +62,35 @@ def _config() -> ProxyConfig: ) -def _body(*, with_thinking: bool) -> dict: - messages: list[dict] = [{"role": "user", "content": "hi"}] +@pytest.fixture +def ccr_marker() -> str: + """A marker this proxy actually owns, so retrieval could really fire. + + The buffered path is only taken when the outgoing body carries a redeemable + marker (#3071) — ``headroom_retrieve`` has nothing to expand otherwise. These + tests are about what happens *on* that path, so they have to earn it. + """ + from headroom.cache.backends import InMemoryBackend + from headroom.cache.compression_store import get_compression_store, reset_compression_store + + reset_compression_store() + store = get_compression_store(backend=InMemoryBackend()) + hash_key = store.store( + "the original, uncompressed tool output", + "<>", + original_tokens=100, + compressed_tokens=5, + tool_name="Read", + ) + try: + yield hash_key + finally: + reset_compression_store() + + +def _body(*, with_thinking: bool, marker: str | None = None) -> dict: + first = "hi" if marker is None else f"hi — earlier output is at <>" + messages: list[dict] = [{"role": "user", "content": first}] if with_thinking: messages.append(SIGNED_THINKING_TURN) messages.append({"role": "user", "content": "continue"}) @@ -85,7 +112,7 @@ def _headers() -> dict[str, str]: [(True, True), (False, False)], ) def test_signed_thinking_history_skips_the_buffered_ccr_path( - with_thinking: bool, expect_plain_streaming: bool + with_thinking: bool, expect_plain_streaming: bool, ccr_marker: str ) -> None: """The buffered path is only chosen when the stream:false flip can land.""" calls: dict[str, object] = {} @@ -115,7 +142,9 @@ def test_signed_thinking_history_skips_the_buffered_ccr_path( 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() + "/v1/messages", + json=_body(with_thinking=with_thinking, marker=ccr_marker), + headers=_headers(), ) assert resp.status_code == 200, resp.text @@ -133,7 +162,7 @@ def test_signed_thinking_history_skips_the_buffered_ccr_path( @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, + upstream_delay: float, ccr_marker: str ) -> None: """A 200 SSE reply on the buffered path reaches the client as a stream. @@ -153,7 +182,11 @@ def test_buffered_ccr_relays_an_unexpected_sse_reply_and_does_not_cache_it( 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()) + resp = client.post( + "/v1/messages", + json=_body(with_thinking=False, marker=ccr_marker), + headers=_headers(), + ) assert resp.status_code == 200, resp.text assert resp.headers["content-type"].startswith("text/event-stream") @@ -164,6 +197,76 @@ def test_buffered_ccr_relays_an_unexpected_sse_reply_and_does_not_cache_it( assert len(proxy.cache._cache) == 0, "an unparseable body must never be cached" +@pytest.mark.parametrize( + ("marker_kind", "expect_buffered"), + [ + ("owned", True), + ("none", False), + ("foreign", False), + ], +) +def test_buffering_is_gated_on_a_redeemable_marker( + marker_kind: str, expect_buffered: bool, ccr_marker: str +) -> None: + """A resident ``headroom_retrieve`` is not on its own a reason to buffer (#3071). + + The tool is injected once and kept resident so the tools array stays + byte-stable for the prompt cache. Buffering on its presence alone meant + every later streaming turn of a sticky session lost incremental delivery — + time-to-first-byte became the whole generation. Retrieval can only expand a + marker that is in the outgoing body *and* redeemable now, so that is what + the wire-format decision keys on. + """ + marker = { + "owned": ccr_marker, + "none": None, + # Correct shape, not ours: adopting it would send the model to a + # retrieval that is guaranteed to miss (#2836). + "foreign": "deadbeefcafe", + }[marker_kind] + + 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=False, marker=marker), + headers=_headers(), + ) + + assert resp.status_code == 200, resp.text + if expect_buffered: + assert "buffered_body" in calls, "a redeemable marker must still buffer" + assert calls["buffered_body"]["stream"] is False + else: + assert "stream_body" in calls, "nothing to retrieve — the client must keep streaming" + assert "buffered_body" not in calls + assert calls["stream_body"]["stream"] is True + + 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 = { diff --git a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py index 478eaceba..70889b647 100644 --- a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py +++ b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py @@ -35,6 +35,37 @@ def _make_config() -> ProxyConfig: ) +@pytest.fixture(autouse=True) +def _fresh_compression_store(): + """Each test gets its own store, so seeded markers cannot leak between them.""" + from headroom.cache.backends import InMemoryBackend + from headroom.cache.compression_store import get_compression_store, reset_compression_store + + reset_compression_store() + get_compression_store(backend=InMemoryBackend()) + try: + yield + finally: + reset_compression_store() + + +def _buffered(text: str) -> str: + """User content carrying a marker this proxy owns. + + The buffered path engages only when retrieval has something to expand + (#3071); a resident ``headroom_retrieve`` with no redeemable marker in the + request keeps streaming. Every test below that is *about* the buffered path + therefore has to earn it with a real marker rather than the tool alone. + """ + store = get_compression_store() + hash_key = store.store( + original=json.dumps({"earlier": "tool output"}), + compressed="{}", + original_item_count=1, + ) + return f"{text} (earlier output at <>)" + + def _message_response(content: list[dict], *, stop_reason: str = "end_turn") -> dict: return { "id": "msg_test", @@ -127,7 +158,7 @@ def test_streaming_headroom_retrieve_is_intercepted_and_returned_as_sse() -> Non "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "retrieve it"}], + "messages": [{"role": "user", "content": _buffered("retrieve it")}], }, ) @@ -213,7 +244,7 @@ def test_streaming_with_headroom_retrieve_available_but_unused_returns_sse() -> "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "hello"}], + "messages": [{"role": "user", "content": _buffered("hello")}], }, ) @@ -282,7 +313,7 @@ def test_mixed_ccr_and_client_tool_streams_both_blocks_as_sse() -> None: "input_schema": {"type": "object", "properties": {}}, }, ], - "messages": [{"role": "user", "content": "use tools"}], + "messages": [{"role": "user", "content": _buffered("use tools")}], }, ) @@ -341,7 +372,7 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None: "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "use tools"}], + "messages": [{"role": "user", "content": _buffered("use tools")}], }, ) @@ -368,7 +399,7 @@ async def test_buffered_ccr_withholds_output_until_delayed_upstream_resolves() - "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "wait"}], + "messages": [{"role": "user", "content": _buffered("wait")}], } request_delivered = False @@ -439,7 +470,7 @@ async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "fail early"}], + "messages": [{"role": "user", "content": _buffered("fail early")}], } async def receive(): @@ -508,7 +539,7 @@ async def test_buffered_ccr_preserves_late_failure_status_and_headers() -> None: "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "fail late"}], + "messages": [{"role": "user", "content": _buffered("fail late")}], } async def receive(): @@ -588,7 +619,7 @@ def test_buffered_ccr_rejects_malformed_success_as_502() -> None: "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "fail safely"}], + "messages": [{"role": "user", "content": _buffered("fail safely")}], }, ) @@ -608,7 +639,7 @@ async def test_buffered_ccr_late_failure_returns_sanitized_json_error() -> None: "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "wait"}], + "messages": [{"role": "user", "content": _buffered("wait")}], } async def receive(): @@ -684,7 +715,7 @@ async def test_buffered_ccr_pre_keepalive_exception_returns_json_error() -> None "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "fail before keepalive"}], + "messages": [{"role": "user", "content": _buffered("fail before keepalive")}], } async def receive(): diff --git a/tests/test_proxy_response_cache_replay.py b/tests/test_proxy_response_cache_replay.py index 776d2569d..738b3ad15 100644 --- a/tests/test_proxy_response_cache_replay.py +++ b/tests/test_proxy_response_cache_replay.py @@ -31,6 +31,11 @@ httpx = pytest.importorskip("httpx") from fastapi.testclient import TestClient # noqa: E402 +from headroom.cache.backends import InMemoryBackend # noqa: E402 +from headroom.cache.compression_store import ( # noqa: E402 + get_compression_store, + reset_compression_store, +) from headroom.ccr.tool_injection import create_ccr_tool_definition # noqa: E402 from headroom.proxy.helpers import sanitize_forwarded_response_headers # noqa: E402 from headroom.proxy.models import CacheEntry # noqa: E402 @@ -245,6 +250,16 @@ def test_buffered_ccr_turn_does_not_write_the_response_cache(): }, } + # The buffered conversion needs a marker retrieval could actually expand; + # a resident `headroom_retrieve` alone keeps the request streaming (#3071). + reset_compression_store() + store = get_compression_store(backend=InMemoryBackend()) + marker = store.store( + original=json.dumps({"earlier": "tool output"}), + compressed="{}", + original_item_count=1, + ) + with patch("headroom.proxy.server.AnyLLMBackend"): app = create_app(_ccr_cache_config()) with TestClient(app) as client: @@ -273,7 +288,9 @@ def test_buffered_ccr_turn_does_not_write_the_response_cache(): "max_tokens": 64, "stream": True, "tools": [create_ccr_tool_definition("anthropic")], - "messages": [{"role": "user", "content": "hello"}], + "messages": [ + {"role": "user", "content": f"hello (earlier at <>)"} + ], }, ) @@ -282,6 +299,7 @@ def test_buffered_ccr_turn_does_not_write_the_response_cache(): assert forwarded_bodies and forwarded_bodies[0]["stream"] is False # ...and nothing was written to the response cache. proxy.cache.set.assert_not_awaited() + reset_compression_store() # --------------------------------------------------------------------------