akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): handle streaming CCR retrieval (#1451)

Browse source 
## Description

Fixes Anthropic-compatible streaming requests that can emit the internal
`headroom_retrieve` CCR tool. When a `stream: true` request includes the
CCR retrieve tool and response handling is enabled, Headroom now buffers
the upstream call as `stream: false`, lets the existing CCR response
handler retrieve and continue, and returns the final result as Anthropic
SSE so streaming clients do not see the internal tool call.

Closes #1450

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Detect direct Anthropic-compatible `stream: true` requests where
`headroom_retrieve` is available and CCR response handling is enabled.
- Route those requests through the existing buffered/non-stream CCR
response handler, then convert the final response back to
`text/event-stream`.
- Fail closed with a 502 SSE error if a buffered response still contains
`headroom_retrieve` after CCR handling, instead of leaking the internal
tool to the client.
- Preserve Anthropic `thinking`, `redacted_thinking`, signatures, and
citations when converting response JSON back to SSE.
- Add regression coverage for handled CCR retrieval, unused CCR tool
availability, normal streaming passthrough, mixed client/CCR tool
fail-closed behavior, and SSE conversion preservation.

## Testing

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

### Test Output

```text
$ rtk gh pr checks 1451 --repo headroomlabs-ai/headroom
CI Checks Summary:
  [ok] Passed: 20
  [FAIL] Failed: 0

Relevant CI commands from .github/workflows/ci.yml:
- ruff check .
- ruff format --check .
- mypy headroom --ignore-missing-imports
- pytest tests scripts/tests

$ rtk python3 -m py_compile headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/streaming.py tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_sse_thinking_blocks.py
# passed, no output

$ rtk pytest tests/test_sse_thinking_blocks.py -q
Pytest: 6 passed

$ MACOSX_DEPLOYMENT_TARGET=15.0 rtk uv run --python 3.13 pytest tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py -q
Failed before test collection while building the local editable package:
esaxx-rs build failed with fatal error: 'cstdint' file not found.
```

## Real Behavior Proof

- Environment: GitHub Actions CI on PR #1451 plus local macOS worktree
`fix/1450-ccr-streaming-retrieve`.
- Exact command / steps: CI ran lint, type checking, build, unit-test
shards, dashboard tests, extras tests, and e2e jobs; locally ran syntax
checks and the SSE conversion regression tests.
- Observed result: CI passed 20 checks with 0 failures; local syntax
checks passed; `tests/test_sse_thinking_blocks.py` passed with 6 tests.
- Not tested: the new proxy-level regression test was not run locally
because the local native extension build fails in `esaxx-rs` before
proxy tests can collect; it is included in the CI-tested suite.

## 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] I have made corresponding changes to the documentation
- [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 have updated the CHANGELOG.md if applicable

## Screenshots (if applicable)

N/A

## Additional Notes

- Scope: this handles the direct Anthropic-compatible HTTP
`/v1/messages` path. The configured Bedrock/backend streaming path does
not share this CCR continuation machinery in this PR.
- Documentation, CHANGELOG, code-comment, and local-full-test checklist
items are N/A for this narrow bug fix or not true locally.
This commit is contained in:
Vinay Gupta 2026-06-30 13:46:34 -05:00 • committed by GitHub
parent ddd4adf911
commit d337e3b828
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 511 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

128
headroom/proxy/handlers/anthropic.py
View file

@ -113,6 +113,21 @@ class AnthropicHandlerMixin:
canonical = str(tool)
return (name, canonical)
@staticmethod
def _has_headroom_retrieve_tool(tools: Any) -> bool:
"""Return True when the final Anthropic tool list includes CCR retrieve."""
if not isinstance(tools, list):
return False
for tool in tools:
if not isinstance(tool, dict):
continue
if tool.get("name") == "headroom_retrieve":
return True
function = tool.get("function")
if isinstance(function, dict) and function.get("name") == "headroom_retrieve":
return True
return False
@staticmethod
def _extract_anthropic_cache_ttl_metrics(usage: dict[str, Any] | None) -> tuple[int, int]:
"""Extract observed Anthropic cache-write TTL bucket usage.
@ -440,7 +455,7 @@ class AnthropicHandlerMixin:
self.pipeline_extensions = PipelineExtensionManager(discover=False)
from fastapi import HTTPException
from fastapi.responses import JSONResponse, Response
from fastapi.responses import JSONResponse, Response, StreamingResponse
from headroom.cache.compression_store import get_compression_store
from headroom.ccr import CCRToolInjector
@ -2148,7 +2163,30 @@ class AnthropicHandlerMixin:
url = f"{url}?{request.url.query}"
try:
if stream:
ccr_handler_config = getattr(self.ccr_response_handler, "config", None)
ccr_response_handler_enabled = bool(
self.ccr_response_handler and getattr(ccr_handler_config, "enabled", True)
)
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")
)
)
if buffered_stream_ccr:
if body.get("stream") is not False:
body["stream"] = False
body_mutation_tracker.mark_mutated(
"ccr_streaming_retrieve_buffered_non_stream"
)
logger.info(
f"[{request_id}] CCR: stream:true request has "
"headroom_retrieve available; using buffered stream:false "
"upstream request for server-side retrieval handling"
)
if stream and not buffered_stream_ccr:
self.pipeline_extensions.emit(
PipelineStage.POST_SEND,
operation="proxy.request",
@ -2220,6 +2258,8 @@ class AnthropicHandlerMixin:
metadata={
"path": pipeline_path,
"stream": False,
"client_stream": buffered_stream_ccr,
"ccr_stream_buffered": buffered_stream_ccr,
"status_code": response.status_code,
},
)
@ -2233,6 +2273,8 @@ class AnthropicHandlerMixin:
metadata={
"path": pipeline_path,
"stream": False,
"client_stream": buffered_stream_ccr,
"ccr_stream_buffered": buffered_stream_ccr,
"status_code": response.status_code,
},
)
@ -2750,16 +2792,88 @@ class AnthropicHandlerMixin:
content=json.dumps(resp_json).encode(),
headers=response_headers,
)
return Response(
content=response.content,
status_code=response.status_code,
headers=response_headers,
)
if not buffered_stream_ccr:
return Response(
content=response.content,
status_code=response.status_code,
headers=response_headers,
)
except Exception as sec_err:
logger.warning(
f"[{request_id}] Security response scan error: {sec_err}"
)
if buffered_stream_ccr and response.status_code == 200 and resp_json:
sse_headers = {
k: v
for k, v in response_headers.items()
if k.lower()
not in (
"content-encoding",
"content-length",
"transfer-encoding",
"content-type",
)
}
def _sse_error_event(message: str) -> bytes:
error_event = {
"type": "error",
"error": {"type": "api_error", "message": message},
}
return f"event: error\ndata: {json.dumps(error_event)}\n\n".encode()
if (
self.ccr_response_handler
and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "anthropic")
):
logger.warning(
f"[{request_id}] CCR: Buffered streaming response still "
"contains headroom_retrieve after handling; failing closed"
)
async def _residual_ccr_error_sse():
yield _sse_error_event(
"Unable to safely complete streamed CCR retrieval."
)
return StreamingResponse(
_residual_ccr_error_sse(),
media_type="text/event-stream",
headers=sse_headers,
status_code=502,
)
try:
sse_events = self._response_to_sse(resp_json, "anthropic")
except ValueError as sse_err:
logger.warning(
f"[{request_id}] CCR: Failed to convert buffered response "
f"to SSE: {sse_err}"
)
async def _conversion_error_sse():
yield _sse_error_event(
"Unable to safely convert buffered response to SSE."
)
return StreamingResponse(
_conversion_error_sse(),
media_type="text/event-stream",
headers=sse_headers,
status_code=502,
)
async def _buffered_ccr_sse():
for event in sse_events:
yield event
return StreamingResponse(
_buffered_ccr_sse(),
media_type="text/event-stream",
headers=sse_headers,
)
return Response(
content=response.content,
status_code=response.status_code,

54
headroom/proxy/handlers/streaming.py
View file

@ -523,8 +523,32 @@ class StreamingMixin:
"input": {},
},
}
elif block.get("type") == "thinking":
content_block = {
"type": "thinking",
"thinking": "",
}
if "signature" in block:
content_block["signature"] = block["signature"]
block_start = {
"type": "content_block_start",
"index": idx,
"content_block": content_block,
}
elif block.get("type") == "redacted_thinking":
block_start = {
"type": "content_block_start",
"index": idx,
"content_block": {
"type": "redacted_thinking",
"data": block.get("data", ""),
},
}
else:
continue
raise ValueError(
f"Unsupported Anthropic content block type for SSE conversion: "
f"{block.get('type')!r}"
)
events.append(
f"event: content_block_start\ndata: {json.dumps(block_start)}\n\n".encode()
@ -538,6 +562,15 @@ class StreamingMixin:
"delta": {"type": "text_delta", "text": block["text"]},
}
events.append(f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode())
for citation in block.get("citations", []) or []:
citation_delta = {
"type": "content_block_delta",
"index": idx,
"delta": {"type": "citations_delta", "citation": citation},
}
events.append(
f"event: content_block_delta\ndata: {json.dumps(citation_delta)}\n\n".encode()
)
elif block.get("type") == "tool_use" and block.get("input"):
delta = {
"type": "content_block_delta",
@ -548,6 +581,25 @@ class StreamingMixin:
},
}
events.append(f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode())
elif block.get("type") == "thinking":
if block.get("thinking"):
delta = {
"type": "content_block_delta",
"index": idx,
"delta": {"type": "thinking_delta", "thinking": block["thinking"]},
}
events.append(
f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode()
)
if block.get("signature"):
delta = {
"type": "content_block_delta",
"index": idx,
"delta": {"type": "signature_delta", "signature": block["signature"]},
}
events.append(
f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode()
)
# content_block_stop
block_stop = {"type": "content_block_stop", "index": idx}

283
tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py Normal file
View file

@ -0,0 +1,283 @@
"""Regression tests for Anthropic streaming CCR retrieval interception."""
from __future__ import annotations
import json
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 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,
)
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": "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": "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_does_not_issue_continuation() -> None:
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
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": "use tools"}],
},
)
assert resp.status_code == 502, resp.text
assert "text/event-stream" in resp.headers["content-type"]
assert "headroom_retrieve" not in resp.text
assert "client_tool" not in resp.text
assert "Unable to safely complete streamed CCR retrieval" in resp.text
assert continuation_client.post_calls == []

54
tests/test_sse_thinking_blocks.py
View file

@ -23,6 +23,8 @@ from __future__ import annotations
import json
from typing import Any
import pytest
from headroom.proxy.handlers.streaming import StreamingMixin
@ -182,3 +184,55 @@ def test_redacted_thinking_data_preserved() -> None:
# `data` field MUST be preserved byte-for-byte for signature
# validation on the next turn.
assert block["data"] == redacted_blob
def test_response_to_sse_preserves_thinking_redacted_and_citations() -> None:
parser = _Parser()
redacted_blob = "ENC:" + ("y" * 200)
response = {
"id": "msg_2",
"model": "claude-opus-4",
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "plan carefully", "signature": "sig_123"},
{
"type": "text",
"text": "Per source A",
"citations": [
{
"type": "page_location",
"cited_text": "abc",
"document_index": 0,
}
],
},
{"type": "redacted_thinking", "data": redacted_blob},
],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 3},
}
sse_text = b"".join(parser._response_to_sse(response, "anthropic")).decode("utf-8")
assert "thinking_delta" in sse_text
assert "signature_delta" in sse_text
assert "citations_delta" in sse_text
assert "redacted_thinking" in sse_text
assert redacted_blob in sse_text
round_tripped = parser._parse_sse_to_response(sse_text, "anthropic")
assert round_tripped is not None
assert round_tripped["content"][0]["thinking"] == "plan carefully"
assert round_tripped["content"][0]["signature"] == "sig_123"
assert round_tripped["content"][1]["citations"][0]["cited_text"] == "abc"
assert round_tripped["content"][2]["data"] == redacted_blob
def test_response_to_sse_rejects_unknown_content_block() -> None:
parser = _Parser()
with pytest.raises(ValueError, match="Unsupported Anthropic content block type"):
parser._response_to_sse(
{"content": [{"type": "future_block", "payload": "preserve me"}]},
"anthropic",
)