headroom/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py
Tejas Chopra c502087db7
fix(ccr): only buffer a stream when a marker is actually redeemable (#3092)
## Description

Closes #3071

`headroom_retrieve` is injected once and kept resident for the session
so the tools array stays byte-stable and the prompt cache survives. The
buffered-CCR path keyed on that tool merely being **present**, so once a
session went sticky, *every* later streaming turn was silently converted
to `stream: false`, buffered whole, and resynthesized as SSE:

```
CCR: stream:true request has headroom_retrieve available; using buffered stream:false upstream request
```

Buffering leaves time-to-last-byte roughly unchanged but makes
**time-to-first-byte the entire generation**. The reporter measured 8s
average and up to 100s across 234 requests in one day — turns that would
have streamed a first token in ~1s instead delivered nothing until done.

Retrieval can only expand a `<<ccr:...>>` marker present in the outgoing
body, so a turn carrying none cannot benefit from the buffered path at
all. Gate on that instead of on the tool.

This is also the root cause #3082 traced independently from the OpenCode
side — its plugin registers `headroom_retrieve` unconditionally, so
*every* turn buffered and neither `--no-ccr` nor `HEADROOM_NO_CCR`
stopped it.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- New `_outgoing_body_has_redeemable_marker()` scans the body **about to
go on the wire** and verifies ownership against the compression store —
the same `exists()` check the retrieve endpoint performs, so a
same-shaped marker from another context tool is not adopted (#2836).
Unexpected shapes answer `True`, keeping the long-standing behavior.
- The buffered-stream decision site gates on it, and logs at INFO when
it skips buffering.
- The correctness detail worth reviewing: the check reads `body`,
**not** the earlier `scan_for_markers(optimized_messages)` result.
`optimized_messages` is reassigned five times after that scan (memory
hooks, pre-send extensions, tool-search repair, CCR repair), so reusing
it would have been stale.
- Two existing test files encoded the very coupling this removes and had
to be repaired — see Testing.

Scope: this narrows *when* buffering happens; it does not make buffered
turns stream. A turn that genuinely carries a marker still loses
incremental delivery — restoring streaming there means wiring
`StreamingCCRHandler`, which is #3069's scope. It does not fix #3088
either, whose requests do carry markers.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] New tests added for new functionality
- [x] Manual testing performed

Two existing files built a request with `headroom_retrieve` and **no**
marker, relying on the tool alone to trigger buffering:

- `tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py` (11 tests)
fell through to the live streaming path, where only `_retry_request` is
stubbed — so the requests reached the network and the file **hung
indefinitely** rather than failing. Seeded real markers; now passes in
~5s.
-
`tests/test_proxy_response_cache_replay.py::test_buffered_ccr_turn_does_not_write_the_response_cache`
asserts its own premise (*"the conversion really happened — otherwise
this test proves nothing"*), so it failed loudly instead of passing
vacuously. Seeded a marker.

New `test_buffering_is_gated_on_a_redeemable_marker` pins all three
directions: owned marker → buffered, no marker → streaming, foreign
marker → streaming.

### Test Output

```text
$ pytest tests/test_ccr_buffered_stream_signed_thinking.py -q
8 passed, 1 warning in 4.45s

$ pytest tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py -q
11 passed, 1 warning in 4.94s        # was: hung indefinitely

$ pytest tests/test_proxy_response_cache_replay.py -q
9 passed, 1 warning in 1.69s

$ pytest tests/ -q
3 failed, 11151 passed, 581 skipped in 398.28s (0:06:38)

Same 3 failures as a clean-main baseline run on this machine:
  tests/test_graceful_shutdown.py::test_run_server_installs_cancelled_error_filter
  tests/test_learn/test_integration.py::TestCodexIntegration::test_full_pipeline
  tests/test_release_workflows.py::test_no_native_tls_in_wheel_build_tree

$ ruff check . && ruff format --check .
All checks passed!
```

## Real Behavior Proof

- Environment: this branch driven through the real FastAPI app with the
outbound HTTP client captured; macOS arm64, Python 3.12.
- Exact command / steps: posted a `stream: true` `/v1/messages` request
carrying a resident `headroom_retrieve` tool in three variants — no
marker, a marker seeded into the compression store, and a
correctly-shaped marker the store does not own — recording whether
`_retry_request` saw a `stream: false` body.
- Observed result: no marker → **streams**, `_retry_request` never sees
a flipped body; owned marker → **buffers**, exactly as before; foreign
marker → streams, honoring #2836 rather than adopting another tool's
hash.
- Not tested: the latency improvement against live client traffic. The
mechanism is verified (the buffered conversion no longer occurs), but
the reported 8s → ~1s TTFB needs the reporter's traffic to confirm.

## Runtime Rollout Safety

- Rollout-managed feature(s): none — this narrows an existing code path
and is not behind a rollout channel.
- Minimum rollout channel: n/a (ships to stable with the fix).
- Stable/default behavior changed: yes. A streaming turn whose body
carries no redeemable marker now stays streaming instead of being
buffered. Turns carrying a marker are unchanged.
- Kill switch / disable path: no new switch. Existing CCR controls still
apply — disabling the CCR response handler bypasses this decision site
entirely, and the helper fails open (returns `True`, i.e. the old
behavior) on any unexpected message shape.
- Unsafe override required: no.
- Qualification impact: none — no qualification-gated surface is
touched.
- Rollback path: revert this commit. Note it also carries two test
repairs; reverting the production change alone would leave those tests
passing but vacuous.

## 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
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`

## Additional Notes

Documentation update is marked N/A: no user-facing flag or endpoint
changes. Type checking (`mypy headroom`) was not run separately; `ruff`
is the gate this repo's CI enforces.

Related: #2836 (marker ownership), #3069 (streaming CCR handler), #3082,
#3088.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Tejas Chopra <tejas@Tejass-MacBook-Pro.local>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 11:19:26 -07:00

763 lines
29 KiB
Python

"""Regression tests for Anthropic streaming CCR retrieval interception."""
from __future__ import annotations
import asyncio
import json
import logging
from unittest.mock import AsyncMock, patch
import pytest
fastapi = pytest.importorskip("fastapi")
httpx = pytest.importorskip("httpx")
from fastapi.responses import StreamingResponse # noqa: E402
from fastapi.testclient import TestClient # noqa: E402
from starlette.requests import Request # noqa: E402
from headroom.cache.compression_store import get_compression_store # noqa: E402
from headroom.ccr.tool_injection import create_ccr_tool_definition # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
def _make_config() -> ProxyConfig:
return ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=True,
ccr_handle_responses=True,
ccr_context_tracking=False,
image_optimize=False,
)
@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 <<ccr:{hash_key}>>)"
def _message_response(content: list[dict], *, stop_reason: str = "end_turn") -> dict:
return {
"id": "msg_test",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": content,
"stop_reason": stop_reason,
"usage": {
"input_tokens": 10,
"output_tokens": 5,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
}
class _ContinuationClient:
def __init__(self, response_json: dict) -> None:
self.response_json = response_json
self.post_calls: list[dict] = []
async def post(self, url, *, content=None, headers=None, timeout=None): # noqa: ANN001
self.post_calls.append(
{
"url": url,
"content": content,
"headers": dict(headers or {}),
"timeout": timeout,
}
)
return httpx.Response(200, json=self.response_json)
async def aclose(self) -> None:
return None
def test_streaming_headroom_retrieve_is_intercepted_and_returned_as_sse() -> None:
config = _make_config()
store = get_compression_store()
hash_key = store.store(
original=json.dumps({"secret": "retrieved answer"}),
compressed="{}",
original_item_count=1,
)
initial_response = _message_response(
[
{
"type": "tool_use",
"id": "toolu_ccr",
"name": "headroom_retrieve",
"input": {"hash": hash_key},
}
],
stop_reason="tool_use",
)
final_response = _message_response(
[{"type": "text", "text": "retrieved answer is now available"}]
)
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
proxy._stream_response = AsyncMock(
side_effect=AssertionError("live streaming path should not be used")
)
continuation_client = _ContinuationClient(final_response)
proxy.http_client = continuation_client
initial_bodies: list[dict] = []
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
initial_bodies.append(json.loads(json.dumps(body)))
assert stream is False
assert body["stream"] is False
return httpx.Response(200, json=initial_response)
proxy._retry_request = _fake_retry # type: ignore[assignment]
resp = client.post(
"/v1/messages",
headers={
"x-api-key": "test-key",
"anthropic-version": "2023-06-01",
"accept": "text/event-stream",
"content-encoding": "identity",
},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("retrieve it")}],
},
)
assert resp.status_code == 200, resp.text
assert "text/event-stream" in resp.headers["content-type"]
assert "retrieved answer is now available" in resp.text
assert "headroom_retrieve" not in resp.text
assert initial_bodies and initial_bodies[0]["stream"] is False
assert len(continuation_client.post_calls) == 1
continuation_body = json.loads(continuation_client.post_calls[0]["content"].decode())
assert continuation_body["stream"] is False
continuation_headers = {
key.lower(): value for key, value in continuation_client.post_calls[0]["headers"].items()
}
assert "content-length" not in continuation_headers
assert "content-encoding" not in continuation_headers
assert "transfer-encoding" not in continuation_headers
assert "accept-encoding" not in continuation_headers
def test_streaming_without_headroom_retrieve_uses_normal_streaming_path() -> None:
config = _make_config()
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
async def _fake_stream_response(*args, **kwargs): # noqa: ANN001, ANN002, ANN003
async def _gen():
yield b"event: message_stop\n"
yield b'data: {"type":"message_stop"}\n\n'
return StreamingResponse(_gen(), media_type="text/event-stream")
proxy._stream_response = AsyncMock(side_effect=_fake_stream_response)
resp = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"messages": [{"role": "user", "content": "hello"}],
},
)
assert resp.status_code == 200, resp.text
assert "text/event-stream" in resp.headers["content-type"]
assert '"message_stop"' in resp.text
proxy._stream_response.assert_awaited_once()
def test_streaming_with_headroom_retrieve_available_but_unused_returns_sse() -> None:
config = _make_config()
text_response = _message_response([{"type": "text", "text": "plain answer"}])
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
proxy._stream_response = AsyncMock(
side_effect=AssertionError("live streaming path should not be used")
)
continuation_client = _ContinuationClient(_message_response([]))
proxy.http_client = continuation_client
initial_bodies: list[dict] = []
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
initial_bodies.append(json.loads(json.dumps(body)))
assert stream is False
assert body["stream"] is False
return httpx.Response(200, json=text_response)
proxy._retry_request = _fake_retry # type: ignore[assignment]
resp = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("hello")}],
},
)
assert resp.status_code == 200, resp.text
assert "text/event-stream" in resp.headers["content-type"]
assert "plain answer" in resp.text
assert "headroom_retrieve" not in resp.text
assert initial_bodies and initial_bodies[0]["stream"] is False
assert continuation_client.post_calls == []
proxy._stream_response.assert_not_awaited()
def test_mixed_ccr_and_client_tool_streams_both_blocks_as_sse() -> None:
"""LEGAL mixed turn (#839, #2089): headroom_retrieve emitted alongside a
client tool. The proxy cannot synthesize the client tool_result, so it must
hand the turn back for the client to resolve — a 200 SSE stream preserving
BOTH tool_use blocks, matching the non-streaming path. It must NOT 502 and
must NOT issue a continuation request."""
config = _make_config()
initial_response = _message_response(
[
{
"type": "tool_use",
"id": "toolu_ccr",
"name": "headroom_retrieve",
"input": {"hash": "abc123"},
},
{
"type": "tool_use",
"id": "toolu_client",
"name": "client_tool",
"input": {"value": 1},
},
],
stop_reason="tool_use",
)
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
proxy._stream_response = AsyncMock(
side_effect=AssertionError("live streaming path should not be used")
)
continuation_client = _ContinuationClient(_message_response([]))
proxy.http_client = continuation_client
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
assert body["stream"] is False
return httpx.Response(200, json=initial_response)
proxy._retry_request = _fake_retry # type: ignore[assignment]
resp = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [
create_ccr_tool_definition("anthropic"),
{
"name": "client_tool",
"description": "Client-owned tool",
"input_schema": {"type": "object", "properties": {}},
},
],
"messages": [{"role": "user", "content": _buffered("use tools")}],
},
)
assert resp.status_code == 200, resp.text
assert "text/event-stream" in resp.headers["content-type"]
# Both tool_use blocks are preserved for the client to resolve.
assert "headroom_retrieve" in resp.text
assert "client_tool" in resp.text
assert "toolu_ccr" in resp.text
assert "toolu_client" in resp.text
assert "Unable to safely complete streamed CCR retrieval" not in resp.text
# No continuation is issued — the client resolves all tool calls.
assert continuation_client.post_calls == []
def test_unresolved_ccr_only_streams_through_as_200() -> None:
"""CCR-only turn that never resolves: the model keeps re-emitting
headroom_retrieve so the continuation exhausts its retrieval rounds with a
residual marker and no accompanying client tool. Per #2089 the streaming
path streams the residual headroom_retrieve back as a 200 SSE so the client
can resolve or retry it, matching the non-streaming path."""
config = _make_config()
persistent_ccr = _message_response(
[
{
"type": "tool_use",
"id": "toolu_ccr",
"name": "headroom_retrieve",
"input": {"hash": "deadbeef"},
},
],
stop_reason="tool_use",
)
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
proxy._stream_response = AsyncMock(
side_effect=AssertionError("live streaming path should not be used")
)
# Every continuation re-emits headroom_retrieve, so it never resolves.
continuation_client = _ContinuationClient(persistent_ccr)
proxy.http_client = continuation_client
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
return httpx.Response(200, json=persistent_ccr)
proxy._retry_request = _fake_retry # type: ignore[assignment]
resp = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("use tools")}],
},
)
assert resp.status_code == 200, resp.text
assert "text/event-stream" in resp.headers["content-type"]
assert "headroom_retrieve" in resp.text
@pytest.mark.asyncio
async def test_buffered_ccr_withholds_output_until_delayed_upstream_resolves() -> None:
"""Nothing is sent — no status, no body — until the buffered result exists.
The response used to commit ``200 text/event-stream`` on a 1s keepalive
timer, which made every later failure unreportable: the client saw a 200
with no ``message_start`` and the real status was gone. See
``test_buffered_ccr_preserves_late_failure_status_and_headers``.
"""
config = _make_config()
final_response = _message_response([{"type": "text", "text": "done"}])
started = asyncio.Event()
release = asyncio.Event()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("wait")}],
}
request_delivered = False
async def receive():
# Mirror a real ASGI server: the body arrives once, then the channel
# stays open because the client is still connected. Returning instantly
# on every call spins `StreamingResponse.listen_for_disconnect`, which
# never yields, so the response body would never be scheduled.
nonlocal request_delivered
if request_delivered:
await asyncio.Event().wait()
request_delivered = True
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
async def delayed_retry(*args, **kwargs): # noqa: ANN002, ANN003
started.set()
await release.wait()
return httpx.Response(200, json=final_response)
proxy._retry_request = delayed_retry
task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive)))
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
response_task = asyncio.create_task(response(scope, receive, send))
# Longer than the deleted 1.0s keepalive deadline: an unresolved
# upstream must still have produced no ASGI message at all.
await asyncio.sleep(1.1)
assert events == []
release.set()
await response_task
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 200
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"event: ping" not in bodies
assert b"done" in bodies
@pytest.mark.asyncio
async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None:
config = _make_config()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("fail early")}],
}
async def receive():
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
async def early_failure(*args, **kwargs): # noqa: ANN002, ANN003
await asyncio.sleep(0.05)
return httpx.Response(
429,
headers={"retry-after": "7"},
json={"error": {"message": "slow down"}},
)
proxy._retry_request = early_failure
response = await proxy.handle_anthropic_messages(Request(scope, receive))
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
await response(scope, receive, send)
start = next(event for event in events if event["type"] == "http.response.start")
headers = dict(start["headers"])
assert start["status"] == 429
assert headers[b"retry-after"] == b"7"
assert b": headroom-keepalive\n\n" not in b"".join(
event["body"] for event in events if event["type"] == "http.response.body"
)
@pytest.mark.asyncio
async def test_buffered_ccr_preserves_late_failure_status_and_headers() -> None:
"""The reported failure: a non-200 landing after the old keepalive deadline.
Headroom had already committed ``200 text/event-stream`` by then, so the 429
reached Claude Code as a 200 whose body carried no ``message_start`` — shown
as "API returned an empty or malformed response (HTTP 200) — check for a
proxy or gateway intercepting the request" — and ``retry-after`` was dropped,
so the client never backed off.
"""
config = _make_config()
started = asyncio.Event()
release = asyncio.Event()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("fail late")}],
}
async def receive():
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003
started.set()
await release.wait()
return httpx.Response(
429,
headers={"retry-after": "7"},
json={"error": {"message": "slow down"}},
)
proxy._retry_request = delayed_failure
task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive)))
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
response_task = asyncio.create_task(response(scope, receive, send))
# Past the deleted 1.0s keepalive deadline before the upstream fails.
await asyncio.sleep(1.1)
assert events == []
release.set()
await response_task
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 429
assert dict(start["headers"])[b"retry-after"] == b"7"
assert b"slow down" in b"".join(
event["body"] for event in events if event["type"] == "http.response.body"
)
def test_buffered_ccr_rejects_malformed_success_as_502() -> None:
"""A non-SSE, non-JSON 200 is an upstream protocol error, not success."""
config = _make_config()
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = app.state.proxy
proxy._retry_request = AsyncMock(
return_value=httpx.Response(
200,
content=b"<html>gateway timeout</html>",
headers={"content-type": "text/html"},
)
)
response = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("fail safely")}],
},
)
assert response.status_code == 502
assert response.json()["error"]["type"] == "upstream_protocol_error"
assert b"gateway timeout" not in response.content
@pytest.mark.asyncio
async def test_buffered_ccr_late_failure_returns_sanitized_json_error() -> None:
"""A slow crash gets the same 502 the fast one does, not a downgraded 200."""
config = _make_config()
started = asyncio.Event()
release = asyncio.Event()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("wait")}],
}
async def receive():
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
proxy_logger = logging.getLogger("headroom.proxy")
error_records: list[logging.LogRecord] = []
log_handler = logging.Handler()
log_handler.setLevel(logging.ERROR)
log_handler.emit = error_records.append
proxy_logger.addHandler(log_handler)
async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003
started.set()
await release.wait()
raise RuntimeError("boom")
with patch.object(
proxy.metrics, "record_failed", new_callable=AsyncMock
) as record_failed:
proxy._retry_request = delayed_failure
task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive)))
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
response_task = asyncio.create_task(response(scope, receive, send))
await asyncio.sleep(0)
assert events == []
release.set()
await response_task
record_failed.assert_awaited_once_with(provider="anthropic")
proxy_logger.removeHandler(log_handler)
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 502
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"An error occurred while processing your request." in bodies
assert b"boom" not in bodies
assert events[-1]["more_body"] is False
assert any(
record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()
for record in error_records
)
@pytest.mark.asyncio
async def test_buffered_ccr_pre_keepalive_exception_returns_json_error() -> None:
config = _make_config()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": _buffered("fail before keepalive")}],
}
async def receive():
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
async def early_exception(*args, **kwargs): # noqa: ANN002, ANN003
raise RuntimeError("boom")
proxy._retry_request = early_exception
response = await proxy.handle_anthropic_messages(Request(scope, receive))
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
await response(scope, receive, send)
start = next(event for event in events if event["type"] == "http.response.start")
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
assert start["status"] == 502
assert dict(start["headers"])[b"content-type"] == b"application/json"
payload = json.loads(bodies[-1].decode())
assert (
payload["error"]["message"]
== "An error occurred while processing your request. Please try again."
)