diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index d864ec04a8d..dbd565f5d5b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -12829,6 +12829,18 @@ "description": "Custom advisory message template used when on_flagged='inject_system_message'. Must contain a {reason} placeholder. Defaults to a generic advisory message if unset.", "title": "Advisory System Message" }, + "affinity_salt": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Salt mixed into the x-headroom-affinity digest, which routes related calls to the same pod so its cache can be reused. Set this to vary or rotate the digest across deployments. Falls back to a fixed default.", + "title": "Affinity Salt" + }, "akto_account_id": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 36d48d49c7e..9ac586579fa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json import math import re @@ -17,7 +18,10 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.compression.compress import get_protected_indices -from litellm.constants import HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS +from litellm.constants import ( + HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, + SESSION_ID_GENERATED_METADATA_KEY, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, @@ -55,6 +59,20 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel BYPASS_HEADER: Final = "x-headroom-bypass" +# Routes related Headroom calls to the same pod so its pod-local cache can be +# reused. The value is derived per tenant (see ``_affinity_key``) and carries no +# credential, so Headroom needs no end-user authentication of its own. +AFFINITY_HEADER: Final = "x-headroom-affinity" +# Fallback salt for the affinity digest when the guardrail config does not set +# one. Deployments that need the digest to differ across instances (or that want +# it rotated) set ``affinity_salt`` in the guardrail config instead. +_DEFAULT_AFFINITY_SALT: Final = "litellm-headroom-affinity-v1" +# How much of the digest to forward. 32 hex chars is far past collision range for +# a routing hint while staying short enough to log. +_AFFINITY_KEY_LENGTH: Final = 32 +# Identities the proxy has already authenticated, most specific first. A client +# supplied tenant name is deliberately absent: it cannot establish identity. +_TENANT_ID_KEYS: Final = ("user_api_key_team_id", "user_api_key_org_id", "user_api_key_user_id", "user_api_key_hash") _STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses) ) @@ -78,6 +96,76 @@ def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isin return isinstance(value, list) +def _explicit_session_id(request_data: dict[str, object]) -> str | None: + """The caller's own session id, or None when there is none worth pinning. + + A session id the proxy generated for a request that arrived without one is + unique per request, so affinity built on it would send every call to a + different pod and the cache would never hit. That case is marked in metadata + and skipped, matching how session affinity is resolved for other providers. + """ + metadata: Final = request_data.get("metadata") + if _is_str_object_dict(metadata) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + return None + for key in ("litellm_session_id", "session_id"): + value: Final = request_data.get(key) + if isinstance(value, str) and value: + return value + if _is_str_object_dict(metadata): + value: Final = metadata.get("session_id") + if isinstance(value, str) and value: + return value + return None + + +def _tenant_id(request_data: dict[str, object]) -> str | None: + """The authenticated identity this request runs as, or None when anonymous. + + Only keys the proxy itself resolved from a validated key are considered: a + client supplied tenant name never establishes identity, so two tenants + cannot end up sharing an affinity value by naming themselves the same thing. + + The identity keys live under ``metadata`` / ``litellm_metadata`` on a proxied + request rather than at the top level (see + ``litellm.proxy.route_llm_request.get_team_id_from_data``), so both nested + containers are searched. The top level is checked too for SDK calls, where + they do sit flat. + """ + containers: Final[list[Mapping[str, object]]] = [request_data] + for container_key in ("metadata", "litellm_metadata"): + nested: Final = request_data.get(container_key) + if _is_str_object_dict(nested): + containers.append(nested) + + for container in containers: + for key in _TENANT_ID_KEYS: + value: Final = container.get(key) + if isinstance(value, str) and value: + return f"{key}={value}" + if isinstance(value, int) and not isinstance(value, bool): + return f"{key}={value}" + return None + + +def _affinity_key(request_data: dict[str, object], salt: str = _DEFAULT_AFFINITY_SALT) -> str | None: + """Opaque, tenant-scoped routing value for related Headroom calls, or None. + + Returns None unless both halves are present: without a session id there is + nothing to keep together, and without a trusted tenant identity the same + session id could belong to different tenants, so affinity is omitted rather + than risk crossing a tenant boundary. The raw session id never leaves the + proxy: Headroom only ever sees a salted digest. + """ + session_id: Final = _explicit_session_id(request_data) + if not session_id: + return None + tenant: Final = _tenant_id(request_data) + if not tenant: + return None + digest: Final = hashlib.sha256(f"{salt or _DEFAULT_AFFINITY_SALT}|{tenant}|{session_id}".encode()).hexdigest() + return digest[:_AFFINITY_KEY_LENGTH] + + def _flatten_messages_for_compression(messages: list[dict[str, object]]) -> list[dict[str, object]]: """Collapse all-text list-of-parts content to plain strings for /v1/compress. @@ -496,6 +584,7 @@ class HeadroomGuardrail(CustomGuardrail): unreachable_fallback: str | None = None, timeout: float | None = None, ccr_retrieval: bool = True, + affinity_salt: str | None = None, ): self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: @@ -509,6 +598,11 @@ class HeadroomGuardrail(CustomGuardrail): "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) self.ccr_retrieval = ccr_retrieval + # Salted so the digest Headroom receives cannot be replayed to recover + # another tenant's session id. Deployments can rotate or vary it per + # instance via the guardrail config; the default keeps the value stable + # for the common single-deployment case. + self.affinity_salt: str = affinity_salt or _DEFAULT_AFFINITY_SALT self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) @@ -531,10 +625,12 @@ class HeadroomGuardrail(CustomGuardrail): value: Final = headers.get(BYPASS_HEADER) return str(value).lower() == "true" - def _request_headers(self) -> dict[str, str]: + def _request_headers(self, affinity: str | None = None) -> dict[str, str]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} if self.headroom_api_key: headers["Authorization"] = f"Bearer {self.headroom_api_key}" + if affinity: + headers[AFFINITY_HEADER] = affinity return headers @staticmethod @@ -587,6 +683,7 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, + affinity: str | None = None, ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: @@ -596,7 +693,7 @@ class HeadroomGuardrail(CustomGuardrail): raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.headroom_api_base}/v1/compress", json=payload, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except httpx.HTTPStatusError as e: @@ -725,7 +822,7 @@ class HeadroomGuardrail(CustomGuardrail): stats["tokens_saved"] = tokens_before - tokens_after return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) - async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: + async def _call_retrieve(self, hash_value: str, query: str | None = None, affinity: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: params["query"] = query @@ -734,7 +831,7 @@ class HeadroomGuardrail(CustomGuardrail): raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get is untyped url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}", params=params, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: @@ -814,6 +911,7 @@ class HeadroomGuardrail(CustomGuardrail): result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, + affinity=_affinity_key(_REQUEST_DATA_ADAPTER.validate_python(request_data), self.affinity_salt), ) end_time: Final = time.time() @@ -934,6 +1032,11 @@ class HeadroomGuardrail(CustomGuardrail): valid_hashes = self._issued_hashes_by_call_id.get(call_id, (frozenset(), 0.0))[0] if call_id else frozenset() retrieved: Final[list[tuple[dict[str, object], str]]] = [] + # The compress call that issued these hashes ran under the originating + # request, so its affinity value is what has to be replayed here for the + # retrieve to land on the pod holding the entry. kwargs carries the same + # request_data shape apply_guardrail sees. + affinity: Final = _affinity_key(_REQUEST_DATA_ADAPTER.validate_python(kwargs), self.affinity_salt) for tc in tool_calls: arguments = tc.get("arguments", {}) raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" @@ -954,6 +1057,7 @@ class HeadroomGuardrail(CustomGuardrail): content = await self._call_retrieve( hash_value=hash_value, query=str(query) if query else None, + affinity=affinity, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) retrieved.append((tc, content)) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..ead16f784e9 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -30,6 +30,14 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): default=True, description="Inject the Headroom retrieval tool for hashes declared by the compression service.", ) + affinity_salt: str | None = Field( + default=None, + description=( + "Salt mixed into the x-headroom-affinity digest, which routes related calls to the same " + "pod so its cache can be reused. Set this to vary or rotate the digest across deployments. " + "Falls back to a fixed default." + ), + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index f3456f6b60c..b98dbe22284 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1,3016 +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 -""" - -import json -import time -from typing import 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 ( - 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 +""" +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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cbaea37d409..1b1d160aeae 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36585,6 +36585,11 @@ export interface components { * @description Custom advisory message template used when on_flagged='inject_system_message'. Must contain a {reason} placeholder. Defaults to a generic advisory message if unset. */ advisory_system_message?: string | null; + /** + * Affinity Salt + * @description Salt mixed into the x-headroom-affinity digest, which routes related calls to the same pod so its cache can be reused. Set this to vary or rotate the digest across deployments. Falls back to a fixed default. + */ + affinity_salt?: string | null; /** * Akto Account Id * @description Akto account ID for multi-tenant deployments. Env: AKTO_ACCOUNT_ID. Default: '1000000'.