diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 1ad3853c5..64c643286 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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, diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 7292ea41e..3723e6352 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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} diff --git a/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py new file mode 100644 index 000000000..9951a08b6 --- /dev/null +++ b/tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py @@ -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 == [] diff --git a/tests/test_sse_thinking_blocks.py b/tests/test_sse_thinking_blocks.py index cf47e923c..762bf2c58 100644 --- a/tests/test_sse_thinking_blocks.py +++ b/tests/test_sse_thinking_blocks.py @@ -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", + )