From a2a7660ba8d1b78a91e897c0930f4d5467cbb96c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 17 Oct 2026 13:49:33 -0700 Subject: [PATCH] style(tests): restore LF line endings in test_headroom.py The file was saved with CRLF line endings, which made git report every existing line as changed and inflated the PR diff to about 6k lines. Converting back to LF leaves only the real change, 244 added lines and 1 removed, against main. --- .../guardrail_hooks/test_headroom.py | 6518 ++++++++--------- 1 file changed, 3259 insertions(+), 3259 deletions(-) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index b98dbe22284..e4dda5f561e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1,3259 +1,3259 @@ -""" -Unit tests for the Headroom guardrail. - -Tests cover: -- apply_guardrail compresses messages via /v1/compress and returns them as structured_messages -- x-headroom-bypass: true header causes guardrail to skip compression -- missing or empty messages are passed through unchanged -- response-type input is passed through unchanged -- /v1/compress HTTP error raises HTTPException (fail_closed, the default) -- /v1/compress returning malformed JSON raises HTTPException -- /v1/compress non-2xx surfaces as httpx.HTTPStatusError (raise_for_status), - not a status_code check on the returned response -- both are handled -- unreachable_fallback="fail_open" forwards the request uncompressed instead of raising -- tokens_saved is derived from tokens_before/tokens_after when the compression - service omits it, passed through verbatim when present, and skipped (without - breaking compression) when the token counts are not numeric -- CCR: headroom_retrieve tool injected when compressed messages contain hashes -- CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls -- CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages -- CCR: streaming /chat/completions is converted to a non-streaming call so the agentic - loop resolves the retrieve tool call, then fake-streamed back to the client -- session affinity: an explicit session id plus an authenticated tenant identity - produce a stable, tenant-scoped x-headroom-affinity header on /v1/compress and - /v1/retrieve, and the header is omitted whenever either half is missing -""" - -import json -import time -from typing import Final, Optional -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest -import respx -from fastapi import HTTPException - -import litellm - -from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( - AFFINITY_HEADER, - HeadroomGuardrail, - has_headroom_retrieve_tool, - HEADROOM_RETRIEVE_TOOL_NAME, -) -from litellm.proxy.spend_tracking.compression_savings import ( - extract_compression_saved_tokens, -) -from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY -from litellm.types.utils import ( - CallTypes, - GenericGuardrailAPIInputs, -) - -FAKE_API_BASE = "https://headroom.example.com" -FAKE_API_KEY = "test-key" - -# The system prompt, the last user turn and the last assistant turn are never -# sent to the compression service, so a fixture needs history for anything to -# be eligible: only index 1 is. -ORIGINAL_MESSAGES = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "A" * 5000}, - {"role": "assistant", "content": "Understood."}, - {"role": "user", "content": "and what about B?"}, -] -COMPRESSIBLE_MESSAGES = [ORIGINAL_MESSAGES[1]] -COMPRESSED_MESSAGES = [{"role": "user", "content": "A" * 500}] -COMPRESSED_MESSAGES_WITH_HASH = [ - { - "role": "user", - "content": "Summary. Retrieve more: hash=b573993006976af767214fac", - }, -] -EXPECTED_MESSAGES = [ - ORIGINAL_MESSAGES[0], - COMPRESSED_MESSAGES[0], - ORIGINAL_MESSAGES[2], - ORIGINAL_MESSAGES[3], -] -EXPECTED_MESSAGES_WITH_HASH = [ - ORIGINAL_MESSAGES[0], - COMPRESSED_MESSAGES_WITH_HASH[0], - ORIGINAL_MESSAGES[2], - ORIGINAL_MESSAGES[3], -] - - -def _make_guardrail(**kwargs) -> HeadroomGuardrail: - defaults = dict( - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - guardrail_name="headroom", - default_on=True, - ) - defaults.update(kwargs) - return HeadroomGuardrail(**defaults) - - -def _make_compress_response( - messages: list, - status: int = 200, - ccr_hashes: list[str] | None = None, -) -> MagicMock: - mock = MagicMock() - mock.status_code = status - mock.json.return_value = { - "messages": messages, - "tokens_before": 1000, - "tokens_after": 100, - "compression_ratio": 0.1, - "transforms_applied": ["router:smart_crusher:0.35"], - **({} if ccr_hashes is None else {"ccr_hashes": ccr_hashes}), - } - mock.text = "" - return mock - - -def _make_retrieve_response(original_content: str, status: int = 200) -> MagicMock: - mock = MagicMock() - mock.status_code = status - mock.json.return_value = {"original_content": original_content} - mock.text = original_content - return mock - - -def _make_openai_response_with_tool_call(tool_name: str, arguments: dict, tool_id: str = "call_abc123") -> MagicMock: - fn = MagicMock() - fn.name = tool_name - fn.arguments = json.dumps(arguments) - - tc = MagicMock() - tc.id = tool_id - tc.type = "function" - tc.function = fn - - message = MagicMock() - message.content = None - message.tool_calls = [tc] - - choice = MagicMock() - choice.message = message - - response = MagicMock() - response.choices = [choice] - return response - - -@pytest.fixture -def guardrail() -> HeadroomGuardrail: - return _make_guardrail() - - -def _recorded_guardrail_entries(request_data: dict) -> list: - for container_key in ("metadata", "litellm_metadata"): - container = request_data.get(container_key) - if isinstance(container, dict): - entries = container.get("standard_logging_guardrail_information") - if isinstance(entries, list): - return entries - return [] - - -def _applied_guardrails(request_data: dict) -> list: - for container_key in ("metadata", "litellm_metadata"): - container = request_data.get(container_key) - if isinstance(container, dict) and isinstance(container.get("applied_guardrails"), list): - return container["applied_guardrails"] - return [] - - -@pytest.mark.asyncio -async def test_apply_guardrail_compresses_and_returns_structured_messages( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data = {"model": "gpt-4o"} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - assert result.get("structured_messages") == EXPECTED_MESSAGES - - entries = _recorded_guardrail_entries(request_data) - assert len(entries) == 1 - assert entries[0]["guardrail_name"] == "headroom" - assert entries[0]["guardrail_status"] == "success" - assert entries[0]["guardrail_provider"] == "headroom" - assert "headroom" in _applied_guardrails(request_data) - - -@pytest.mark.asyncio -async def test_apply_guardrail_leaves_background_requests_uncompressed( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - request_data = {"model": "gpt-4o", "background": True} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES), - ) as post: - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - assert result is inputs - post.assert_not_awaited() - assert _recorded_guardrail_entries(request_data) == [] - - -def _recorded_guardrail_response(request_data: dict) -> dict: - entries = request_data["metadata"]["standard_logging_guardrail_information"] - assert len(entries) == 1 - return entries[0]["guardrail_response"] - - -@pytest.mark.asyncio -async def test_apply_guardrail_derives_tokens_saved_when_service_omits_it( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - # _make_compress_response omits tokens_saved, matching the live service. - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data: dict = {"model": "gpt-4o"} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - stats = _recorded_guardrail_response(request_data) - assert stats["tokens_saved"] == 900 - - # Spend tracking reads the entry under the spend-log metadata key. - entry = request_data["metadata"]["standard_logging_guardrail_information"][0] - assert extract_compression_saved_tokens({"guardrail_information": [entry]}) == 900 - - -@pytest.mark.asyncio -async def test_apply_guardrail_passes_through_service_sent_tokens_saved( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - # Deliberately different from tokens_before - tokens_after (900): the - # service-sent value must win over the derived one. - mock_response.json.return_value["tokens_saved"] = 123 - request_data: dict = {"model": "gpt-4o"} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - assert _recorded_guardrail_response(request_data)["tokens_saved"] == 123 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "tokens_before, tokens_after", - [ - ("1000", "100"), - (True, False), - (None, None), - ], -) -async def test_apply_guardrail_skips_derivation_for_non_numeric_token_counts( - guardrail: HeadroomGuardrail, - tokens_before, - tokens_after, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - mock_response.json.return_value["tokens_before"] = tokens_before - mock_response.json.return_value["tokens_after"] = tokens_after - request_data: dict = {"model": "gpt-4o"} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - assert "tokens_saved" not in _recorded_guardrail_response(request_data) - # Compression itself is unaffected by the skipped derivation. - assert result.get("structured_messages") == EXPECTED_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - tools = result.get("tools") - assert tools is not None - assert has_headroom_retrieve_tool(tools) - - -@pytest.mark.asyncio -async def test_apply_guardrail_no_tool_injected_when_no_hashes( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - tools = result.get("tools") - assert not has_headroom_retrieve_tool(tools or []) - - -@pytest.mark.asyncio -async def test_apply_guardrail_preserves_existing_tools_when_injecting( - guardrail: HeadroomGuardrail, -): - existing_tool = {"type": "function", "function": {"name": "my_tool", "parameters": {}}} - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - tools=[existing_tool], - ) - mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - tools = result.get("tools") - assert tools is not None - assert isinstance(tools, list) - assert any(isinstance(t, dict) and t.get("function", {}).get("name") == "my_tool" for t in tools) - assert has_headroom_retrieve_tool(tools) - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_returns_true_for_retrieve_call( - guardrail: HeadroomGuardrail, -): - retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] - response = _make_openai_response_with_tool_call( - tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "b573993006976af767214fac"}, - ) - - should_run, ctx = await guardrail.async_should_run_agentic_loop( - response=response, - model="gpt-4o", - messages=[], - tools=retrieve_tool_def, - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) - - assert should_run is True - assert len(ctx["tool_calls"]) == 1 - assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_returns_false_without_retrieve_tool( - guardrail: HeadroomGuardrail, -): - other_tools = [{"type": "function", "function": {"name": "other_tool"}}] - response = _make_openai_response_with_tool_call( - tool_name="other_tool", - arguments={}, - ) - - should_run, _ = await guardrail.async_should_run_agentic_loop( - response=response, - model="gpt-4o", - messages=[], - tools=other_tools, - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) - - assert should_run is False - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_returns_false_when_no_retrieve_calls( - guardrail: HeadroomGuardrail, -): - retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] - response = _make_openai_response_with_tool_call( - tool_name="some_other_function", - arguments={}, - ) - - should_run, _ = await guardrail.async_should_run_agentic_loop( - response=response, - model="gpt-4o", - messages=[], - tools=retrieve_tool_def, - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) - - assert should_run is False - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( - guardrail: HeadroomGuardrail, -): - original_content = "This is the full compressed content." - mock_retrieve = _make_retrieve_response(original_content) - - # Registered hashes are lowercase; a model may echo the marker's hex in uppercase. - tool_calls = [ - { - "id": "call_abc123", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "B573993006976AF767214FAC"}, - } - ] - response = _make_openai_response_with_tool_call( - tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "B573993006976AF767214FAC"}, - tool_id="call_abc123", - ) - messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=mock_retrieve, - ) as mock_get: - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="gpt-4o", - messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-1"}, - ) - - assert plan.run_agentic_loop is True - assert plan.request_patch is not None - - follow_up = plan.request_patch.messages - assert follow_up is not None - - tool_result_message = next((m for m in follow_up if m.get("role") == "tool"), None) - assert tool_result_message is not None - assert tool_result_message["content"] == original_content - assert tool_result_message["tool_call_id"] == "call_abc123" - - mock_get.assert_called_once() - call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] - assert "b573993006976af767214fac" in call_url - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_handles_retrieve_404( - guardrail: HeadroomGuardrail, -): - mock_retrieve = MagicMock() - mock_retrieve.status_code = 404 - - tool_calls = [ - { - "id": "call_xyz", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "deadbeef000000000000dead"}, - } - ] - response = _make_openai_response_with_tool_call( - tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "deadbeef000000000000dead"}, - tool_id="call_xyz", - ) - - messages = [ - { - "role": "user", - "content": "Retrieve more: hash=deadbeef000000000000dead", - } - ] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"deadbeef000000000000dead"}), - time.monotonic() + 999, - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=mock_retrieve, - ): - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="gpt-4o", - messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-1"}, - ) - - follow_up = plan.request_patch.messages # type: ignore[union-attr] - tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) - assert tool_result is not None - assert "not found" in tool_result["content"] or "expired" in tool_result["content"] - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, -): - """A hash-shaped string planted in message text must not be honored when - this guardrail has no record of ever issuing it, even if it's echoed back - in the current request's own messages (e.g. via prompt injection).""" - tool_calls = [ - { - "id": "call_xyz", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "deadbeef000000000000dead"}, - } - ] - response = _make_openai_response_with_tool_call( - tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "deadbeef000000000000dead"}, - tool_id="call_xyz", - ) - assert not guardrail._issued_hashes_by_call_id - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - ) as mock_get: - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="gpt-4o", - messages=[{"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me"}], - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-unknown"}, - ) - - mock_get.assert_not_called() - - follow_up = plan.request_patch.messages # type: ignore[union-attr] - tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) - assert tool_result is not None - assert "was not produced by the current request" in tool_result["content"] - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( - guardrail: HeadroomGuardrail, -): - """A hash issued for one request must not be retrievable by a different - request just because the second request echoes that hash-shaped string - back in its own messages -- retrieval must be scoped per litellm_call_id, - not derived by re-scanning attacker-controlled message text.""" - guardrail._issued_hashes_by_call_id["call-A"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) - - tool_calls = [ - { - "id": "call_xyz", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "b573993006976af767214fac"}, - } - ] - response = _make_openai_response_with_tool_call( - tool_name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments={"hash": "b573993006976af767214fac"}, - tool_id="call_xyz", - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - ) as mock_get: - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="gpt-4o", - messages=[{"role": "user", "content": "Please fetch hash=b573993006976af767214fac for me"}], - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-B"}, - ) - - mock_get.assert_not_called() - - follow_up = plan.request_patch.messages # type: ignore[union-attr] - tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) - assert tool_result is not None - assert "was not produced by the current request" in tool_result["content"] - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_items( - guardrail: HeadroomGuardrail, -): - """For the Responses API, follow-up input must echo a function_call paired - with a function_call_output keyed by the same call_id -- chat-style - assistant/tool messages are not valid Responses API input items.""" - original_content = "This is the full compressed content." - mock_retrieve = _make_retrieve_response(original_content) - - response = MagicMock() - response.choices = None - response.content = None - response.output = [ - { - "type": "function_call", - "id": "fc_abc123", - "call_id": "call_abc123", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": "b573993006976af767214fac"}), - } - ] - - tool_calls = [ - { - "id": "call_abc123", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "b573993006976af767214fac"}, - } - ] - messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=mock_retrieve, - ): - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="gpt-4o", - messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-1"}, - ) - - follow_up = plan.request_patch.messages # type: ignore[union-attr] - assert all("role" not in item for item in follow_up if item not in messages) - - function_call_item = next((i for i in follow_up if i.get("type") == "function_call"), None) - assert function_call_item is not None - assert function_call_item["call_id"] == "call_abc123" - assert function_call_item["name"] == HEADROOM_RETRIEVE_TOOL_NAME - - output_item = next((i for i in follow_up if i.get("type") == "function_call_output"), None) - assert output_item is not None - assert output_item["call_id"] == "call_abc123" - assert output_item["output"] == original_content - - -@pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messages( - guardrail: HeadroomGuardrail, -): - """For the Anthropic Messages API, follow-up must echo a tool_use content - block in an assistant message paired with a tool_result content block in a - user message keyed by the same tool_use_id -- chat-style tool-role - messages are not valid Anthropic input. - - AnthropicMessagesResponse is a TypedDict, so real responses are plain - dicts at runtime; a MagicMock response here would pass even if branch - selection used bare getattr() and silently fell through to the - chat-completions replay shape for every real Anthropic response. - """ - original_content = "This is the full compressed content." - mock_retrieve = _make_retrieve_response(original_content) - - response = { - "content": [ - { - "type": "tool_use", - "id": "toolu_abc123", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input": {"hash": "b573993006976af767214fac"}, - } - ] - } - - tool_calls = [ - { - "id": "toolu_abc123", - "type": "function", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": {"hash": "b573993006976af767214fac"}, - } - ] - messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=mock_retrieve, - ): - plan = await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": tool_calls}, - model="claude-sonnet-4-5", - messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs={"litellm_call_id": "call-1"}, - ) - - follow_up = plan.request_patch.messages # type: ignore[union-attr] - assert all(m.get("role") != "tool" for m in follow_up) - - assistant_message = next((m for m in follow_up if m.get("role") == "assistant"), None) - assert assistant_message is not None - tool_use_block = next((b for b in assistant_message["content"] if b.get("type") == "tool_use"), None) - assert tool_use_block is not None - assert tool_use_block["id"] == "toolu_abc123" - - user_message = follow_up[-1] - assert user_message["role"] == "user" - tool_result_block = next((b for b in user_message["content"] if b.get("type") == "tool_result"), None) - assert tool_result_block is not None - assert tool_result_block["tool_use_id"] == "toolu_abc123" - assert tool_result_block["content"] == original_content - - -HASH_SHAPED_HISTORY = [ - {"role": "system", "content": "You are Claude Code."}, - {"role": "user", "content": [{"type": "text", "text": "Run git log."}]}, - { - "role": "assistant", - "content": [{"type": "text", "text": "Done."}], - "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "tu_1", "content": "hash=3f2a9c1d7e5b4a6f8c0d1e2f9a8b7c6d5e4f3a2b"}, - {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me."}, -] - - -async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: - request_data = {"model": "claude-sonnet-5"} - - def _echo(**kwargs): - return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): - return await guardrail.apply_guardrail( - inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), - request_data=request_data, - input_type="request", - ) - - -@pytest.mark.asyncio -async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail: HeadroomGuardrail): - """Regression for LIT-7086: a git SHA in a tool result and a hash= string the - caller typed both look like markers, but the service stored nothing, so the - tool must not appear and no hash may be registered as issued. Covers a - service that omits ccr_hashes, returns it empty, or returns a non-list.""" - for ccr_hashes in (None, [], "b573993006976af767214fac"): - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) - - assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id - - -@pytest.mark.asyncio -async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardrail: HeadroomGuardrail): - """Only the hashes the service reports in ccr_hashes are honored, in the - service's own 12 to 24 hex grammar; anything else is dropped because each - entry is interpolated into the /v1/retrieve URL.""" - result = await _apply( - guardrail, - HASH_SHAPED_HISTORY, - ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], - ) - - assert has_headroom_retrieve_tool(result.get("tools") or []) - ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() - assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) - - -@pytest.mark.asyncio -async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatch: pytest.MonkeyPatch): - """`ccr_retrieval: false` in config.yaml compresses without any retrieval - round trip, so it has to reach the instance through the initializer.""" - from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail - from litellm.types.guardrails import LitellmParams - - monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) - params = LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, ccr_retrieval=False) - guardrail = initialize_guardrail(params, {"guardrail_name": "headroom", "litellm_params": params}) - - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["b573993006976af767214fac"]) - - assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] - assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id - - -@pytest.mark.asyncio -async def test_anthropic_assistant_history_never_reaches_compression_service(guardrail: HeadroomGuardrail): - """The public Anthropic handler translates assistant content blocks to a - string before Headroom sees them, so model-authored rows must be excluded - from the compression payload rather than protected by their content shape.""" - from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler - - table = "| Guardrail | Model |\n|---|---|\n" + "\n".join(f"| gr-{i} | model-{i} |" for i in range(40)) - messages = [ - {"role": "user", "content": [{"type": "text", "text": "List the guardrails."}]}, - {"role": "assistant", "content": [{"type": "text", "text": table}]}, - {"role": "user", "content": [{"type": "text", "text": "Earlier follow-up. " + "B" * 5000}]}, - {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]}, - {"role": "user", "content": "Re-print the table."}, - ] - sent: dict = {} - - def _echo(**kwargs): - sent["messages"] = kwargs["json"]["messages"] - return _make_compress_response(json.loads(json.dumps(sent["messages"]))) - - data = {"model": "claude-sonnet-5", "messages": json.loads(json.dumps(messages))} - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): - result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) - - assert [row["role"] for row in sent["messages"]] == ["user", "user"] - assert sent["messages"][1]["content"] == "Earlier follow-up. " + "B" * 5000 - assert table not in json.dumps(sent["messages"]) - assert result["messages"][1]["content"] == [{"type": "text", "text": table}] - - -def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape(): - """By the time an Anthropic Messages API response reaches the agentic-loop - gate, the OpenAI-shaped tool this guardrail injects (type: "function") - has already been transformed into Anthropic's native tool shape - (type: "custom", top-level "name", no nested "function" object).""" - anthropic_native_tools = [ - { - "type": "custom", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, - } - ] - assert has_headroom_retrieve_tool(anthropic_native_tools) - assert not has_headroom_retrieve_tool([{"type": "custom", "name": "some_other_tool"}]) - - -@pytest.mark.asyncio -async def test_apply_guardrail_bypass_header_skips_compression( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - request_data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "true"}}} - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - mock_post.assert_not_called() - - assert result.get("structured_messages") == ORIGINAL_MESSAGES - assert _recorded_guardrail_entries(request_data) == [] - - -@pytest.mark.asyncio -async def test_apply_guardrail_response_type_passthrough( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["some response text"], - structured_messages=ORIGINAL_MESSAGES, - ) - request_data: dict = {"model": "gpt-4o"} - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="response", - ) - mock_post.assert_not_called() - - assert result is inputs - assert _recorded_guardrail_entries(request_data) == [] - - -@pytest.mark.asyncio -async def test_apply_guardrail_empty_structured_messages_passthrough( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs(texts=["hello"]) - request_data: dict = {"model": "gpt-4o"} - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - mock_post.assert_not_called() - - assert result is inputs - assert _recorded_guardrail_entries(request_data) == [] - assert "headroom" not in _applied_guardrails(request_data) - - -@pytest.mark.asyncio -async def test_passthrough_handler_does_not_log_headroom_as_run( - guardrail: HeadroomGuardrail, -): - """Regression for LIT-4650. - - A passthrough request drives headroom through PassThroughEndpointHandler, which - only supplies `texts` (no `structured_messages`). Headroom cannot compress that - shape and no-ops, so it must not appear in the spend log's - standard_logging_guardrail_information as a successful run. - """ - from litellm.llms.pass_through.guardrail_translation.handler import ( - PassThroughEndpointHandler, - ) - - data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: - await PassThroughEndpointHandler().process_input_messages( - data=data, - guardrail_to_apply=guardrail, - litellm_logging_obj=None, - ) - mock_post.assert_not_called() - - assert _recorded_guardrail_entries(data) == [] - assert "headroom" not in _applied_guardrails(data) - - -@pytest.mark.asyncio -async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream( - guardrail: HeadroomGuardrail, -): - """Regression for LIT-6494: on /v1/responses the compressed messages must be - written back into `input`, not only the retrieve tool into `tools`, or the - model keeps reading the full document and never calls headroom_retrieve.""" - from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler - - data = { - "model": "gpt-5.6", - "instructions": ORIGINAL_MESSAGES[0]["content"], - "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]], - "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}], - } - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), - ): - result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) - - assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] - assert [item["content"] for item in result["input"]] == [ - COMPRESSED_MESSAGES_WITH_HASH[0]["content"], - ORIGINAL_MESSAGES[2]["content"], - ORIGINAL_MESSAGES[3]["content"], - ] - assert "A" * 5000 not in json.dumps(result["input"]) - assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME] - - -@pytest.mark.asyncio -async def test_apply_guardrail_http_error_raises(): - guardrail = _make_guardrail() - mock_response = _make_compress_response([], status=500) - mock_response.text = "Internal Server Error" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - - -@pytest.mark.asyncio -async def test_apply_guardrail_transport_error_raises(): - guardrail = _make_guardrail() - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ConnectError("Connection refused"), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert "unreachable" in str(exc_info.value.detail) - - -def _make_http_status_error(status: int, body: str) -> httpx.HTTPStatusError: - request = httpx.Request("POST", f"{FAKE_API_BASE}/v1/compress") - response = httpx.Response(status, request=request, text=body) - return httpx.HTTPStatusError( - f"Server error '{status}' for url", - request=request, - response=response, - ) - - -@pytest.mark.asyncio -async def test_apply_guardrail_http_status_error_raises(): - """Regression test: litellm's async httpx client calls raise_for_status() - internally, so a non-2xx /v1/compress response surfaces as - httpx.HTTPStatusError, not as a returned MagicMock with status_code set. - A prior version of _call_compress only checked response.status_code and - never caught this exception, so it went unhandled instead of blocking - the request per fail_closed policy.""" - guardrail = _make_guardrail() - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=_make_http_status_error(500, "headroom internal error"), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - - -@pytest.mark.asyncio -async def test_apply_guardrail_404_error_includes_troubleshooting_hint(): - """404 responses include a troubleshooting hint for self-hosted Headroom deployments.""" - guardrail = _make_guardrail() - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=_make_http_status_error(404, "Not Found"), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert exc_info.value.detail["status_code"] == 404 - assert exc_info.value.detail["body"] == "Not Found" - assert "hint" in exc_info.value.detail - assert "HEADROOM_COMPRESS_ALLOW_REMOTE=1" in exc_info.value.detail["hint"] - - -@pytest.mark.asyncio -async def test_apply_guardrail_non_404_error_omits_troubleshooting_hint(): - guardrail = _make_guardrail() - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=_make_http_status_error(500, "headroom internal error"), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert exc_info.value.detail["status_code"] == 500 - assert exc_info.value.detail["body"] == "headroom internal error" - assert "hint" not in exc_info.value.detail - - -@pytest.mark.asyncio -async def test_apply_guardrail_http_status_error_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=_make_http_status_error(500, "headroom internal error"), - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_transport_error_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - request_data = {"model": "gpt-4o"} - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ConnectError("Connection refused"), - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - entries = _recorded_guardrail_entries(request_data) - assert len(entries) == 1 - assert entries[0]["guardrail_name"] == "headroom" - assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" - assert "headroom" in _applied_guardrails(request_data) - - -@pytest.mark.asyncio -async def test_apply_guardrail_http_error_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - mock_response = _make_compress_response([], status=500) - mock_response.text = "Internal Server Error" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_non_json_response_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.side_effect = ValueError("not JSON") - mock_response.text = "not json" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_missing_messages_key_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = {"tokens_before": 100, "tokens_after": 10} - mock_response.text = "{}" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_empty_compressed_messages_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = { - "messages": ["not-a-dict", 42, None], - "tokens_before": 1000, - "tokens_after": 0, - "compression_ratio": 0, - } - mock_response.text = "{}" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -@pytest.mark.asyncio -async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_messages(): - """When compression fails with fail_open, user-supplied messages that - happen to contain hash-shaped strings must NOT cause those hashes to be - registered as valid for CCR retrieval. Otherwise an attacker can plant a - hash= string in their prompt, trigger a compression failure, and have - that hash honored by a later headroom_retrieve tool call.""" - messages_with_fake_hash = [ - {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me"}, - ] - guardrail = _make_guardrail(unreachable_fallback="fail_open") - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=messages_with_fake_hash, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ConnectError("Connection refused"), - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == messages_with_fake_hash - assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id - - -@pytest.mark.asyncio -async def test_apply_guardrail_missing_messages_key_raises(): - guardrail = _make_guardrail() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = {"tokens_before": 100, "tokens_after": 10} - mock_response.text = "{}" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - - -@pytest.mark.asyncio -async def test_apply_guardrail_empty_compressed_messages_raises(): - guardrail = _make_guardrail() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = { - "messages": ["not-a-dict", 42, None], - "tokens_before": 1000, - "tokens_after": 0, - "compression_ratio": 0, - } - mock_response.text = "{}" - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert "empty message list" in str(exc_info.value.detail) - - -def test_init_raises_without_api_base(): - with pytest.raises(ValueError, match="API base URL"): - HeadroomGuardrail(api_base=None) - - -def test_init_defaults_to_fail_closed(): - guardrail = _make_guardrail() - assert guardrail.unreachable_fallback == "fail_closed" - - -def test_init_rejects_invalid_unreachable_fallback_value(): - guardrail = _make_guardrail(unreachable_fallback="not-a-real-mode") - assert guardrail.unreachable_fallback == "fail_closed" - - -def test_bypass_header_case_insensitive(): - guardrail = _make_guardrail() - - for header_value in ("true", "True", "TRUE"): - data = {"proxy_server_request": {"headers": {"x-headroom-bypass": header_value}}} - assert guardrail._should_bypass(data) is True - - data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "false"}}} - assert guardrail._should_bypass(data) is False - - data = {"proxy_server_request": {"headers": {}}} - assert guardrail._should_bypass(data) is False - - data = {} - assert guardrail._should_bypass(data) is False - - -@pytest.mark.asyncio -async def test_apply_guardrail_sends_model_from_config(): - guardrail = _make_guardrail(model="gpt-4o-mini") - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_post: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - call_kwargs = mock_post.call_args - sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] - assert sent_payload.get("model") == "gpt-4o-mini" - - -@pytest.mark.asyncio -async def test_apply_guardrail_sends_model_from_request_data_when_no_config_model(): - guardrail = _make_guardrail() - mock_response = _make_compress_response(COMPRESSED_MESSAGES) - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_post: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - call_kwargs = mock_post.call_args - sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] - assert sent_payload.get("model") == "gpt-4o" - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_detects_anthropic_content_block_format( - guardrail: HeadroomGuardrail, -): - # Anthropic's native tool format (type: "custom", top-level "name") -- - # by the time a Messages API response reaches this gate, the OpenAI-shaped - # tool this guardrail injects has already been transformed into this shape. - retrieve_tool_def = [ - { - "type": "custom", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, - } - ] - - response = MagicMock() - response.choices = None - response.content = [ - { - "type": "tool_use", - "id": "toolu_abc", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input": {"hash": "b573993006976af767214fac"}, - } - ] - - should_run, ctx = await guardrail.async_should_run_agentic_loop( - response=response, - model="claude-sonnet-4-6", - messages=[], - tools=retrieve_tool_def, - stream=False, - custom_llm_provider="anthropic", - kwargs={}, - ) - - assert should_run is True - assert len(ctx["tool_calls"]) == 1 - assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_detects_anthropic_response_as_plain_dict( - guardrail: HeadroomGuardrail, -): - """AnthropicMessagesResponse is a TypedDict -- real Messages API responses - are plain dicts at runtime, not objects with attribute access. A - MagicMock-only test would pass even if detection used bare getattr() and - silently treated every real response as having no tool calls.""" - retrieve_tool_def = [ - { - "type": "custom", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, - } - ] - response = { - "content": [ - { - "type": "tool_use", - "id": "toolu_abc", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "input": {"hash": "b573993006976af767214fac"}, - } - ] - } - - should_run, ctx = await guardrail.async_should_run_agentic_loop( - response=response, - model="claude-sonnet-4-6", - messages=[], - tools=retrieve_tool_def, - stream=False, - custom_llm_provider="anthropic", - kwargs={}, - ) - - assert should_run is True - assert len(ctx["tool_calls"]) == 1 - assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" - - -@pytest.mark.asyncio -async def test_async_should_run_agentic_loop_detects_responses_api_output_format( - guardrail: HeadroomGuardrail, -): - retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] - - response = MagicMock() - response.choices = None - response.content = None - response.output = [ - { - "type": "function_call", - "id": "fc_abc123", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": "b573993006976af767214fac"}), - } - ] - - should_run, ctx = await guardrail.async_should_run_agentic_loop( - response=response, - model="gpt-4o", - messages=[], - tools=retrieve_tool_def, - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) - - assert should_run is True - assert len(ctx["tool_calls"]) == 1 - assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" - - -@pytest.mark.asyncio -async def test_apply_guardrail_litellm_timeout_raises_when_fail_closed(): - guardrail = _make_guardrail() - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=litellm.Timeout( - message="Connection timed out after 10 seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert "unreachable" in str(exc_info.value.detail) - - -@pytest.mark.asyncio -async def test_apply_guardrail_litellm_timeout_fail_open_forwards_uncompressed(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - - inputs = GenericGuardrailAPIInputs( - texts=["hello"], - structured_messages=ORIGINAL_MESSAGES, - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=litellm.Timeout( - message="Connection timed out after 10 seconds.", - model="default-model-name", - llm_provider="litellm-httpx-handler", - ), - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - assert result["structured_messages"] == ORIGINAL_MESSAGES - - -# --------------------------------------------------------------------------- -# Content-parts flattening (LIT-4795) -# -# Anthropic-format requests translate to messages whose content is a list of -# part dicts. The compression service only rewrites string content, so the -# guardrail flattens ALL-TEXT part lists on the wire and restores the -# original shapes afterwards. Rows with non-text parts are never flattened: -# cache_control breakpoints are positional, and merging text across a -# non-text part would move a later breakpoint to the other side of it. -# --------------------------------------------------------------------------- - -PARTS_MESSAGES = [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "Earlier turn."}, - {"type": "text", "text": "Second block. " + "B" * 5000}, - ], - }, - { - "role": "user", - "content": [ - {"type": "text", "text": "Mixed row text."}, - {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}, - ], - }, - {"role": "tool", "content": "tool output " + "C" * 500}, - {"role": "user", "content": "what does that file do?"}, -] - -FLATTENED_HISTORY_TEXT = "Earlier turn.\n\nSecond block. " + "B" * 5000 - - -def _parts_copy() -> list: - return json.loads(json.dumps(PARTS_MESSAGES)) - - -def _echo_wire_view() -> list: - """What the service receives (and echoes back when it changes nothing). - - The system row and the trailing user row are never sent. - """ - return [ - {"role": "user", "content": FLATTENED_HISTORY_TEXT}, - json.loads(json.dumps(PARTS_MESSAGES[2])), - {"role": "tool", "content": "tool output " + "C" * 500}, - ] - - -@pytest.mark.asyncio -async def test_apply_guardrail_flattens_all_text_rows_only( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["B" * 5000], - structured_messages=_parts_copy(), - ) - mock_response = _make_compress_response(_echo_wire_view()) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_post: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - wire_messages = mock_post.call_args.kwargs["json"]["messages"] - assert wire_messages[0]["content"] == FLATTENED_HISTORY_TEXT - # Mixed text+image row is never flattened: merging its text would move a - # later cache_control breakpoint across the image part. - assert isinstance(wire_messages[1]["content"], list) - assert wire_messages[2]["content"] == "tool output " + "C" * 500 - - -@pytest.mark.asyncio -async def test_apply_guardrail_restores_rewritten_all_text_row( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["B" * 5000], - structured_messages=_parts_copy(), - ) - compressed = _echo_wire_view() - compressed[0]["content"] = "compressed history. Retrieve more: hash=b573993006976af767214fac" - mock_response = _make_compress_response(compressed, ccr_hashes=["b573993006976af767214fac"]) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - messages = result["structured_messages"] - history_content = messages[1]["content"] - assert isinstance(history_content, list) - assert len(history_content) == 1 - assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" - # Mixed row passes through byte-identical. - assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] - # The service-declared hash still drives retrieve-tool injection on a restored row. - assert has_headroom_retrieve_tool(result.get("tools") or []) - - -@pytest.mark.asyncio -async def test_apply_guardrail_keeps_originals_when_service_echoes_unchanged( - guardrail: HeadroomGuardrail, -): - inputs = GenericGuardrailAPIInputs( - texts=["B" * 5000], - structured_messages=_parts_copy(), - ) - mock_response = _make_compress_response(_echo_wire_view()) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - messages = result["structured_messages"] - assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] - - -@pytest.mark.asyncio -async def test_apply_guardrail_rejects_service_output_when_rows_dropped( - guardrail: HeadroomGuardrail, -): - """A reshaped conversation cannot be applied at all: the rows held back from - compression are matched positionally, so a response with a different row - count goes through the fail policy instead of being adopted.""" - inputs = GenericGuardrailAPIInputs( - texts=["B" * 5000], - structured_messages=_parts_copy(), - ) - dropped = [{"role": "user", "content": "B" * 50}] - mock_response = _make_compress_response(dropped) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - assert exc_info.value.status_code == 502 - assert "changed the message count" in str(exc_info.value.detail) - - -@pytest.mark.asyncio -async def test_apply_guardrail_forwards_original_when_rows_dropped_and_fail_open(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - original = _parts_copy() - inputs = GenericGuardrailAPIInputs(texts=["B" * 5000], structured_messages=original) - mock_response = _make_compress_response([{"role": "user", "content": "B" * 50}]) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - # Same object back, so translation handlers that detect a rewrite by - # identity leave the request alone instead of round-tripping it. - assert result is inputs - assert result["structured_messages"] is original - - -@pytest.mark.asyncio -async def test_apply_guardrail_sends_textless_parts_rows_unflattened( - guardrail: HeadroomGuardrail, -): - image_only = [ - {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}]}, - {"role": "user", "content": "D" * 5000}, - ] - inputs = GenericGuardrailAPIInputs( - texts=["D" * 5000], - structured_messages=json.loads(json.dumps(image_only)) + [{"role": "user", "content": "and now?"}], - ) - mock_response = _make_compress_response(json.loads(json.dumps(image_only))) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_post: - await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-fable-5"}, - input_type="request", - ) - - wire_messages = mock_post.call_args.kwargs["json"]["messages"] - assert isinstance(wire_messages[0]["content"], list) - assert wire_messages[1]["content"] == "D" * 5000 - - -@pytest.mark.asyncio -async def test_fail_open_returns_original_parts_shapes(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - inputs = GenericGuardrailAPIInputs( - texts=["B" * 5000], - structured_messages=_parts_copy(), - ) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ConnectError("boom"), - ): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={}, - input_type="request", - ) - - messages = result["structured_messages"] - assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] - - -CCR_HASH = "b573993006976af767214fac" - - -def _retrieve_tool_definition() -> dict: - return { - "type": "function", - "function": { - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "description": "retrieve compressed content", - "parameters": {"type": "object", "properties": {"hash": {"type": "string"}}}, - }, - } - - -def _openai_completion_payload(message: dict, finish_reason: str) -> dict: - return { - "id": "chatcmpl-ccr", - "object": "chat.completion", - "created": 1700000000, - "model": "gpt-4o", - "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - } - - -def _openai_tool_call_payload() -> dict: - return _openai_completion_payload( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "call_ccr", - "type": "function", - "function": { - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), - }, - } - ], - }, - "tool_calls", - ) - - -def _openai_text_payload(content: str) -> dict: - return _openai_completion_payload({"role": "assistant", "content": content}, "stop") - - -def _responses_retrieve_tool_definition() -> dict: - return {"type": "function", **_retrieve_tool_definition()["function"]} - - -def _openai_responses_payload(output_item: dict) -> dict: - return { - "id": "resp_ccr", - "object": "response", - "created_at": 1700000000, - "status": "completed", - "model": "gpt-4o", - "output": [output_item], - "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - "parallel_tool_calls": True, - "tool_choice": "auto", - "tools": [], - "error": None, - "incomplete_details": None, - "instructions": None, - "metadata": {}, - "temperature": 1.0, - "top_p": 1.0, - "text": {"format": {"type": "text"}}, - "truncation": "disabled", - } - - -def _openai_responses_retrieve_call_payload() -> dict: - return _openai_responses_payload( - { - "type": "function_call", - "id": "fc_ccr", - "call_id": "call_ccr", - "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), - "status": "completed", - } - ) - - -def _openai_responses_text_payload(text: str) -> dict: - return _openai_responses_payload( - { - "type": "message", - "id": "msg_ccr", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": text, "annotations": []}], - } - ) - - -@pytest.mark.parametrize( - "call_type, stream, tools, expect_conversion", - [ - (CallTypes.acompletion, True, [_retrieve_tool_definition()], True), - (CallTypes.completion, True, [_retrieve_tool_definition()], True), - (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), - (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), - (CallTypes.acompletion, True, None, False), - (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), - (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), - (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), - (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), - ], -) -@pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( - guardrail: HeadroomGuardrail, - call_type: CallTypes, - stream: bool, - tools: Optional[list], - expect_conversion: bool, -): - kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools} - - result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type) - - if not expect_conversion: - assert result is kwargs - assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs - assert kwargs["stream"] is stream - return - - assert result is not None - assert result["stream"] is False - assert result[HEADROOM_CONVERTED_STREAM_KEY] is True - assert kwargs["stream"] is True - - -@pytest.mark.asyncio -async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): - kwargs = { - "model": "gpt-4o", - "stream": True, - "background": True, - "tools": [_responses_retrieve_tool_definition()], - } - - result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses) - - assert result is kwargs - assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs - assert kwargs["stream"] is True - - -@pytest.mark.asyncio -async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( - guardrail: HeadroomGuardrail, -): - """Regression for the stream-conversion override swallowing the parent hook: - when the guardrail is attached at the deployment level and proxy pre_call never - ran, the deployment hook is the only place compression executes, so the - override must delegate to CustomGuardrail.async_pre_call_deployment_hook.""" - kwargs = { - "model": "gpt-4o", - "messages": [dict(m) for m in ORIGINAL_MESSAGES], - "stream": False, - "guardrails": ["headroom"], - "metadata": {}, - } - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES), - ): - result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) - - assert result is not None - assert result["messages"] == EXPECTED_MESSAGES - - -@pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_compression( - guardrail: HeadroomGuardrail, -): - kwargs = { - "model": "gpt-4o", - "messages": [dict(m) for m in ORIGINAL_MESSAGES], - "stream": True, - "guardrails": ["headroom"], - "metadata": {}, - } - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), - ): - result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) - - assert result is not None - assert has_headroom_retrieve_tool(result["tools"]) - assert result["stream"] is False - assert result[HEADROOM_CONVERTED_STREAM_KEY] is True - - -@pytest.mark.asyncio -async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, - monkeypatch: pytest.MonkeyPatch, -): - """Regression test for streaming /chat/completions: the retrieve tool call the - model emits must be resolved by the agentic loop instead of being streamed back - to a client that never declared the tool.""" - original_content = "the full uncompressed document" - final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) - - monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - side_effect=[ - httpx.Response(200, json=_openai_tool_call_payload()), - httpx.Response(200, json=_openai_text_payload(final_answer)), - ] - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response(original_content), - ) as mock_get: - response = await litellm.acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], - tools=[_retrieve_tool_definition()], - stream=True, - litellm_call_id="ccr-call-id", - ) - chunks = [chunk async for chunk in response] - - streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) - assert streamed_text == final_answer - assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) - - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) - assert not followup_body.get("stream") - assert original_content in json.dumps(followup_body["messages"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) - - -@pytest.mark.asyncio -async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, - monkeypatch: pytest.MonkeyPatch, -): - """Regression test for LIT-6481: streaming /v1/responses must resolve the - retrieve tool call server-side exactly like streaming /chat/completions does, - instead of streaming a headroom_retrieve function_call to the client.""" - original_content = "the full uncompressed document" - final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) - - monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( - side_effect=[ - httpx.Response(200, json=_openai_responses_retrieve_call_payload()), - httpx.Response(200, json=_openai_responses_text_payload(final_answer)), - ] - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response(original_content), - ) as mock_get: - response = await litellm.aresponses( - model="openai/gpt-4o", - input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], - tools=[_responses_retrieve_tool_definition()], - stream=True, - litellm_call_id="ccr-call-id", - ) - events = [event async for event in response] - - streamed_text = "".join( - getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) - assert streamed_text == final_answer - assert not any("function_call" in str(getattr(event, "type", "")) for event in events) - assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) - - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) - assert not followup_body.get("stream") - assert original_content in json.dumps(followup_body["input"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) - - -def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, - monkeypatch: pytest.MonkeyPatch, -): - """The synchronous responses() path converts the stream the same way, so it - must hand back a stream iterator with the resolved answer rather than the - completed response object.""" - original_content = "the full uncompressed document" - final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) - - monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( - side_effect=[ - httpx.Response(200, json=_openai_responses_retrieve_call_payload()), - httpx.Response(200, json=_openai_responses_text_payload(final_answer)), - ] - ) - - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response(original_content), - ) as mock_get: - response = litellm.responses( - model="openai/gpt-4o", - input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], - tools=[_responses_retrieve_tool_definition()], - stream=True, - litellm_call_id="ccr-call-id", - ) - events = list(response) - - streamed_text = "".join( - getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) - assert streamed_text == final_answer - assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert len(upstream.calls) == 2 - assert not json.loads(upstream.calls[1].request.content).get("stream") - - -# --------------------------------------------------------------------------- -# LIT-5018: the turn the model is being asked to act on is never compressed. -# -# A Claude Code request ends with the live instruction, preceded by the tool -# result answering the assistant's last tool call. Replacing either with a -# marker makes the model answer a retrieval result instead of the request. -# --------------------------------------------------------------------------- - -AGENTIC_MESSAGES = [ - {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, - {"role": "user", "content": "H" * 5000}, - { - "role": "assistant", - "content": "Older answer. " + "O" * 5000, - "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "old_1", "content": "older tool output " + "T" * 5000}, - { - "role": "assistant", - "content": "Reading the file now.", - "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "tu_1", "content": "FILE BODY " + "F" * 5000}, - { - "role": "user", - "content": [ - {"type": "text", "text": " " + "E" * 5000}, - {"type": "text", "text": "can we run /team to fix this"}, - ], - }, -] - - -async def _wire_and_result(guardrail: HeadroomGuardrail, messages: list, returned: list | None = None): - inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))) - sent: dict = {} - - def _echo(**kwargs): - sent["messages"] = kwargs["json"]["messages"] - return _make_compress_response( - returned if returned is not None else json.loads(json.dumps(kwargs["json"]["messages"])) - ) - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-sonnet-4-5-20250929"}, - input_type="request", - ) - return sent["messages"], result - - -@pytest.mark.asyncio -async def test_live_user_turn_is_never_sent_for_compression(guardrail: HeadroomGuardrail): - wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) - - live_turn = AGENTIC_MESSAGES[-1] - assert live_turn not in wire - assert not any("can we run /team to fix this" in json.dumps(row) for row in wire) - # It reaches the model byte-identical, both text parts intact, so no - # marker and no retrieval round-trip stands in for the instruction. - assert result["structured_messages"][-1] == live_turn - - -@pytest.mark.asyncio -async def test_system_prompt_is_never_sent_for_compression(guardrail: HeadroomGuardrail): - wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) - - assert not any(row.get("role") == "system" for row in wire) - # The Anthropic write-back drops compressed system rows, so sending it - # only inflates the savings the service reports back. - assert result["structured_messages"][0] == AGENTIC_MESSAGES[0] - - -@pytest.mark.asyncio -async def test_trailing_tool_exchange_is_never_sent_for_compression(guardrail: HeadroomGuardrail): - """The tool result answering the last assistant's tool call is protected - with it: a marker there stands in for the result of the call the model just - made, forcing an immediate retrieval of data it already asked for.""" - wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) - - assert not any(row.get("tool_call_id") == "tu_1" for row in wire) - assert result["structured_messages"][5] == AGENTIC_MESSAGES[5] - - -@pytest.mark.asyncio -async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): - """Negative control: protection must not turn compression into a no-op.""" - compressed_history = [ - {"role": "user", "content": "hist. hash=b573993006976af767214fac"}, - {"role": "tool", "tool_call_id": "old_1", "content": "older tool. hash=c73993006976af767214fac2"}, - ] - wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES, returned=compressed_history) - - # Older user and tool rows go to the service, in order. Every assistant row - # stays out, but the tool results those turns asked for remain compressible. - assert [row["role"] for row in wire] == ["user", "tool"] - assert wire[0]["content"] == "H" * 5000 - assert wire[1]["tool_call_id"] == "old_1" - - messages = result["structured_messages"] - assert len(messages) == len(AGENTIC_MESSAGES) - assert messages[1] == compressed_history[0] - assert messages[2] == AGENTIC_MESSAGES[2] - assert messages[3] == compressed_history[1] - - -CACHED_PREFIX_MESSAGES = [ - {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, - {"role": "user", "content": "old question " + "Q" * 5000}, - { - "role": "assistant", - "content": "Reading the file now.", - "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000}, - { - "role": "user", - "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}], - }, - { - "role": "assistant", - "content": "Listing now.", - "tool_calls": [{"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000}, - {"role": "assistant", "content": "Finished listing."}, - {"role": "user", "content": "live instruction"}, -] - - -@pytest.mark.asyncio -async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail): - wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES) - - assert [row.get("tool_call_id") for row in wire] == ["new_1"] - assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5] - - -CACHE_MARKED_HISTORY_MESSAGES = [ - {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, - {"role": "user", "content": "old question " + "Q" * 5000}, - { - "role": "assistant", - "content": "Reading the file now.", - "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], - }, - { - "role": "tool", - "tool_call_id": "old_1", - "content": [{"type": "text", "text": "large cached file body " + "F" * 5000}], - "cache_control": {"type": "ephemeral"}, - }, - {"role": "assistant", "content": "Summarized the file for you."}, - {"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000}, - {"role": "user", "content": "live instruction"}, -] - - -@pytest.mark.asyncio -async def test_mid_history_cache_control_row_is_never_sent_for_compression(guardrail: HeadroomGuardrail): - wire, result = await _wire_and_result(guardrail, CACHE_MARKED_HISTORY_MESSAGES) - - cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] - assert cached_row not in wire - assert not any(row.get("tool_call_id") == "old_1" for row in wire) - assert result["structured_messages"][3] == cached_row - - -# --------------------------------------------------------------------------- -# #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP -# gateway) executes headroom_retrieve and echoes the recovered original content -# back as a tool result. Compressing that row re-derives the same content hash -# it was just retrieved from -- the marker returns and the agent loops. The -# retrieved row must be held back from the compression service. -# --------------------------------------------------------------------------- - -RETRIEVE_ECHO_MESSAGES = [ - {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, - {"role": "user", "content": "H" * 5000}, - { - "role": "assistant", - "content": "Expanding the marker.", - "tool_calls": [ - { - "id": "hr_1", - "type": "function", - "function": { - "name": "mcp__headroom__headroom_retrieve", - "arguments": '{"hash": "b573993006976af767214fac"}', - }, - } - ], - }, - {"role": "tool", "tool_call_id": "hr_1", "content": "RETRIEVED BODY " + "R" * 5000}, - {"role": "assistant", "content": "Older answer. " + "O" * 5000}, - {"role": "user", "content": "now summarize the description"}, -] - - -@pytest.mark.asyncio -async def test_retrieved_content_is_never_recompressed(guardrail: HeadroomGuardrail): - """The tool result carrying headroom_retrieve output is held back, so it can - never collapse back to the hash it was just retrieved from.""" - wire, result = await _wire_and_result(guardrail, RETRIEVE_ECHO_MESSAGES) - - assert not any(row.get("tool_call_id") == "hr_1" for row in wire) - assert not any("RETRIEVED BODY" in json.dumps(row) for row in wire) - # Reaches the model byte-identical, so no marker stands in for the expansion. - assert result["structured_messages"][3] == RETRIEVE_ECHO_MESSAGES[3] - # Negative control: unrelated history is still compressed, not a no-op. - assert any(row.get("content") == "H" * 5000 for row in wire) - - -@pytest.mark.asyncio -async def test_retrieved_content_guard_matches_direct_tool_name(guardrail: HeadroomGuardrail): - """Server-side the tool is named headroom_retrieve (no MCP prefix); its - result must be protected the same way.""" - messages = [ - {"role": "system", "content": "sys " + "S" * 5000}, - {"role": "user", "content": "H" * 5000}, - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "hr_direct", - "type": "function", - "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, - } - ], - }, - {"role": "tool", "tool_call_id": "hr_direct", "content": "RETRIEVED BODY " + "R" * 5000}, - {"role": "assistant", "content": "Older. " + "O" * 5000}, - {"role": "user", "content": "summarize"}, - ] - wire, result = await _wire_and_result(guardrail, messages) - - assert not any(row.get("tool_call_id") == "hr_direct" for row in wire) - assert result["structured_messages"][3] == messages[3] - - -@pytest.mark.asyncio -async def test_retrieved_content_protected_when_mcp_tool_name_is_truncated(guardrail: HeadroomGuardrail): - """A long mcp____headroom_retrieve name is truncated past 64 chars in - the OpenAI-translated view the guardrail scans, dropping the suffix. The call - id read from the request's own Anthropic tool_use (never truncated) still - pairs the retrieved row so it is held back.""" - from litellm.llms.anthropic.pass_through.adapters.transformation import ( - truncate_tool_name, - ) - - long_name = "mcp__" + "s" * 45 + "__" + HEADROOM_RETRIEVE_TOOL_NAME - assert len(long_name) > 64 - truncated = truncate_tool_name(long_name) - assert not truncated.endswith(HEADROOM_RETRIEVE_TOOL_NAME) - - # What the guardrail scans: OpenAI-translated messages with the truncated name. - structured = [ - {"role": "system", "content": "sys " + "S" * 5000}, - {"role": "user", "content": "H" * 5000}, - { - "role": "assistant", - "content": "", - "tool_calls": [{"id": "hr_long", "type": "function", "function": {"name": truncated, "arguments": "{}"}}], - }, - {"role": "tool", "tool_call_id": "hr_long", "content": "RETRIEVED BODY " + "R" * 5000}, - {"role": "assistant", "content": "Older. " + "O" * 5000}, - {"role": "user", "content": "summarize"}, - ] - # The request's own messages, untranslated: Anthropic tool_use carries the full name. - raw_messages = [ - {"role": "assistant", "content": [{"type": "tool_use", "id": "hr_long", "name": long_name, "input": {}}]}, - {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "hr_long", "content": "RETRIEVED BODY"}]}, - ] - - inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(structured))) - sent: dict = {} - - def _echo(**kwargs): - sent["messages"] = kwargs["json"]["messages"] - return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"]))) - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "claude-sonnet-4-5-20250929", "messages": raw_messages}, - input_type="request", - ) - - assert not any(row.get("tool_call_id") == "hr_long" for row in sent["messages"]) - assert result["structured_messages"][3] == structured[3] - assert any(row.get("content") == "H" * 5000 for row in sent["messages"]) - - -def test_raw_retrieve_call_ids_covers_both_shapes_and_ignores_others(): - """Retrieve ids are read from OpenAI tool_calls and Anthropic tool_use blocks; - non-retrieve calls, non-tool_use blocks, string content, and non-list inputs - yield nothing.""" - from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import _raw_retrieve_call_ids - - messages = [ - { - "role": "assistant", - "tool_calls": [ - {"id": "oa1", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}, - {"id": "other", "function": {"name": "get_weather"}}, - {"id": "malformed", "function": {"name": 123}}, - {"id": "nofunc"}, - ], - }, - { - "role": "assistant", - "content": [ - {"type": "tool_use", "id": "an1", "name": "mcp__hr__headroom_retrieve", "input": {}}, - {"type": "tool_use", "id": "an2", "name": "jira_get_issue", "input": {}}, - {"type": "text", "text": "noise"}, - ], - }, - {"role": "user", "content": "plain string content, not a list"}, - ] - - assert _raw_retrieve_call_ids(messages) == frozenset({"oa1", "an1"}) - assert _raw_retrieve_call_ids("not a list") == frozenset() - assert _raw_retrieve_call_ids(None) == frozenset() - - -@pytest.mark.asyncio -async def test_nothing_compressible_returns_inputs_untouched(guardrail: HeadroomGuardrail): - """A single-turn request is all protected, so there is nothing to send and - the caller's own inputs object comes back.""" - inputs = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=[ - {"role": "system", "content": "sys"}, - {"role": "user", "content": "A" * 5000}, - ], - ) - - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: - result = await guardrail.apply_guardrail( - inputs=inputs, - request_data={"model": "gpt-4o"}, - input_type="request", - ) - - mock_post.assert_not_called() - assert result is inputs - - -@pytest.mark.asyncio -async def test_fail_open_returns_the_caller_inputs_object(): - """Translation handlers detect a rewrite by object identity, so a request - that was not compressed must come back as the same object or it is - round-tripped through the write-back for nothing.""" - guardrail = _make_guardrail(unreachable_fallback="fail_open") - original = json.loads(json.dumps(AGENTIC_MESSAGES)) - inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=original) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ConnectError("boom"), - ): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") - - assert result is inputs - assert result["structured_messages"] is original - - -# --------------------------------------------------------------------------- -# LIT-5018: the retrieval follow-up keeps the model's own text. -# --------------------------------------------------------------------------- - - -def _anthropic_response_with_text_and_tool_call() -> dict: - return { - "content": [ - {"type": "text", "text": "Let me pull the original back."}, - {"type": "tool_use", "id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "input": {"hash": "h" * 24}}, - ] - } - - -async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) - logging_obj = MagicMock() - logging_obj.litellm_call_id = "call-1" - logging_obj.model_call_details = {} - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response("ORIGINAL CONTENT"), - ): - return await guardrail.async_build_agentic_loop_plan( - tools={ - "tool_calls": [{"id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": {"hash": "h" * 24}}] - }, - model="claude-sonnet-4-5-20250929", - messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=logging_obj, - stream=False, - kwargs={}, - ) - - -@pytest.mark.asyncio -async def test_anthropic_followup_preserves_assistant_text(guardrail: HeadroomGuardrail): - plan = await _plan_for(guardrail, _anthropic_response_with_text_and_tool_call(), [{"role": "user", "content": "q"}]) - - assistant = plan.request_patch.messages[-2] # type: ignore[union-attr] - assert assistant["role"] == "assistant" - # Text first, then the tool_use it accompanied: dropping it loses the - # model's stated reason for the retrieval from its own transcript. - assert assistant["content"][0] == {"type": "text", "text": "Let me pull the original back."} - assert assistant["content"][1]["type"] == "tool_use" - - -@pytest.mark.asyncio -async def test_responses_followup_preserves_assistant_text(guardrail: HeadroomGuardrail): - response = { - "output": [ - {"type": "message", "content": [{"type": "output_text", "text": "Fetching the original."}]}, - {"type": "function_call", "call_id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, - ] - } - - plan = await _plan_for(guardrail, response, [{"role": "user", "content": "q"}]) - - items = plan.request_patch.messages # type: ignore[union-attr] - assert items[1] == {"role": "assistant", "content": "Fetching the original."} - assert items[2]["type"] == "function_call" - - -@pytest.mark.asyncio -async def test_chat_followup_echoes_only_the_retrieve_call(guardrail: HeadroomGuardrail): - """A turn that called another tool alongside headroom_retrieve must not - echo that call: only the retrieve call gets a tool result, and a tool_call - without one is rejected by the provider.""" - other = MagicMock() - other.id = "call_other" - other.type = "function" - other.function = MagicMock() - other.function.name = "Write" - other.function.arguments = "{}" - - response = _make_openai_response_with_tool_call(HEADROOM_RETRIEVE_TOOL_NAME, {"hash": "h" * 24}, "call_1") - response.choices[0].message.content = "Getting the original first." - response.choices[0].message.tool_calls = [response.choices[0].message.tool_calls[0], other] - - plan = await _plan_for(guardrail, response, [{"role": "user", "content": "q"}]) - - messages = plan.request_patch.messages # type: ignore[union-attr] - assistant = messages[1] - assert assistant["content"] == "Getting the original first." - assert [tc["id"] for tc in assistant["tool_calls"]] == ["call_1"] - assert [m["tool_call_id"] for m in messages[2:]] == ["call_1"] - - -# --- LIT-5881: the calls to the compression service must be time-bounded --- - - -def _timeout_of(mock_call) -> httpx.Timeout: - timeout = mock_call.kwargs["timeout"] - assert isinstance(timeout, httpx.Timeout), timeout - return timeout - - -@pytest.mark.asyncio -async def test_compress_call_passes_bounded_timeout(guardrail: HeadroomGuardrail): - """Without an explicit timeout the call inherits the shared client's 600s read leg.""" - inputs = GenericGuardrailAPIInputs(texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES), - ) as mock_post: - await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") - - timeout = _timeout_of(mock_post.call_args) - assert timeout.read == 60.0 - assert timeout.write == 60.0 - assert timeout.pool == 60.0 - assert timeout.connect == 5.0 - - -@pytest.mark.asyncio -async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail): - """The retrieval leg runs on the same request and needs the same bound.""" - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response("original"), - ) as mock_get: - result = await guardrail._call_retrieve("a" * 24) - - assert result == "original" - timeout = _timeout_of(mock_get.call_args) - assert timeout.read == 60.0 - assert timeout.connect == 5.0 - - -@pytest.mark.asyncio -async def test_configured_timeout_overrides_the_default(): - """Headroom accepted litellm_params.timeout and ignored it.""" - guardrail = _make_guardrail(timeout=3.5) - inputs = GenericGuardrailAPIInputs(texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - return_value=_make_compress_response(COMPRESSED_MESSAGES), - ) as mock_post: - await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") - - timeout = _timeout_of(mock_post.call_args) - assert timeout.read == 3.5 - assert timeout.connect == 3.5 - - -@pytest.mark.asyncio -async def test_read_timeout_is_surfaced_as_unreachable_under_fail_closed(): - """A stalled service must reach the fail policy, not escape as a 500.""" - guardrail = _make_guardrail() - inputs = GenericGuardrailAPIInputs(texts=["hello"], structured_messages=ORIGINAL_MESSAGES) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ReadTimeout("timed out"), - ): - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") - - assert exc_info.value.status_code == 502 - assert "unreachable" in str(exc_info.value.detail) - - -@pytest.mark.asyncio -async def test_read_timeout_forwards_uncompressed_under_fail_open(): - guardrail = _make_guardrail(unreachable_fallback="fail_open") - inputs = GenericGuardrailAPIInputs(texts=["hello"], structured_messages=ORIGINAL_MESSAGES) - - with patch.object( - guardrail.async_handler, - "post", - new_callable=AsyncMock, - side_effect=httpx.ReadTimeout("timed out"), - ): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") - - assert result.get("structured_messages") == ORIGINAL_MESSAGES - - -def test_initializer_forwards_configured_timeout(monkeypatch: pytest.MonkeyPatch): - """Wiring it only in __init__ leaves `timeout:` in config.yaml silently ignored.""" - from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail - from litellm.types.guardrails import LitellmParams - - monkeypatch.setattr( - litellm.logging_callback_manager, - "add_litellm_callback", - lambda callback: None, - ) - params = LitellmParams( - guardrail="headroom", - mode="pre_call", - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - timeout=7.0, - ) - callback = initialize_guardrail(params, {"guardrail_name": "headroom"}) # type: ignore[arg-type] - - assert callback.timeout.read == 7.0 - - -def test_in_place_update_keeps_the_timeout_resolved(): - """The base implementation copies every attribute over, nulling an unset timeout.""" - from litellm.types.guardrails import LitellmParams - - guardrail = _make_guardrail(timeout=5.0) - assert guardrail.timeout.read == 5.0 - - guardrail.update_in_memory_litellm_params( - LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE) - ) - assert isinstance(guardrail.timeout, httpx.Timeout) - assert guardrail.timeout.read == 60.0 - - guardrail.update_in_memory_litellm_params( - LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, timeout=7.0) - ) - assert guardrail.timeout.read == 7.0 - - -@pytest.mark.parametrize("configured", [0, 0.0, -1, -30.0, float("inf"), float("-inf"), float("nan")]) -def test_unusable_timeout_falls_back_to_the_default(configured: float): - """0 and inf read as no deadline at all, a negative one as a deadline already past.""" - guardrail = _make_guardrail(timeout=configured) - - assert guardrail.timeout.read == 60.0 - assert guardrail.timeout.connect == 5.0 - - -def _affinity_of(call: "respx.models.Call") -> str | None: - value: Final = call.request.headers.get(AFFINITY_HEADER) - return str(value) if value else None - - -def _captured_affinity(route: respx.Route) -> str | None: - """Read the affinity header off the most recent call on this route. - - route.calls[0] would always be the first request, which is wrong when a test - drives several requests through one route and wants to compare them. - """ - return _affinity_of(route.calls[-1]) - - -@pytest.mark.asyncio -async def test_compress_sends_stable_tenant_scoped_affinity( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, -): - """Related calls in one conversation reach the same pod, and the raw ids never leave.""" - route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={ - "messages": COMPRESSED_MESSAGES, - "tokens_before": 1000, - "tokens_after": 100, - }, - ) - ) - inputs: Final = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - request_data: Final = { - "model": "gpt-4o", - "litellm_session_id": "sess-abc", - "metadata": {"user_api_key_team_id": "team-1"}, - } - - for _ in range(2): - await guardrail.apply_guardrail( - inputs=inputs, - request_data=dict(request_data), - input_type="request", - ) - - sent: Final = [_affinity_of(call) for call in route.calls] - assert len(sent) == 2, f"expected two compress calls, got {len(sent)}" - first: Final = sent[0] - assert first is not None, "an authenticated session must produce an affinity value" - assert sent[0] == sent[1], f"the same conversation must map to one affinity, got {sent}" - assert len(first) == 32 - assert "sess-abc" not in first - assert "team-1" not in first - - -@pytest.mark.asyncio -async def test_affinity_differs_per_tenant_for_the_same_session_id( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, -): - """A session id is reusable across tenants, so it cannot be the affinity on its own.""" - route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, - ) - ) - inputs: Final = GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ) - - seen: Final[list[str | None]] = [] - for team_id in ("team-A", "team-B"): - await guardrail.apply_guardrail( - inputs=inputs, - request_data={ - "model": "gpt-4o", - "litellm_session_id": "shared-session", - "metadata": {"user_api_key_team_id": team_id}, - }, - input_type="request", - ) - seen.append(_affinity_of(route.calls[-1])) - - assert seen[0] is not None and seen[1] is not None - assert seen[0] != seen[1], f"different tenants must differ, got {seen}" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("request_extra", [{}, {"user_api_key_team_alias": "chosen-by-caller"}]) -async def test_affinity_omitted_without_an_authenticated_tenant( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, - request_extra: dict, -): - """Without a trusted identity the header is dropped rather than risking a cross-tenant route.""" - route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, - ) - ) - - await guardrail.apply_guardrail( - inputs=GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ), - request_data={"model": "gpt-4o", "litellm_session_id": "sess-abc", **request_extra}, - input_type="request", - ) - - assert _captured_affinity(route) is None - - -@pytest.mark.asyncio -async def test_affinity_omitted_for_a_proxy_generated_session_id( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, -): - """A generated id is unique per request, so pinning on it would defeat the cache.""" - route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, - ) - ) - - await guardrail.apply_guardrail( - inputs=GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ), - request_data={ - "model": "gpt-4o", - "litellm_session_id": "generated-per-request", - "metadata": { - "litellm_session_id_generated": True, - "user_api_key_team_id": "team-1", - }, - }, - input_type="request", - ) - - assert _captured_affinity(route) is None - - -@pytest.mark.asyncio -@pytest.mark.parametrize("container", ["metadata", "litellm_metadata"]) -async def test_affinity_uses_the_nested_identity_a_proxied_request_carries( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, - container: str, -): - """A proxied request keeps its identity nested, not at the top level. - - Reading only the top level would silently disable affinity for every real - proxy call, so each container the proxy actually populates is covered. - """ - route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, - ) - ) - - await guardrail.apply_guardrail( - inputs=GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ), - request_data={ - "model": "gpt-4o", - "litellm_session_id": "sess-abc", - container: {"user_api_key_team_id": "team-1"}, - }, - input_type="request", - ) - - assert _captured_affinity(route) is not None, f"{container} must supply the tenant identity" - - -@pytest.mark.asyncio -async def test_retrieve_replays_the_same_affinity_the_compress_call_sent( - guardrail: HeadroomGuardrail, - respx_mock: respx.MockRouter, -): - """The retrieve must carry the identical value, or it lands on a different pod. - - Asserting the header merely exists would not catch a replay that recomputes a - different digest, which is the failure this is here to prevent. - """ - compress: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( - return_value=httpx.Response( - 200, - json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, - ) - ) - request_data: Final = { - "litellm_call_id": "call-1", - "litellm_session_id": "sess-abc", - "metadata": {"user_api_key_team_id": "team-1"}, - } - - await guardrail.apply_guardrail( - inputs=GenericGuardrailAPIInputs( - texts=["A" * 5000], - structured_messages=ORIGINAL_MESSAGES, - ), - request_data=dict(request_data), - input_type="request", - ) - - compress_affinity: Final = _captured_affinity(compress) - assert compress_affinity is not None, "authenticated session must produce an affinity value" - - retrieve: Final = respx_mock.get(f"{FAKE_API_BASE}/v1/retrieve/{CCR_HASH}").mock( - return_value=httpx.Response(200, json={"original_content": "restored"}) - ) - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({CCR_HASH}), time.monotonic() + 999) - - await guardrail.async_build_agentic_loop_plan( - tools={"tool_calls": [{"arguments": {"hash": CCR_HASH}}]}, - model="gpt-4o", - messages=list(ORIGINAL_MESSAGES), - response={}, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params={}, - logging_obj=None, - stream=False, - kwargs=dict(request_data), - ) - - assert _captured_affinity(retrieve) == compress_affinity +""" +Unit tests for the Headroom guardrail. + +Tests cover: +- apply_guardrail compresses messages via /v1/compress and returns them as structured_messages +- x-headroom-bypass: true header causes guardrail to skip compression +- missing or empty messages are passed through unchanged +- response-type input is passed through unchanged +- /v1/compress HTTP error raises HTTPException (fail_closed, the default) +- /v1/compress returning malformed JSON raises HTTPException +- /v1/compress non-2xx surfaces as httpx.HTTPStatusError (raise_for_status), + not a status_code check on the returned response -- both are handled +- unreachable_fallback="fail_open" forwards the request uncompressed instead of raising +- tokens_saved is derived from tokens_before/tokens_after when the compression + service omits it, passed through verbatim when present, and skipped (without + breaking compression) when the token counts are not numeric +- CCR: headroom_retrieve tool injected when compressed messages contain hashes +- CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls +- CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages +- CCR: streaming /chat/completions is converted to a non-streaming call so the agentic + loop resolves the retrieve tool call, then fake-streamed back to the client +- session affinity: an explicit session id plus an authenticated tenant identity + produce a stable, tenant-scoped x-headroom-affinity header on /v1/compress and + /v1/retrieve, and the header is omitted whenever either half is missing +""" + +import json +import time +from typing import Final, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +import respx +from fastapi import HTTPException + +import litellm + +from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( + AFFINITY_HEADER, + HeadroomGuardrail, + has_headroom_retrieve_tool, + HEADROOM_RETRIEVE_TOOL_NAME, +) +from litellm.proxy.spend_tracking.compression_savings import ( + extract_compression_saved_tokens, +) +from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY +from litellm.types.utils import ( + CallTypes, + GenericGuardrailAPIInputs, +) + +FAKE_API_BASE = "https://headroom.example.com" +FAKE_API_KEY = "test-key" + +# The system prompt, the last user turn and the last assistant turn are never +# sent to the compression service, so a fixture needs history for anything to +# be eligible: only index 1 is. +ORIGINAL_MESSAGES = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "A" * 5000}, + {"role": "assistant", "content": "Understood."}, + {"role": "user", "content": "and what about B?"}, +] +COMPRESSIBLE_MESSAGES = [ORIGINAL_MESSAGES[1]] +COMPRESSED_MESSAGES = [{"role": "user", "content": "A" * 500}] +COMPRESSED_MESSAGES_WITH_HASH = [ + { + "role": "user", + "content": "Summary. Retrieve more: hash=b573993006976af767214fac", + }, +] +EXPECTED_MESSAGES = [ + ORIGINAL_MESSAGES[0], + COMPRESSED_MESSAGES[0], + ORIGINAL_MESSAGES[2], + ORIGINAL_MESSAGES[3], +] +EXPECTED_MESSAGES_WITH_HASH = [ + ORIGINAL_MESSAGES[0], + COMPRESSED_MESSAGES_WITH_HASH[0], + ORIGINAL_MESSAGES[2], + ORIGINAL_MESSAGES[3], +] + + +def _make_guardrail(**kwargs) -> HeadroomGuardrail: + defaults = dict( + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + guardrail_name="headroom", + default_on=True, + ) + defaults.update(kwargs) + return HeadroomGuardrail(**defaults) + + +def _make_compress_response( + messages: list, + status: int = 200, + ccr_hashes: list[str] | None = None, +) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "messages": messages, + "tokens_before": 1000, + "tokens_after": 100, + "compression_ratio": 0.1, + "transforms_applied": ["router:smart_crusher:0.35"], + **({} if ccr_hashes is None else {"ccr_hashes": ccr_hashes}), + } + mock.text = "" + return mock + + +def _make_retrieve_response(original_content: str, status: int = 200) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = {"original_content": original_content} + mock.text = original_content + return mock + + +def _make_openai_response_with_tool_call(tool_name: str, arguments: dict, tool_id: str = "call_abc123") -> MagicMock: + fn = MagicMock() + fn.name = tool_name + fn.arguments = json.dumps(arguments) + + tc = MagicMock() + tc.id = tool_id + tc.type = "function" + tc.function = fn + + message = MagicMock() + message.content = None + message.tool_calls = [tc] + + choice = MagicMock() + choice.message = message + + response = MagicMock() + response.choices = [choice] + return response + + +@pytest.fixture +def guardrail() -> HeadroomGuardrail: + return _make_guardrail() + + +def _recorded_guardrail_entries(request_data: dict) -> list: + for container_key in ("metadata", "litellm_metadata"): + container = request_data.get(container_key) + if isinstance(container, dict): + entries = container.get("standard_logging_guardrail_information") + if isinstance(entries, list): + return entries + return [] + + +def _applied_guardrails(request_data: dict) -> list: + for container_key in ("metadata", "litellm_metadata"): + container = request_data.get(container_key) + if isinstance(container, dict) and isinstance(container.get("applied_guardrails"), list): + return container["applied_guardrails"] + return [] + + +@pytest.mark.asyncio +async def test_apply_guardrail_compresses_and_returns_structured_messages( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + request_data = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result.get("structured_messages") == EXPECTED_MESSAGES + + entries = _recorded_guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "headroom" + assert entries[0]["guardrail_status"] == "success" + assert entries[0]["guardrail_provider"] == "headroom" + assert "headroom" in _applied_guardrails(request_data) + + +@pytest.mark.asyncio +async def test_apply_guardrail_leaves_background_requests_uncompressed( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"model": "gpt-4o", "background": True} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result is inputs + post.assert_not_awaited() + assert _recorded_guardrail_entries(request_data) == [] + + +def _recorded_guardrail_response(request_data: dict) -> dict: + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + return entries[0]["guardrail_response"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_derives_tokens_saved_when_service_omits_it( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + # _make_compress_response omits tokens_saved, matching the live service. + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + stats = _recorded_guardrail_response(request_data) + assert stats["tokens_saved"] == 900 + + # Spend tracking reads the entry under the spend-log metadata key. + entry = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert extract_compression_saved_tokens({"guardrail_information": [entry]}) == 900 + + +@pytest.mark.asyncio +async def test_apply_guardrail_passes_through_service_sent_tokens_saved( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + # Deliberately different from tokens_before - tokens_after (900): the + # service-sent value must win over the derived one. + mock_response.json.return_value["tokens_saved"] = 123 + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert _recorded_guardrail_response(request_data)["tokens_saved"] == 123 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tokens_before, tokens_after", + [ + ("1000", "100"), + (True, False), + (None, None), + ], +) +async def test_apply_guardrail_skips_derivation_for_non_numeric_token_counts( + guardrail: HeadroomGuardrail, + tokens_before, + tokens_after, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + mock_response.json.return_value["tokens_before"] = tokens_before + mock_response.json.return_value["tokens_after"] = tokens_after + request_data: dict = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert "tokens_saved" not in _recorded_guardrail_response(request_data) + # Compression itself is unaffected by the skipped derivation. + assert result.get("structured_messages") == EXPECTED_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + tools = result.get("tools") + assert tools is not None + assert has_headroom_retrieve_tool(tools) + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_tool_injected_when_no_hashes( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + tools = result.get("tools") + assert not has_headroom_retrieve_tool(tools or []) + + +@pytest.mark.asyncio +async def test_apply_guardrail_preserves_existing_tools_when_injecting( + guardrail: HeadroomGuardrail, +): + existing_tool = {"type": "function", "function": {"name": "my_tool", "parameters": {}}} + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + tools=[existing_tool], + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + tools = result.get("tools") + assert tools is not None + assert isinstance(tools, list) + assert any(isinstance(t, dict) and t.get("function", {}).get("name") == "my_tool" for t in tools) + assert has_headroom_retrieve_tool(tools) + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_returns_true_for_retrieve_call( + guardrail: HeadroomGuardrail, +): + retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] + response = _make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "b573993006976af767214fac"}, + ) + + should_run, ctx = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=retrieve_tool_def, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is True + assert len(ctx["tool_calls"]) == 1 + assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_returns_false_without_retrieve_tool( + guardrail: HeadroomGuardrail, +): + other_tools = [{"type": "function", "function": {"name": "other_tool"}}] + response = _make_openai_response_with_tool_call( + tool_name="other_tool", + arguments={}, + ) + + should_run, _ = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=other_tools, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is False + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_returns_false_when_no_retrieve_calls( + guardrail: HeadroomGuardrail, +): + retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] + response = _make_openai_response_with_tool_call( + tool_name="some_other_function", + arguments={}, + ) + + should_run, _ = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=retrieve_tool_def, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is False + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( + guardrail: HeadroomGuardrail, +): + original_content = "This is the full compressed content." + mock_retrieve = _make_retrieve_response(original_content) + + # Registered hashes are lowercase; a model may echo the marker's hex in uppercase. + tool_calls = [ + { + "id": "call_abc123", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "B573993006976AF767214FAC"}, + } + ] + response = _make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "B573993006976AF767214FAC"}, + tool_id="call_abc123", + ) + messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] + guardrail._issued_hashes_by_call_id["call-1"] = ( + frozenset({"b573993006976af767214fac"}), + time.monotonic() + 999, + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=mock_retrieve, + ) as mock_get: + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="gpt-4o", + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-1"}, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + + follow_up = plan.request_patch.messages + assert follow_up is not None + + tool_result_message = next((m for m in follow_up if m.get("role") == "tool"), None) + assert tool_result_message is not None + assert tool_result_message["content"] == original_content + assert tool_result_message["tool_call_id"] == "call_abc123" + + mock_get.assert_called_once() + call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] + assert "b573993006976af767214fac" in call_url + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_handles_retrieve_404( + guardrail: HeadroomGuardrail, +): + mock_retrieve = MagicMock() + mock_retrieve.status_code = 404 + + tool_calls = [ + { + "id": "call_xyz", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "deadbeef000000000000dead"}, + } + ] + response = _make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "deadbeef000000000000dead"}, + tool_id="call_xyz", + ) + + messages = [ + { + "role": "user", + "content": "Retrieve more: hash=deadbeef000000000000dead", + } + ] + guardrail._issued_hashes_by_call_id["call-1"] = ( + frozenset({"deadbeef000000000000dead"}), + time.monotonic() + 999, + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=mock_retrieve, + ): + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="gpt-4o", + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-1"}, + ) + + follow_up = plan.request_patch.messages # type: ignore[union-attr] + tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) + assert tool_result is not None + assert "not found" in tool_result["content"] or "expired" in tool_result["content"] + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( + guardrail: HeadroomGuardrail, +): + """A hash-shaped string planted in message text must not be honored when + this guardrail has no record of ever issuing it, even if it's echoed back + in the current request's own messages (e.g. via prompt injection).""" + tool_calls = [ + { + "id": "call_xyz", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "deadbeef000000000000dead"}, + } + ] + response = _make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "deadbeef000000000000dead"}, + tool_id="call_xyz", + ) + assert not guardrail._issued_hashes_by_call_id + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + ) as mock_get: + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="gpt-4o", + messages=[{"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-unknown"}, + ) + + mock_get.assert_not_called() + + follow_up = plan.request_patch.messages # type: ignore[union-attr] + tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) + assert tool_result is not None + assert "was not produced by the current request" in tool_result["content"] + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( + guardrail: HeadroomGuardrail, +): + """A hash issued for one request must not be retrievable by a different + request just because the second request echoes that hash-shaped string + back in its own messages -- retrieval must be scoped per litellm_call_id, + not derived by re-scanning attacker-controlled message text.""" + guardrail._issued_hashes_by_call_id["call-A"] = ( + frozenset({"b573993006976af767214fac"}), + time.monotonic() + 999, + ) + + tool_calls = [ + { + "id": "call_xyz", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "b573993006976af767214fac"}, + } + ] + response = _make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "b573993006976af767214fac"}, + tool_id="call_xyz", + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + ) as mock_get: + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="gpt-4o", + messages=[{"role": "user", "content": "Please fetch hash=b573993006976af767214fac for me"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-B"}, + ) + + mock_get.assert_not_called() + + follow_up = plan.request_patch.messages # type: ignore[union-attr] + tool_result = next((m for m in follow_up if m.get("role") == "tool"), None) + assert tool_result is not None + assert "was not produced by the current request" in tool_result["content"] + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_items( + guardrail: HeadroomGuardrail, +): + """For the Responses API, follow-up input must echo a function_call paired + with a function_call_output keyed by the same call_id -- chat-style + assistant/tool messages are not valid Responses API input items.""" + original_content = "This is the full compressed content." + mock_retrieve = _make_retrieve_response(original_content) + + response = MagicMock() + response.choices = None + response.content = None + response.output = [ + { + "type": "function_call", + "id": "fc_abc123", + "call_id": "call_abc123", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": "b573993006976af767214fac"}), + } + ] + + tool_calls = [ + { + "id": "call_abc123", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "b573993006976af767214fac"}, + } + ] + messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] + guardrail._issued_hashes_by_call_id["call-1"] = ( + frozenset({"b573993006976af767214fac"}), + time.monotonic() + 999, + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=mock_retrieve, + ): + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="gpt-4o", + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-1"}, + ) + + follow_up = plan.request_patch.messages # type: ignore[union-attr] + assert all("role" not in item for item in follow_up if item not in messages) + + function_call_item = next((i for i in follow_up if i.get("type") == "function_call"), None) + assert function_call_item is not None + assert function_call_item["call_id"] == "call_abc123" + assert function_call_item["name"] == HEADROOM_RETRIEVE_TOOL_NAME + + output_item = next((i for i in follow_up if i.get("type") == "function_call_output"), None) + assert output_item is not None + assert output_item["call_id"] == "call_abc123" + assert output_item["output"] == original_content + + +@pytest.mark.asyncio +async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messages( + guardrail: HeadroomGuardrail, +): + """For the Anthropic Messages API, follow-up must echo a tool_use content + block in an assistant message paired with a tool_result content block in a + user message keyed by the same tool_use_id -- chat-style tool-role + messages are not valid Anthropic input. + + AnthropicMessagesResponse is a TypedDict, so real responses are plain + dicts at runtime; a MagicMock response here would pass even if branch + selection used bare getattr() and silently fell through to the + chat-completions replay shape for every real Anthropic response. + """ + original_content = "This is the full compressed content." + mock_retrieve = _make_retrieve_response(original_content) + + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_abc123", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input": {"hash": "b573993006976af767214fac"}, + } + ] + } + + tool_calls = [ + { + "id": "toolu_abc123", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "b573993006976af767214fac"}, + } + ] + messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] + guardrail._issued_hashes_by_call_id["call-1"] = ( + frozenset({"b573993006976af767214fac"}), + time.monotonic() + 999, + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=mock_retrieve, + ): + plan = await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls}, + model="claude-sonnet-4-5", + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-1"}, + ) + + follow_up = plan.request_patch.messages # type: ignore[union-attr] + assert all(m.get("role") != "tool" for m in follow_up) + + assistant_message = next((m for m in follow_up if m.get("role") == "assistant"), None) + assert assistant_message is not None + tool_use_block = next((b for b in assistant_message["content"] if b.get("type") == "tool_use"), None) + assert tool_use_block is not None + assert tool_use_block["id"] == "toolu_abc123" + + user_message = follow_up[-1] + assert user_message["role"] == "user" + tool_result_block = next((b for b in user_message["content"] if b.get("type") == "tool_result"), None) + assert tool_result_block is not None + assert tool_result_block["tool_use_id"] == "toolu_abc123" + assert tool_result_block["content"] == original_content + + +HASH_SHAPED_HISTORY = [ + {"role": "system", "content": "You are Claude Code."}, + {"role": "user", "content": [{"type": "text", "text": "Run git log."}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Done."}], + "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "tu_1", "content": "hash=3f2a9c1d7e5b4a6f8c0d1e2f9a8b7c6d5e4f3a2b"}, + {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me."}, +] + + +async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: + request_data = {"model": "claude-sonnet-5"} + + def _echo(**kwargs): + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + return await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), + request_data=request_data, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail: HeadroomGuardrail): + """Regression for LIT-7086: a git SHA in a tool result and a hash= string the + caller typed both look like markers, but the service stored nothing, so the + tool must not appear and no hash may be registered as issued. Covers a + service that omits ccr_hashes, returns it empty, or returns a non-list.""" + for ccr_hashes in (None, [], "b573993006976af767214fac"): + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardrail: HeadroomGuardrail): + """Only the hashes the service reports in ccr_hashes are honored, in the + service's own 12 to 24 hex grammar; anything else is dropped because each + entry is interpolated into the /v1/retrieve URL.""" + result = await _apply( + guardrail, + HASH_SHAPED_HISTORY, + ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + ) + + assert has_headroom_retrieve_tool(result.get("tools") or []) + ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() + assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + + +@pytest.mark.asyncio +async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatch: pytest.MonkeyPatch): + """`ccr_retrieval: false` in config.yaml compresses without any retrieval + round trip, so it has to reach the instance through the initializer.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None) + params = LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, ccr_retrieval=False) + guardrail = initialize_guardrail(params, {"guardrail_name": "headroom", "litellm_params": params}) + + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["b573993006976af767214fac"]) + + assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_anthropic_assistant_history_never_reaches_compression_service(guardrail: HeadroomGuardrail): + """The public Anthropic handler translates assistant content blocks to a + string before Headroom sees them, so model-authored rows must be excluded + from the compression payload rather than protected by their content shape.""" + from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler + + table = "| Guardrail | Model |\n|---|---|\n" + "\n".join(f"| gr-{i} | model-{i} |" for i in range(40)) + messages = [ + {"role": "user", "content": [{"type": "text", "text": "List the guardrails."}]}, + {"role": "assistant", "content": [{"type": "text", "text": table}]}, + {"role": "user", "content": [{"type": "text", "text": "Earlier follow-up. " + "B" * 5000}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]}, + {"role": "user", "content": "Re-print the table."}, + ] + sent: dict = {} + + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(sent["messages"]))) + + data = {"model": "claude-sonnet-5", "messages": json.loads(json.dumps(messages))} + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [row["role"] for row in sent["messages"]] == ["user", "user"] + assert sent["messages"][1]["content"] == "Earlier follow-up. " + "B" * 5000 + assert table not in json.dumps(sent["messages"]) + assert result["messages"][1]["content"] == [{"type": "text", "text": table}] + + +def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape(): + """By the time an Anthropic Messages API response reaches the agentic-loop + gate, the OpenAI-shaped tool this guardrail injects (type: "function") + has already been transformed into Anthropic's native tool shape + (type: "custom", top-level "name", no nested "function" object).""" + anthropic_native_tools = [ + { + "type": "custom", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, + } + ] + assert has_headroom_retrieve_tool(anthropic_native_tools) + assert not has_headroom_retrieve_tool([{"type": "custom", "name": "some_other_tool"}]) + + +@pytest.mark.asyncio +async def test_apply_guardrail_bypass_header_skips_compression( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "true"}}} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + mock_post.assert_not_called() + + assert result.get("structured_messages") == ORIGINAL_MESSAGES + assert _recorded_guardrail_entries(request_data) == [] + + +@pytest.mark.asyncio +async def test_apply_guardrail_response_type_passthrough( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["some response text"], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data: dict = {"model": "gpt-4o"} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + mock_post.assert_not_called() + + assert result is inputs + assert _recorded_guardrail_entries(request_data) == [] + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_structured_messages_passthrough( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + request_data: dict = {"model": "gpt-4o"} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + mock_post.assert_not_called() + + assert result is inputs + assert _recorded_guardrail_entries(request_data) == [] + assert "headroom" not in _applied_guardrails(request_data) + + +@pytest.mark.asyncio +async def test_passthrough_handler_does_not_log_headroom_as_run( + guardrail: HeadroomGuardrail, +): + """Regression for LIT-4650. + + A passthrough request drives headroom through PassThroughEndpointHandler, which + only supplies `texts` (no `structured_messages`). Headroom cannot compress that + shape and no-ops, so it must not appear in the spend log's + standard_logging_guardrail_information as a successful run. + """ + from litellm.llms.pass_through.guardrail_translation.handler import ( + PassThroughEndpointHandler, + ) + + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + await PassThroughEndpointHandler().process_input_messages( + data=data, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + mock_post.assert_not_called() + + assert _recorded_guardrail_entries(data) == [] + assert "headroom" not in _applied_guardrails(data) + + +@pytest.mark.asyncio +async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream( + guardrail: HeadroomGuardrail, +): + """Regression for LIT-6494: on /v1/responses the compressed messages must be + written back into `input`, not only the retrieve tool into `tools`, or the + model keeps reading the full document and never calls headroom_retrieve.""" + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + + data = { + "model": "gpt-5.6", + "instructions": ORIGINAL_MESSAGES[0]["content"], + "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]], + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}], + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), + ): + result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] + assert [item["content"] for item in result["input"]] == [ + COMPRESSED_MESSAGES_WITH_HASH[0]["content"], + ORIGINAL_MESSAGES[2]["content"], + ORIGINAL_MESSAGES[3]["content"], + ] + assert "A" * 5000 not in json.dumps(result["input"]) + assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME] + + +@pytest.mark.asyncio +async def test_apply_guardrail_http_error_raises(): + guardrail = _make_guardrail() + mock_response = _make_compress_response([], status=500) + mock_response.text = "Internal Server Error" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_apply_guardrail_transport_error_raises(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("Connection refused"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "unreachable" in str(exc_info.value.detail) + + +def _make_http_status_error(status: int, body: str) -> httpx.HTTPStatusError: + request = httpx.Request("POST", f"{FAKE_API_BASE}/v1/compress") + response = httpx.Response(status, request=request, text=body) + return httpx.HTTPStatusError( + f"Server error '{status}' for url", + request=request, + response=response, + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_http_status_error_raises(): + """Regression test: litellm's async httpx client calls raise_for_status() + internally, so a non-2xx /v1/compress response surfaces as + httpx.HTTPStatusError, not as a returned MagicMock with status_code set. + A prior version of _call_compress only checked response.status_code and + never caught this exception, so it went unhandled instead of blocking + the request per fail_closed policy.""" + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(500, "headroom internal error"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_apply_guardrail_404_error_includes_troubleshooting_hint(): + """404 responses include a troubleshooting hint for self-hosted Headroom deployments.""" + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(404, "Not Found"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 404 + assert exc_info.value.detail["body"] == "Not Found" + assert "hint" in exc_info.value.detail + assert "HEADROOM_COMPRESS_ALLOW_REMOTE=1" in exc_info.value.detail["hint"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_404_error_omits_troubleshooting_hint(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(500, "headroom internal error"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 500 + assert exc_info.value.detail["body"] == "headroom internal error" + assert "hint" not in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_apply_guardrail_http_status_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(500, "headroom internal error"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_transport_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"model": "gpt-4o"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("Connection refused"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + entries = _recorded_guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "headroom" + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + assert "headroom" in _applied_guardrails(request_data) + + +@pytest.mark.asyncio +async def test_apply_guardrail_http_error_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + mock_response = _make_compress_response([], status=500) + mock_response.text = "Internal Server Error" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_json_response_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.side_effect = ValueError("not JSON") + mock_response.text = "not json" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_messages_key_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"tokens_before": 100, "tokens_after": 10} + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_compressed_messages_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "messages": ["not-a-dict", 42, None], + "tokens_before": 1000, + "tokens_after": 0, + "compression_ratio": 0, + } + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_messages(): + """When compression fails with fail_open, user-supplied messages that + happen to contain hash-shaped strings must NOT cause those hashes to be + registered as valid for CCR retrieval. Otherwise an attacker can plant a + hash= string in their prompt, trigger a compression failure, and have + that hash honored by a later headroom_retrieve tool call.""" + messages_with_fake_hash = [ + {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me"}, + ] + guardrail = _make_guardrail(unreachable_fallback="fail_open") + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=messages_with_fake_hash, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("Connection refused"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == messages_with_fake_hash + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert not guardrail._issued_hashes_by_call_id + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_messages_key_raises(): + guardrail = _make_guardrail() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"tokens_before": 100, "tokens_after": 10} + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_compressed_messages_raises(): + guardrail = _make_guardrail() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "messages": ["not-a-dict", 42, None], + "tokens_before": 1000, + "tokens_after": 0, + "compression_ratio": 0, + } + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "empty message list" in str(exc_info.value.detail) + + +def test_init_raises_without_api_base(): + with pytest.raises(ValueError, match="API base URL"): + HeadroomGuardrail(api_base=None) + + +def test_init_defaults_to_fail_closed(): + guardrail = _make_guardrail() + assert guardrail.unreachable_fallback == "fail_closed" + + +def test_init_rejects_invalid_unreachable_fallback_value(): + guardrail = _make_guardrail(unreachable_fallback="not-a-real-mode") + assert guardrail.unreachable_fallback == "fail_closed" + + +def test_bypass_header_case_insensitive(): + guardrail = _make_guardrail() + + for header_value in ("true", "True", "TRUE"): + data = {"proxy_server_request": {"headers": {"x-headroom-bypass": header_value}}} + assert guardrail._should_bypass(data) is True + + data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "false"}}} + assert guardrail._should_bypass(data) is False + + data = {"proxy_server_request": {"headers": {}}} + assert guardrail._should_bypass(data) is False + + data = {} + assert guardrail._should_bypass(data) is False + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_model_from_config(): + guardrail = _make_guardrail(model="gpt-4o-mini") + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args + sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] + assert sent_payload.get("model") == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_model_from_request_data_when_no_config_model(): + guardrail = _make_guardrail() + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args + sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] + assert sent_payload.get("model") == "gpt-4o" + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_detects_anthropic_content_block_format( + guardrail: HeadroomGuardrail, +): + # Anthropic's native tool format (type: "custom", top-level "name") -- + # by the time a Messages API response reaches this gate, the OpenAI-shaped + # tool this guardrail injects has already been transformed into this shape. + retrieve_tool_def = [ + { + "type": "custom", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, + } + ] + + response = MagicMock() + response.choices = None + response.content = [ + { + "type": "tool_use", + "id": "toolu_abc", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input": {"hash": "b573993006976af767214fac"}, + } + ] + + should_run, ctx = await guardrail.async_should_run_agentic_loop( + response=response, + model="claude-sonnet-4-6", + messages=[], + tools=retrieve_tool_def, + stream=False, + custom_llm_provider="anthropic", + kwargs={}, + ) + + assert should_run is True + assert len(ctx["tool_calls"]) == 1 + assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_detects_anthropic_response_as_plain_dict( + guardrail: HeadroomGuardrail, +): + """AnthropicMessagesResponse is a TypedDict -- real Messages API responses + are plain dicts at runtime, not objects with attribute access. A + MagicMock-only test would pass even if detection used bare getattr() and + silently treated every real response as having no tool calls.""" + retrieve_tool_def = [ + { + "type": "custom", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input_schema": {"type": "object", "properties": {"hash": {"type": "string"}}}, + } + ] + response = { + "content": [ + { + "type": "tool_use", + "id": "toolu_abc", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "input": {"hash": "b573993006976af767214fac"}, + } + ] + } + + should_run, ctx = await guardrail.async_should_run_agentic_loop( + response=response, + model="claude-sonnet-4-6", + messages=[], + tools=retrieve_tool_def, + stream=False, + custom_llm_provider="anthropic", + kwargs={}, + ) + + assert should_run is True + assert len(ctx["tool_calls"]) == 1 + assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" + + +@pytest.mark.asyncio +async def test_async_should_run_agentic_loop_detects_responses_api_output_format( + guardrail: HeadroomGuardrail, +): + retrieve_tool_def = [{"type": "function", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}] + + response = MagicMock() + response.choices = None + response.content = None + response.output = [ + { + "type": "function_call", + "id": "fc_abc123", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": "b573993006976af767214fac"}), + } + ] + + should_run, ctx = await guardrail.async_should_run_agentic_loop( + response=response, + model="gpt-4o", + messages=[], + tools=retrieve_tool_def, + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is True + assert len(ctx["tool_calls"]) == 1 + assert ctx["tool_calls"][0]["arguments"]["hash"] == "b573993006976af767214fac" + + +@pytest.mark.asyncio +async def test_apply_guardrail_litellm_timeout_raises_when_fail_closed(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=litellm.Timeout( + message="Connection timed out after 10 seconds.", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "unreachable" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_litellm_timeout_fail_open_forwards_uncompressed(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=litellm.Timeout( + message="Connection timed out after 10 seconds.", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result["structured_messages"] == ORIGINAL_MESSAGES + + +# --------------------------------------------------------------------------- +# Content-parts flattening (LIT-4795) +# +# Anthropic-format requests translate to messages whose content is a list of +# part dicts. The compression service only rewrites string content, so the +# guardrail flattens ALL-TEXT part lists on the wire and restores the +# original shapes afterwards. Rows with non-text parts are never flattened: +# cache_control breakpoints are positional, and merging text across a +# non-text part would move a later breakpoint to the other side of it. +# --------------------------------------------------------------------------- + +PARTS_MESSAGES = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Earlier turn."}, + {"type": "text", "text": "Second block. " + "B" * 5000}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Mixed row text."}, + {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}, + ], + }, + {"role": "tool", "content": "tool output " + "C" * 500}, + {"role": "user", "content": "what does that file do?"}, +] + +FLATTENED_HISTORY_TEXT = "Earlier turn.\n\nSecond block. " + "B" * 5000 + + +def _parts_copy() -> list: + return json.loads(json.dumps(PARTS_MESSAGES)) + + +def _echo_wire_view() -> list: + """What the service receives (and echoes back when it changes nothing). + + The system row and the trailing user row are never sent. + """ + return [ + {"role": "user", "content": FLATTENED_HISTORY_TEXT}, + json.loads(json.dumps(PARTS_MESSAGES[2])), + {"role": "tool", "content": "tool output " + "C" * 500}, + ] + + +@pytest.mark.asyncio +async def test_apply_guardrail_flattens_all_text_rows_only( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + mock_response = _make_compress_response(_echo_wire_view()) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + wire_messages = mock_post.call_args.kwargs["json"]["messages"] + assert wire_messages[0]["content"] == FLATTENED_HISTORY_TEXT + # Mixed text+image row is never flattened: merging its text would move a + # later cache_control breakpoint across the image part. + assert isinstance(wire_messages[1]["content"], list) + assert wire_messages[2]["content"] == "tool output " + "C" * 500 + + +@pytest.mark.asyncio +async def test_apply_guardrail_restores_rewritten_all_text_row( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + compressed = _echo_wire_view() + compressed[0]["content"] = "compressed history. Retrieve more: hash=b573993006976af767214fac" + mock_response = _make_compress_response(compressed, ccr_hashes=["b573993006976af767214fac"]) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + messages = result["structured_messages"] + history_content = messages[1]["content"] + assert isinstance(history_content, list) + assert len(history_content) == 1 + assert history_content[0]["text"] == "compressed history. Retrieve more: hash=b573993006976af767214fac" + # Mixed row passes through byte-identical. + assert messages[2]["content"] == PARTS_MESSAGES[2]["content"] + # The service-declared hash still drives retrieve-tool injection on a restored row. + assert has_headroom_retrieve_tool(result.get("tools") or []) + + +@pytest.mark.asyncio +async def test_apply_guardrail_keeps_originals_when_service_echoes_unchanged( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + mock_response = _make_compress_response(_echo_wire_view()) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + messages = result["structured_messages"] + assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] + + +@pytest.mark.asyncio +async def test_apply_guardrail_rejects_service_output_when_rows_dropped( + guardrail: HeadroomGuardrail, +): + """A reshaped conversation cannot be applied at all: the rows held back from + compression are matched positionally, so a response with a different row + count goes through the fail policy instead of being adopted.""" + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + dropped = [{"role": "user", "content": "B" * 50}] + mock_response = _make_compress_response(dropped) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "changed the message count" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_forwards_original_when_rows_dropped_and_fail_open(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + original = _parts_copy() + inputs = GenericGuardrailAPIInputs(texts=["B" * 5000], structured_messages=original) + mock_response = _make_compress_response([{"role": "user", "content": "B" * 50}]) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + # Same object back, so translation handlers that detect a rewrite by + # identity leave the request alone instead of round-tripping it. + assert result is inputs + assert result["structured_messages"] is original + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_textless_parts_rows_unflattened( + guardrail: HeadroomGuardrail, +): + image_only = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}]}, + {"role": "user", "content": "D" * 5000}, + ] + inputs = GenericGuardrailAPIInputs( + texts=["D" * 5000], + structured_messages=json.loads(json.dumps(image_only)) + [{"role": "user", "content": "and now?"}], + ) + mock_response = _make_compress_response(json.loads(json.dumps(image_only))) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-fable-5"}, + input_type="request", + ) + + wire_messages = mock_post.call_args.kwargs["json"]["messages"] + assert isinstance(wire_messages[0]["content"], list) + assert wire_messages[1]["content"] == "D" * 5000 + + +@pytest.mark.asyncio +async def test_fail_open_returns_original_parts_shapes(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = GenericGuardrailAPIInputs( + texts=["B" * 5000], + structured_messages=_parts_copy(), + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("boom"), + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + messages = result["structured_messages"] + assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] + + +CCR_HASH = "b573993006976af767214fac" + + +def _retrieve_tool_definition() -> dict: + return { + "type": "function", + "function": { + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "description": "retrieve compressed content", + "parameters": {"type": "object", "properties": {"hash": {"type": "string"}}}, + }, + } + + +def _openai_completion_payload(message: dict, finish_reason: str) -> dict: + return { + "id": "chatcmpl-ccr", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _openai_tool_call_payload() -> dict: + return _openai_completion_payload( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_ccr", + "type": "function", + "function": { + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + }, + } + ], + }, + "tool_calls", + ) + + +def _openai_text_payload(content: str) -> dict: + return _openai_completion_payload({"role": "assistant", "content": content}, "stop") + + +def _responses_retrieve_tool_definition() -> dict: + return {"type": "function", **_retrieve_tool_definition()["function"]} + + +def _openai_responses_payload(output_item: dict) -> dict: + return { + "id": "resp_ccr", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-4o", + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "text": {"format": {"type": "text"}}, + "truncation": "disabled", + } + + +def _openai_responses_retrieve_call_payload() -> dict: + return _openai_responses_payload( + { + "type": "function_call", + "id": "fc_ccr", + "call_id": "call_ccr", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + "status": "completed", + } + ) + + +def _openai_responses_text_payload(text: str) -> dict: + return _openai_responses_payload( + { + "type": "message", + "id": "msg_ccr", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ) + + +@pytest.mark.parametrize( + "call_type, stream, tools, expect_conversion", + [ + (CallTypes.acompletion, True, [_retrieve_tool_definition()], True), + (CallTypes.completion, True, [_retrieve_tool_definition()], True), + (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), + (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), + (CallTypes.acompletion, True, None, False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), + (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), + (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), + (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), + ], +) +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( + guardrail: HeadroomGuardrail, + call_type: CallTypes, + stream: bool, + tools: Optional[list], + expect_conversion: bool, +): + kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools} + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type) + + if not expect_conversion: + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs + assert kwargs["stream"] is stream + return + + assert result is not None + assert result["stream"] is False + assert result[HEADROOM_CONVERTED_STREAM_KEY] is True + assert kwargs["stream"] is True + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): + kwargs = { + "model": "gpt-4o", + "stream": True, + "background": True, + "tools": [_responses_retrieve_tool_definition()], + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses) + + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs + assert kwargs["stream"] is True + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( + guardrail: HeadroomGuardrail, +): + """Regression for the stream-conversion override swallowing the parent hook: + when the guardrail is attached at the deployment level and proxy pre_call never + ran, the deployment hook is the only place compression executes, so the + override must delegate to CustomGuardrail.async_pre_call_deployment_hook.""" + kwargs = { + "model": "gpt-4o", + "messages": [dict(m) for m in ORIGINAL_MESSAGES], + "stream": False, + "guardrails": ["headroom"], + "metadata": {}, + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ): + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) + + assert result is not None + assert result["messages"] == EXPECTED_MESSAGES + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_compression( + guardrail: HeadroomGuardrail, +): + kwargs = { + "model": "gpt-4o", + "messages": [dict(m) for m in ORIGINAL_MESSAGES], + "stream": True, + "guardrails": ["headroom"], + "metadata": {}, + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]), + ): + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) + + assert result is not None + assert has_headroom_retrieve_tool(result["tools"]) + assert result["stream"] is False + assert result[HEADROOM_CONVERTED_STREAM_KEY] is True + + +@pytest.mark.asyncio +async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for streaming /chat/completions: the retrieve tool call the + model emits must be resolved by the agentic loop instead of being streamed back + to a client that never declared the tool.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[ + httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + chunks = [chunk async for chunk in response] + + streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + assert streamed_text == final_answer + assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["messages"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + +@pytest.mark.asyncio +async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for LIT-6481: streaming /v1/responses must resolve the + retrieve tool call server-side exactly like streaming /chat/completions does, + instead of streaming a headroom_retrieve function_call to the client.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = [event async for event in response] + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any("function_call" in str(getattr(event, "type", "")) for event in events) + assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["input"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + +def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """The synchronous responses() path converts the stream the same way, so it + must hand back a stream iterator with the resolved answer rather than the + completed response object.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = litellm.responses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = list(response) + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) + mock_get.assert_called_once() + assert len(upstream.calls) == 2 + assert not json.loads(upstream.calls[1].request.content).get("stream") + + +# --------------------------------------------------------------------------- +# LIT-5018: the turn the model is being asked to act on is never compressed. +# +# A Claude Code request ends with the live instruction, preceded by the tool +# result answering the assistant's last tool call. Replacing either with a +# marker makes the model answer a retrieval result instead of the request. +# --------------------------------------------------------------------------- + +AGENTIC_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "Older answer. " + "O" * 5000, + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "old_1", "content": "older tool output " + "T" * 5000}, + { + "role": "assistant", + "content": "Reading the file now.", + "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "tu_1", "content": "FILE BODY " + "F" * 5000}, + { + "role": "user", + "content": [ + {"type": "text", "text": " " + "E" * 5000}, + {"type": "text", "text": "can we run /team to fix this"}, + ], + }, +] + + +async def _wire_and_result(guardrail: HeadroomGuardrail, messages: list, returned: list | None = None): + inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))) + sent: dict = {} + + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response( + returned if returned is not None else json.loads(json.dumps(kwargs["json"]["messages"])) + ) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-sonnet-4-5-20250929"}, + input_type="request", + ) + return sent["messages"], result + + +@pytest.mark.asyncio +async def test_live_user_turn_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) + + live_turn = AGENTIC_MESSAGES[-1] + assert live_turn not in wire + assert not any("can we run /team to fix this" in json.dumps(row) for row in wire) + # It reaches the model byte-identical, both text parts intact, so no + # marker and no retrieval round-trip stands in for the instruction. + assert result["structured_messages"][-1] == live_turn + + +@pytest.mark.asyncio +async def test_system_prompt_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) + + assert not any(row.get("role") == "system" for row in wire) + # The Anthropic write-back drops compressed system rows, so sending it + # only inflates the savings the service reports back. + assert result["structured_messages"][0] == AGENTIC_MESSAGES[0] + + +@pytest.mark.asyncio +async def test_trailing_tool_exchange_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + """The tool result answering the last assistant's tool call is protected + with it: a marker there stands in for the result of the call the model just + made, forcing an immediate retrieval of data it already asked for.""" + wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES) + + assert not any(row.get("tool_call_id") == "tu_1" for row in wire) + assert result["structured_messages"][5] == AGENTIC_MESSAGES[5] + + +@pytest.mark.asyncio +async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): + """Negative control: protection must not turn compression into a no-op.""" + compressed_history = [ + {"role": "user", "content": "hist. hash=b573993006976af767214fac"}, + {"role": "tool", "tool_call_id": "old_1", "content": "older tool. hash=c73993006976af767214fac2"}, + ] + wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES, returned=compressed_history) + + # Older user and tool rows go to the service, in order. Every assistant row + # stays out, but the tool results those turns asked for remain compressible. + assert [row["role"] for row in wire] == ["user", "tool"] + assert wire[0]["content"] == "H" * 5000 + assert wire[1]["tool_call_id"] == "old_1" + + messages = result["structured_messages"] + assert len(messages) == len(AGENTIC_MESSAGES) + assert messages[1] == compressed_history[0] + assert messages[2] == AGENTIC_MESSAGES[2] + assert messages[3] == compressed_history[1] + + +CACHED_PREFIX_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "old question " + "Q" * 5000}, + { + "role": "assistant", + "content": "Reading the file now.", + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000}, + { + "role": "user", + "content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}], + }, + { + "role": "assistant", + "content": "Listing now.", + "tool_calls": [{"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000}, + {"role": "assistant", "content": "Finished listing."}, + {"role": "user", "content": "live instruction"}, +] + + +@pytest.mark.asyncio +async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES) + + assert [row.get("tool_call_id") for row in wire] == ["new_1"] + assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5] + + +CACHE_MARKED_HISTORY_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "old question " + "Q" * 5000}, + { + "role": "assistant", + "content": "Reading the file now.", + "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "old_1", + "content": [{"type": "text", "text": "large cached file body " + "F" * 5000}], + "cache_control": {"type": "ephemeral"}, + }, + {"role": "assistant", "content": "Summarized the file for you."}, + {"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000}, + {"role": "user", "content": "live instruction"}, +] + + +@pytest.mark.asyncio +async def test_mid_history_cache_control_row_is_never_sent_for_compression(guardrail: HeadroomGuardrail): + wire, result = await _wire_and_result(guardrail, CACHE_MARKED_HISTORY_MESSAGES) + + cached_row = CACHE_MARKED_HISTORY_MESSAGES[3] + assert cached_row not in wire + assert not any(row.get("tool_call_id") == "old_1" for row in wire) + assert result["structured_messages"][3] == cached_row + + +# --------------------------------------------------------------------------- +# #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP +# gateway) executes headroom_retrieve and echoes the recovered original content +# back as a tool result. Compressing that row re-derives the same content hash +# it was just retrieved from -- the marker returns and the agent loops. The +# retrieved row must be held back from the compression service. +# --------------------------------------------------------------------------- + +RETRIEVE_ECHO_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "Expanding the marker.", + "tool_calls": [ + { + "id": "hr_1", + "type": "function", + "function": { + "name": "mcp__headroom__headroom_retrieve", + "arguments": '{"hash": "b573993006976af767214fac"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_1", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older answer. " + "O" * 5000}, + {"role": "user", "content": "now summarize the description"}, +] + + +@pytest.mark.asyncio +async def test_retrieved_content_is_never_recompressed(guardrail: HeadroomGuardrail): + """The tool result carrying headroom_retrieve output is held back, so it can + never collapse back to the hash it was just retrieved from.""" + wire, result = await _wire_and_result(guardrail, RETRIEVE_ECHO_MESSAGES) + + assert not any(row.get("tool_call_id") == "hr_1" for row in wire) + assert not any("RETRIEVED BODY" in json.dumps(row) for row in wire) + # Reaches the model byte-identical, so no marker stands in for the expansion. + assert result["structured_messages"][3] == RETRIEVE_ECHO_MESSAGES[3] + # Negative control: unrelated history is still compressed, not a no-op. + assert any(row.get("content") == "H" * 5000 for row in wire) + + +@pytest.mark.asyncio +async def test_retrieved_content_guard_matches_direct_tool_name(guardrail: HeadroomGuardrail): + """Server-side the tool is named headroom_retrieve (no MCP prefix); its + result must be protected the same way.""" + messages = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "hr_direct", + "type": "function", + "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_direct", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + wire, result = await _wire_and_result(guardrail, messages) + + assert not any(row.get("tool_call_id") == "hr_direct" for row in wire) + assert result["structured_messages"][3] == messages[3] + + +@pytest.mark.asyncio +async def test_retrieved_content_protected_when_mcp_tool_name_is_truncated(guardrail: HeadroomGuardrail): + """A long mcp____headroom_retrieve name is truncated past 64 chars in + the OpenAI-translated view the guardrail scans, dropping the suffix. The call + id read from the request's own Anthropic tool_use (never truncated) still + pairs the retrieved row so it is held back.""" + from litellm.llms.anthropic.pass_through.adapters.transformation import ( + truncate_tool_name, + ) + + long_name = "mcp__" + "s" * 45 + "__" + HEADROOM_RETRIEVE_TOOL_NAME + assert len(long_name) > 64 + truncated = truncate_tool_name(long_name) + assert not truncated.endswith(HEADROOM_RETRIEVE_TOOL_NAME) + + # What the guardrail scans: OpenAI-translated messages with the truncated name. + structured = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "hr_long", "type": "function", "function": {"name": truncated, "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "hr_long", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + # The request's own messages, untranslated: Anthropic tool_use carries the full name. + raw_messages = [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "hr_long", "name": long_name, "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "hr_long", "content": "RETRIEVED BODY"}]}, + ] + + inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(structured))) + sent: dict = {} + + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"]))) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-sonnet-4-5-20250929", "messages": raw_messages}, + input_type="request", + ) + + assert not any(row.get("tool_call_id") == "hr_long" for row in sent["messages"]) + assert result["structured_messages"][3] == structured[3] + assert any(row.get("content") == "H" * 5000 for row in sent["messages"]) + + +def test_raw_retrieve_call_ids_covers_both_shapes_and_ignores_others(): + """Retrieve ids are read from OpenAI tool_calls and Anthropic tool_use blocks; + non-retrieve calls, non-tool_use blocks, string content, and non-list inputs + yield nothing.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import _raw_retrieve_call_ids + + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"id": "oa1", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}, + {"id": "other", "function": {"name": "get_weather"}}, + {"id": "malformed", "function": {"name": 123}}, + {"id": "nofunc"}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "an1", "name": "mcp__hr__headroom_retrieve", "input": {}}, + {"type": "tool_use", "id": "an2", "name": "jira_get_issue", "input": {}}, + {"type": "text", "text": "noise"}, + ], + }, + {"role": "user", "content": "plain string content, not a list"}, + ] + + assert _raw_retrieve_call_ids(messages) == frozenset({"oa1", "an1"}) + assert _raw_retrieve_call_ids("not a list") == frozenset() + assert _raw_retrieve_call_ids(None) == frozenset() + + +@pytest.mark.asyncio +async def test_nothing_compressible_returns_inputs_untouched(guardrail: HeadroomGuardrail): + """A single-turn request is all protected, so there is nothing to send and + the caller's own inputs object comes back.""" + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=[ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "A" * 5000}, + ], + ) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + mock_post.assert_not_called() + assert result is inputs + + +@pytest.mark.asyncio +async def test_fail_open_returns_the_caller_inputs_object(): + """Translation handlers detect a rewrite by object identity, so a request + that was not compressed must come back as the same object or it is + round-tripped through the write-back for nothing.""" + guardrail = _make_guardrail(unreachable_fallback="fail_open") + original = json.loads(json.dumps(AGENTIC_MESSAGES)) + inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=original) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("boom"), + ): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result is inputs + assert result["structured_messages"] is original + + +# --------------------------------------------------------------------------- +# LIT-5018: the retrieval follow-up keeps the model's own text. +# --------------------------------------------------------------------------- + + +def _anthropic_response_with_text_and_tool_call() -> dict: + return { + "content": [ + {"type": "text", "text": "Let me pull the original back."}, + {"type": "tool_use", "id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "input": {"hash": "h" * 24}}, + ] + } + + +async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): + guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-1" + logging_obj.model_call_details = {} + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response("ORIGINAL CONTENT"), + ): + return await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [{"id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": {"hash": "h" * 24}}] + }, + model="claude-sonnet-4-5-20250929", + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={}, + ) + + +@pytest.mark.asyncio +async def test_anthropic_followup_preserves_assistant_text(guardrail: HeadroomGuardrail): + plan = await _plan_for(guardrail, _anthropic_response_with_text_and_tool_call(), [{"role": "user", "content": "q"}]) + + assistant = plan.request_patch.messages[-2] # type: ignore[union-attr] + assert assistant["role"] == "assistant" + # Text first, then the tool_use it accompanied: dropping it loses the + # model's stated reason for the retrieval from its own transcript. + assert assistant["content"][0] == {"type": "text", "text": "Let me pull the original back."} + assert assistant["content"][1]["type"] == "tool_use" + + +@pytest.mark.asyncio +async def test_responses_followup_preserves_assistant_text(guardrail: HeadroomGuardrail): + response = { + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Fetching the original."}]}, + {"type": "function_call", "call_id": "call_1", "name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, + ] + } + + plan = await _plan_for(guardrail, response, [{"role": "user", "content": "q"}]) + + items = plan.request_patch.messages # type: ignore[union-attr] + assert items[1] == {"role": "assistant", "content": "Fetching the original."} + assert items[2]["type"] == "function_call" + + +@pytest.mark.asyncio +async def test_chat_followup_echoes_only_the_retrieve_call(guardrail: HeadroomGuardrail): + """A turn that called another tool alongside headroom_retrieve must not + echo that call: only the retrieve call gets a tool result, and a tool_call + without one is rejected by the provider.""" + other = MagicMock() + other.id = "call_other" + other.type = "function" + other.function = MagicMock() + other.function.name = "Write" + other.function.arguments = "{}" + + response = _make_openai_response_with_tool_call(HEADROOM_RETRIEVE_TOOL_NAME, {"hash": "h" * 24}, "call_1") + response.choices[0].message.content = "Getting the original first." + response.choices[0].message.tool_calls = [response.choices[0].message.tool_calls[0], other] + + plan = await _plan_for(guardrail, response, [{"role": "user", "content": "q"}]) + + messages = plan.request_patch.messages # type: ignore[union-attr] + assistant = messages[1] + assert assistant["content"] == "Getting the original first." + assert [tc["id"] for tc in assistant["tool_calls"]] == ["call_1"] + assert [m["tool_call_id"] for m in messages[2:]] == ["call_1"] + + +# --- LIT-5881: the calls to the compression service must be time-bounded --- + + +def _timeout_of(mock_call) -> httpx.Timeout: + timeout = mock_call.kwargs["timeout"] + assert isinstance(timeout, httpx.Timeout), timeout + return timeout + + +@pytest.mark.asyncio +async def test_compress_call_passes_bounded_timeout(guardrail: HeadroomGuardrail): + """Without an explicit timeout the call inherits the shared client's 600s read leg.""" + inputs = GenericGuardrailAPIInputs(texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as mock_post: + await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + timeout = _timeout_of(mock_post.call_args) + assert timeout.read == 60.0 + assert timeout.write == 60.0 + assert timeout.pool == 60.0 + assert timeout.connect == 5.0 + + +@pytest.mark.asyncio +async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail): + """The retrieval leg runs on the same request and needs the same bound.""" + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response("original"), + ) as mock_get: + result = await guardrail._call_retrieve("a" * 24) + + assert result == "original" + timeout = _timeout_of(mock_get.call_args) + assert timeout.read == 60.0 + assert timeout.connect == 5.0 + + +@pytest.mark.asyncio +async def test_configured_timeout_overrides_the_default(): + """Headroom accepted litellm_params.timeout and ignored it.""" + guardrail = _make_guardrail(timeout=3.5) + inputs = GenericGuardrailAPIInputs(texts=["A" * 5000], structured_messages=ORIGINAL_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as mock_post: + await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + timeout = _timeout_of(mock_post.call_args) + assert timeout.read == 3.5 + assert timeout.connect == 3.5 + + +@pytest.mark.asyncio +async def test_read_timeout_is_surfaced_as_unreachable_under_fail_closed(): + """A stalled service must reach the fail policy, not escape as a 500.""" + guardrail = _make_guardrail() + inputs = GenericGuardrailAPIInputs(texts=["hello"], structured_messages=ORIGINAL_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ReadTimeout("timed out"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert exc_info.value.status_code == 502 + assert "unreachable" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_read_timeout_forwards_uncompressed_under_fail_open(): + guardrail = _make_guardrail(unreachable_fallback="fail_open") + inputs = GenericGuardrailAPIInputs(texts=["hello"], structured_messages=ORIGINAL_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ReadTimeout("timed out"), + ): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result.get("structured_messages") == ORIGINAL_MESSAGES + + +def test_initializer_forwards_configured_timeout(monkeypatch: pytest.MonkeyPatch): + """Wiring it only in __init__ leaves `timeout:` in config.yaml silently ignored.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda callback: None, + ) + params = LitellmParams( + guardrail="headroom", + mode="pre_call", + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + timeout=7.0, + ) + callback = initialize_guardrail(params, {"guardrail_name": "headroom"}) # type: ignore[arg-type] + + assert callback.timeout.read == 7.0 + + +def test_in_place_update_keeps_the_timeout_resolved(): + """The base implementation copies every attribute over, nulling an unset timeout.""" + from litellm.types.guardrails import LitellmParams + + guardrail = _make_guardrail(timeout=5.0) + assert guardrail.timeout.read == 5.0 + + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE) + ) + assert isinstance(guardrail.timeout, httpx.Timeout) + assert guardrail.timeout.read == 60.0 + + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, timeout=7.0) + ) + assert guardrail.timeout.read == 7.0 + + +@pytest.mark.parametrize("configured", [0, 0.0, -1, -30.0, float("inf"), float("-inf"), float("nan")]) +def test_unusable_timeout_falls_back_to_the_default(configured: float): + """0 and inf read as no deadline at all, a negative one as a deadline already past.""" + guardrail = _make_guardrail(timeout=configured) + + assert guardrail.timeout.read == 60.0 + assert guardrail.timeout.connect == 5.0 + + +def _affinity_of(call: "respx.models.Call") -> str | None: + value: Final = call.request.headers.get(AFFINITY_HEADER) + return str(value) if value else None + + +def _captured_affinity(route: respx.Route) -> str | None: + """Read the affinity header off the most recent call on this route. + + route.calls[0] would always be the first request, which is wrong when a test + drives several requests through one route and wants to compare them. + """ + return _affinity_of(route.calls[-1]) + + +@pytest.mark.asyncio +async def test_compress_sends_stable_tenant_scoped_affinity( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, +): + """Related calls in one conversation reach the same pod, and the raw ids never leave.""" + route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={ + "messages": COMPRESSED_MESSAGES, + "tokens_before": 1000, + "tokens_after": 100, + }, + ) + ) + inputs: Final = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data: Final = { + "model": "gpt-4o", + "litellm_session_id": "sess-abc", + "metadata": {"user_api_key_team_id": "team-1"}, + } + + for _ in range(2): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=dict(request_data), + input_type="request", + ) + + sent: Final = [_affinity_of(call) for call in route.calls] + assert len(sent) == 2, f"expected two compress calls, got {len(sent)}" + first: Final = sent[0] + assert first is not None, "an authenticated session must produce an affinity value" + assert sent[0] == sent[1], f"the same conversation must map to one affinity, got {sent}" + assert len(first) == 32 + assert "sess-abc" not in first + assert "team-1" not in first + + +@pytest.mark.asyncio +async def test_affinity_differs_per_tenant_for_the_same_session_id( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, +): + """A session id is reusable across tenants, so it cannot be the affinity on its own.""" + route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, + ) + ) + inputs: Final = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + + seen: Final[list[str | None]] = [] + for team_id in ("team-A", "team-B"): + await guardrail.apply_guardrail( + inputs=inputs, + request_data={ + "model": "gpt-4o", + "litellm_session_id": "shared-session", + "metadata": {"user_api_key_team_id": team_id}, + }, + input_type="request", + ) + seen.append(_affinity_of(route.calls[-1])) + + assert seen[0] is not None and seen[1] is not None + assert seen[0] != seen[1], f"different tenants must differ, got {seen}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_extra", [{}, {"user_api_key_team_alias": "chosen-by-caller"}]) +async def test_affinity_omitted_without_an_authenticated_tenant( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + request_extra: dict, +): + """Without a trusted identity the header is dropped rather than risking a cross-tenant route.""" + route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, + ) + ) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ), + request_data={"model": "gpt-4o", "litellm_session_id": "sess-abc", **request_extra}, + input_type="request", + ) + + assert _captured_affinity(route) is None + + +@pytest.mark.asyncio +async def test_affinity_omitted_for_a_proxy_generated_session_id( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, +): + """A generated id is unique per request, so pinning on it would defeat the cache.""" + route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, + ) + ) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ), + request_data={ + "model": "gpt-4o", + "litellm_session_id": "generated-per-request", + "metadata": { + "litellm_session_id_generated": True, + "user_api_key_team_id": "team-1", + }, + }, + input_type="request", + ) + + assert _captured_affinity(route) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("container", ["metadata", "litellm_metadata"]) +async def test_affinity_uses_the_nested_identity_a_proxied_request_carries( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + container: str, +): + """A proxied request keeps its identity nested, not at the top level. + + Reading only the top level would silently disable affinity for every real + proxy call, so each container the proxy actually populates is covered. + """ + route: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, + ) + ) + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ), + request_data={ + "model": "gpt-4o", + "litellm_session_id": "sess-abc", + container: {"user_api_key_team_id": "team-1"}, + }, + input_type="request", + ) + + assert _captured_affinity(route) is not None, f"{container} must supply the tenant identity" + + +@pytest.mark.asyncio +async def test_retrieve_replays_the_same_affinity_the_compress_call_sent( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, +): + """The retrieve must carry the identical value, or it lands on a different pod. + + Asserting the header merely exists would not catch a replay that recomputes a + different digest, which is the failure this is here to prevent. + """ + compress: Final = respx_mock.post(f"{FAKE_API_BASE}/v1/compress").mock( + return_value=httpx.Response( + 200, + json={"messages": COMPRESSED_MESSAGES, "tokens_before": 1000, "tokens_after": 100}, + ) + ) + request_data: Final = { + "litellm_call_id": "call-1", + "litellm_session_id": "sess-abc", + "metadata": {"user_api_key_team_id": "team-1"}, + } + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ), + request_data=dict(request_data), + input_type="request", + ) + + compress_affinity: Final = _captured_affinity(compress) + assert compress_affinity is not None, "authenticated session must produce an affinity value" + + retrieve: Final = respx_mock.get(f"{FAKE_API_BASE}/v1/retrieve/{CCR_HASH}").mock( + return_value=httpx.Response(200, json={"original_content": "restored"}) + ) + guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({CCR_HASH}), time.monotonic() + 999) + + await guardrail.async_build_agentic_loop_plan( + tools={"tool_calls": [{"arguments": {"hash": CCR_HASH}}]}, + model="gpt-4o", + messages=list(ORIGINAL_MESSAGES), + response={}, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs=dict(request_data), + ) + + assert _captured_affinity(retrieve) == compress_affinity