From e120419dc5d76e7d08ae4d73a99d355ff61d4749 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 3 Nov 2026 08:52:03 +0530 Subject: [PATCH] fix(headroom): send the whole request on Headroom's gateway turn The guardrail sent Headroom only a subset of chat rows, flattened to strings, and withheld the newest results. Claude Code traffic got no compression, other clients paid more than with no guardrail because history was rewritten and the prompt cache busted, and Read results were compressed lossily. Now the guardrail sends the provider request (messages or input, system, tools) with Headroom's gateway block and config.mode "lossless", and writes back what Headroom returns. History stays byte-stable, nothing is lost, and server-side plugins can rewrite tools and skills. Requests with no authenticated caller and older Headroom servers keep the row-subset path, and an older server is asked again every 5 minutes. --- litellm/proxy/_lazy_openapi_snapshot.json | 26 +- .../guardrail_hooks/headroom/__init__.py | 2 + .../guardrail_hooks/headroom/headroom.py | 371 ++++++++++++++---- .../guardrails/guardrail_hooks/headroom.py | 18 +- .../guardrail_hooks/test_headroom.py | 341 +++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 14 +- 6 files changed, 690 insertions(+), 82 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5a80583b207..4592a1e23b7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -17877,7 +17877,7 @@ }, "ccr_retrieval": { "default": true, - "description": "Inject the Headroom retrieval tool for hashes declared by the compression service.", + "description": "Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false.", "title": "Ccr Retrieval", "type": "boolean" }, @@ -17930,6 +17930,24 @@ "description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.", "title": "Client Secret" }, + "compression_mode": { + "anyOf": [ + { + "enum": [ + "lossless", + "lossy_inline", + "lossless_then_lossy" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "default": "lossless", + "description": "Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; 'lossy_inline' saves more without a way back. Null uses the Headroom server's default.", + "title": "Compression Mode" + }, "confidence_threshold": { "default": 0.5, "default_value": 0.5, @@ -18271,6 +18289,12 @@ "description": "noma_v2 only: name of this gateway, used as the gateway_host label on Noma scans", "title": "Gateway Name" }, + "gateway_turn": { + "default": true, + "description": "Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail compresses a subset of rows instead.", + "title": "Gateway Turn", + "type": "boolean" + }, "grounding_check": { "anyOf": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py index b7e4d9275fa..8e5b8e422bb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py @@ -37,6 +37,8 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> unreachable_fallback=litellm_params.unreachable_fallback, timeout=litellm_params.timeout, ccr_retrieval=litellm_params.ccr_retrieval, + gateway_turn=litellm_params.gateway_turn, + compression_mode=litellm_params.compression_mode, ) litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 0b762b2d5fe..ad9c98c7031 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -54,7 +54,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( HEADROOM_REQUEST_CONTEXT_KEY, HeadroomRequestContext, ) -from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, GuardrailStatus if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -72,6 +72,21 @@ _COMPRESS_TIMEOUT_SECONDS: Final = 60.0 HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"[a-f0-9]{12,24}") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 +# The provider request Headroom's gateway turn reads, and the part it may rewrite. +# Everything else in request_data is LiteLLM's own or the client's to keep. +_GATEWAY_REWRITABLE_KEYS: Final = ("messages", "input", "system", "instructions", "tools", "tool_choice") +_GATEWAY_BODY_KEYS: Final = (*_GATEWAY_REWRITABLE_KEYS, "model", "max_tokens", "max_output_tokens", "stream") +# How long a Headroom without the gateway turn is left alone before it's asked again, +# so a rolling upgrade or one old replica doesn't switch the feature off for good. +_GATEWAY_UNSUPPORTED_RETRY_SECONDS: Final = 300.0 +_COMPRESS_STAT_KEYS: Final = ( + "tokens_before", + "tokens_after", + "tokens_saved", + "tool_tokens_saved", + "compression_ratio", + "transforms_applied", +) # Narrows the base class's bare-dict ``request_data`` at the boundary so its # untranslated messages can be read with concrete types (values pass through by # reference, so this is a shallow top-level reconstruction). @@ -337,6 +352,37 @@ class _CompressResult: ccr_hashes: frozenset[str] = frozenset() +@dataclass(frozen=True, slots=True) +class _TurnAnswered: + answer: Mapping[str, object] + body: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class _TurnUnsupported: + status_code: int + + +@dataclass(frozen=True, slots=True) +class _TurnFailed: + """Headroom failed, and fail_open lets the request through as it was.""" + + +def _is_count(value: object) -> TypeGuard[int | float]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _compression_stats(body: Mapping[str, object]) -> Mapping[str, object]: + """The savings Headroom reported. Spend tracking (extract_compression_saved_tokens) + reads only tokens_saved, which the service can omit, so it's derived then.""" + stats: Final = {key: body[key] for key in _COMPRESS_STAT_KEYS if key in body} + before: Final = stats.get("tokens_before") + after: Final = stats.get("tokens_after") + if "tokens_saved" in stats or not _is_count(before) or not _is_count(after): + return stats + return {**stats, "tokens_saved": before - after} + + def _build_headroom_retrieve_tool() -> dict[str, object]: return { "type": "function", @@ -364,20 +410,110 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: } -def _session_affinity(request_data: Mapping[str, object]) -> str | None: +def _request_tenant(request_data: Mapping[str, object]) -> tuple[object, ...] | None: + """Who is calling, or None for a request with no caller on it.""" auth: Final = request_user_api_key_auth(request_data) - session_id: Final = get_caller_session_id(request_data) - if not isinstance(auth, UserAPIKeyAuth) or not session_id: + if not isinstance(auth, UserAPIKeyAuth): return None tenant: Final = ( ("key", auth.api_key) if auth.api_key else ("principal", auth.org_id, auth.team_id, auth.user_id, auth.agent_id) ) - if not any(tenant[1:]): + return tenant if any(tenant[1:]) else None + + +def _session_affinity(request_data: Mapping[str, object]) -> str | None: + session_id: Final = get_caller_session_id(request_data) + tenant: Final = _request_tenant(request_data) + if not session_id or tenant is None: return None identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":")) return hashlib.sha256(identity.encode("utf-8")).hexdigest() +def _without_cache_control(part: object) -> object: + return {key: value for key, value in part.items() if key != "cache_control"} if _is_str_object_dict(part) else part + + +def _conversation_identity(row: Mapping[str, object]) -> object: + """A row as a conversation key. Clients move cache_control between turns, and + with it can resend the same text as a part list instead of a string.""" + content: Final = row.get("content") + if isinstance(content, str) or is_all_text_parts(content): + return content_to_text(content) + return [_without_cache_control(part) for part in content] if _is_object_list(content) else content + + +def _gateway_session_id( + request_data: Mapping[str, object], rows: Sequence[object], model: object, tenant: tuple[object, ...] +) -> str | None: + """One Headroom session per conversation, or None to stay stateless. + + Headroom replays what it already sent in a session, which keeps the provider's + prompt cache warm, and turns in one session queue on one lock. So only the + caller's own session id joins requests into a session: one app key can serve + many end users whose conversations all start the same way. Claude Code's + subagents share their parent's session id, and the first user row tells them apart. + """ + caller_session: Final = get_caller_session_id(request_data) + if not caller_session: + return None + first_user_row: Final = next((row for row in rows if _is_str_object_dict(row) and row.get("role") == "user"), None) + identity: Final = json.dumps( + [ + "headroom-session-v2", + tenant, + request_data.get("user"), + caller_session, + model, + _conversation_identity(first_user_row) if first_user_row is not None else None, + ], + sort_keys=True, + default=str, + ) + return f"litellm-{hashlib.sha256(identity.encode('utf-8')).hexdigest()[:32]}" + + +def _client_anthropic_beta(request_data: Mapping[str, object]) -> str | None: + proxy_request: Final = request_data.get("proxy_server_request") + headers: Final = proxy_request.get("headers") if _is_str_object_dict(proxy_request) else None + if not _is_str_object_dict(headers): + return None + return next( + (value for name, value in headers.items() if name.lower() == "anthropic-beta" and isinstance(value, str)), + None, + ) + + +def _is_gateway_turn_unsupported(response: HttpxResponse, body: object) -> bool: + """An older Headroom: it rejects the mode, or answers in the legacy messages-only shape.""" + if response.status_code == 400: + return "config.mode" in response.text + return _is_str_object_dict(body) and "messages" in body and "body" not in body + + +def _json_or_none(response: HttpxResponse) -> object: + """A 200 answer's JSON body, or None for any other status or a body that isn't JSON.""" + if response.status_code != 200: + return None + try: + return _json_body(response) + except ValueError: + return None + + +def _merged_extra_headers(existing: object, added: object) -> Mapping[str, object] | None: + """The request's extra headers plus the ones Headroom asks for, anthropic-beta values combined.""" + if not _is_str_object_dict(added) or not added: + return None + current: Final = existing if _is_str_object_dict(existing) else {} + if "anthropic-beta" not in current or "anthropic-beta" not in added: + return {**current, **added} + from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers + + beta: Final = merge_anthropic_beta_headers(str(current["anthropic-beta"]), str(added["anthropic-beta"])) + return {**current, **added, "anthropic-beta": beta} + + def has_headroom_retrieve_tool(tools: object) -> bool: return has_tool_with_name(tools, HEADROOM_RETRIEVE_TOOL_NAME) @@ -515,6 +651,8 @@ class HeadroomGuardrail(CustomGuardrail): unreachable_fallback: str | None = None, timeout: float | None = None, ccr_retrieval: bool = True, + gateway_turn: bool = True, + compression_mode: Literal["lossless", "lossy_inline", "lossless_then_lossy"] | None = "lossless", ): self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: @@ -527,6 +665,9 @@ class HeadroomGuardrail(CustomGuardrail): self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) + self.gateway_turn = gateway_turn + self._gateway_turn_retry_at = 0.0 + self.compression_mode = compression_mode self.ccr_retrieval = ccr_retrieval self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, @@ -607,34 +748,17 @@ class HeadroomGuardrail(CustomGuardrail): if self.ccr_retrieval: payload["config"] = {"mode": "ccr"} - try: - 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(affinity), - timeout=self.timeout, - ) - except httpx.HTTPStatusError as e: - return _CompressResult( - self._handle_compress_failure( - messages, - "Headroom compression service returned an error", - _build_compress_failure_detail(e.response.status_code, e.response.text), - ), - False, - {}, - ) - except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: + response: Final = await self._post_compress(payload, affinity) + if isinstance(response, str): return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service unreachable", - {"detail": str(e)}, + {"detail": response}, ), False, {}, ) - response: Final[HttpxResponse] = raw_response if response.status_code != 200: return _CompressResult( @@ -714,32 +838,146 @@ class HeadroomGuardrail(CustomGuardrail): body.get("compression_ratio", 0), ) - stats: Final = { - key: body[key] - for key in ( - "tokens_before", - "tokens_after", - "tokens_saved", - "compression_ratio", - "transforms_applied", - ) - if key in body - } - tokens_before: Final = stats.get("tokens_before") - tokens_after: Final = stats.get("tokens_after") - if ( - "tokens_saved" not in stats - and isinstance(tokens_before, (int, float)) - and not isinstance(tokens_before, bool) - and isinstance(tokens_after, (int, float)) - and not isinstance(tokens_after, bool) - ): - # Spend tracking (extract_compression_saved_tokens) reads only - # tokens_saved, which the live compression service omits; derive it - # so savings are counted, but let a service-sent value win. - stats["tokens_saved"] = tokens_before - tokens_after + stats: Final = dict(_compression_stats(body)) return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) + async def _post_compress(self, payload: Mapping[str, object], affinity: str | None) -> HttpxResponse | str: + """Headroom's HTTP answer, or why it could not be reached.""" + try: + return await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped + url=f"{self.headroom_api_base}/v1/compress", + json=dict(payload), + headers=self._request_headers(affinity), + timeout=self.timeout, + ) + except httpx.HTTPStatusError as e: + # AsyncHTTPHandler raises on any non-2xx, and the status decides what happens next. + return e.response + except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: + return str(e) + + async def _call_gateway_turn( + self, payload: Mapping[str, object], affinity: str | None + ) -> _TurnAnswered | _TurnUnsupported | _TurnFailed: + """A failed call goes through the configured fallback, so fail_closed raises here.""" + first: Final = await self._post_compress(payload, affinity) + # A 503 means a session's turn lock timed out, and Headroom asks the gateway to retry. + retry: Final = not isinstance(first, str) and first.status_code == 503 + sent: Final = await self._post_compress(payload, affinity) if retry else first + if isinstance(sent, str): + _ = self._handle_compress_failure([], "Headroom compression service unreachable", {"detail": sent}) + return _TurnFailed() + answer: Final = _json_or_none(sent) + if _is_gateway_turn_unsupported(sent, answer): + return _TurnUnsupported(sent.status_code) + body: Final = answer.get("body") if _is_str_object_dict(answer) else None + if not _is_str_object_dict(answer) or not _is_str_object_dict(body): + _ = self._handle_compress_failure( + [], + "Headroom compression service returned an error", + _build_compress_failure_detail(sent.status_code, sent.text), + ) + return _TurnFailed() + return _TurnAnswered(answer, body) + + async def _apply_gateway_turn( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], # mutable-ok: the rewritten request is written back into request_data + ) -> GenericGuardrailAPIInputs | None: + """Send Headroom the whole provider request and take back what it rewrites. + + Reads and writes request_data, not the translated inputs: on the Anthropic + route those are OpenAI-shaped copies, while Headroom's folds and hooks need + the wire shape. The inputs handed back carry no texts or tools and the same + structured_messages object, so the translation handlers write nothing over + the rewritten request. + + None means this request takes the row-subset path: it isn't a caller's live + request (the auto-router's routing hop passes a throwaway request_data and + reads the result's structured_messages), or Headroom is too old. + """ + tenant: Final = _request_tenant(request_data) + rows: Final = request_data["messages"] if "messages" in request_data else request_data.get("input") + if tenant is None or "structured_messages" not in inputs or not _is_object_list(rows) or not rows: + return None + + model: Final = self.headroom_model or request_data.get("model") + session_id: Final = _gateway_session_id(request_data, rows, model, tenant) + config: Final = { + **({"session_id": session_id} if session_id else {}), + **({"mode": self.compression_mode} if self.compression_mode else {}), + } + client_beta: Final = _client_anthropic_beta(request_data) + payload: Final = { + **{key: request_data[key] for key in _GATEWAY_BODY_KEYS if key in request_data}, + **({"model": model} if model else {}), + "config": config, + "gateway": { + "can_redrive": False, + "can_relay_response": False, + "request_headers": {"anthropic-beta": client_beta} if client_beta else {}, + }, + } + affinity: Final = _session_affinity(request_data) or session_id + start_time: Final = time.time() + outcome: Final = await self._call_gateway_turn(payload, affinity) + end_time: Final = time.time() + match outcome: + case _TurnUnsupported(status_code=status_code): + self._gateway_turn_retry_at = time.monotonic() + _GATEWAY_UNSUPPORTED_RETRY_SECONDS + verbose_proxy_logger.warning( + "Headroom at %s has no lossless gateway turn (HTTP %s); using row-subset compression and " + "asking again in %d seconds. Upgrade headroom-ai to get it.", + self.headroom_api_base, + status_code, + _GATEWAY_UNSUPPORTED_RETRY_SECONDS, + ) + return None + case _TurnFailed(): + self._record_outcome( + request_data, + "guardrail_failed_to_respond", + {"error": "headroom compression unavailable; request forwarded uncompressed"}, + start_time, + end_time, + ) + return inputs + case _TurnAnswered(answer=answer, body=body): + # A key Headroom sent back without was dropped on purpose (a tool_choice naming + # a tool it deferred, say), so it goes from the request too. + dropped: Final = tuple(key for key in _GATEWAY_REWRITABLE_KEYS if key in payload and key not in body) + for key in dropped: + _ = request_data.pop(key, None) + rewrites: Final = {key: body[key] for key in _GATEWAY_REWRITABLE_KEYS if key in body} + extra_headers: Final = _merged_extra_headers(request_data.get("extra_headers"), answer.get("headers")) + request_data.update({**rewrites, **({"extra_headers": extra_headers} if extra_headers else {})}) + self._record_outcome(request_data, "success", _compression_stats(answer), start_time, end_time) + return GenericGuardrailAPIInputs(structured_messages=inputs["structured_messages"]) + + def _record_outcome( + self, + request_data: dict[str, object], # mutable-ok: LiteLLM's guardrail logging writes into request_data + status: GuardrailStatus, + response: Mapping[str, object], + start_time: float, + end_time: float, + ) -> None: + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=response, + request_data=request_data, + guardrail_status=status, + guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) + async def _call_retrieve(self, hash_value: str, affinity: str | None, query: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: @@ -799,6 +1037,11 @@ class HeadroomGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Headroom: background request; skipping compression") return inputs + if self.gateway_turn and time.monotonic() >= self._gateway_turn_retry_at: + turned: Final = await self._apply_gateway_turn(inputs, request_data) + if turned is not None: + return turned + structured_messages: Final = inputs.get("structured_messages") if not _is_object_list(structured_messages) or not structured_messages: return inputs @@ -833,21 +1076,14 @@ class HeadroomGuardrail(CustomGuardrail): ) end_time: Final = time.time() - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) - if not result.succeeded: - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"}, - request_data=request_data, - guardrail_status="guardrail_failed_to_respond", - guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, - start_time=start_time, - end_time=end_time, - duration=end_time - start_time, + self._record_outcome( + request_data, + "guardrail_failed_to_respond", + {"error": "headroom compression unavailable; request forwarded uncompressed"}, + start_time, + end_time, ) - add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) # Hand back the caller's own inputs object. Translation handlers # detect "the guardrail rewrote the messages" by identity, so # returning a rebuilt copy sends an unchanged request through the @@ -860,16 +1096,7 @@ class HeadroomGuardrail(CustomGuardrail): protected_indices=protected_indices, ) - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=result.stats, - request_data=request_data, - guardrail_status="success", - guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, - start_time=start_time, - end_time=end_time, - duration=end_time - start_time, - ) - add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) + self._record_outcome(request_data, "success", result.stats, start_time, end_time) hashes: Final = result.ccr_hashes if self.ccr_retrieval else frozenset() if not hashes: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 76e806e3899..60e8b978bc9 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -38,7 +38,23 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): ) ccr_retrieval: bool = Field( default=True, - description="Inject the Headroom retrieval tool for hashes declared by the compression service.", + description="Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false.", + ) + gateway_turn: bool = Field( + default=True, + description=( + "Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway " + "contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools " + "and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail " + "compresses a subset of rows instead." + ), + ) + compression_mode: Literal["lossless", "lossy_inline", "lossless_then_lossy"] | None = Field( + default="lossless", + description=( + "Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; " + "'lossy_inline' saves more without a way back. Null uses the Headroom server's default." + ), ) @staticmethod diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 0cdcb2e2d07..cea9acdb496 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -33,6 +33,9 @@ from fastapi import HTTPException import litellm from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler +from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( HEADROOM_RETRIEVE_TOOL_NAME, HeadroomGuardrail, @@ -95,7 +98,7 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail: default_on=True, ) defaults.update(kwargs) - return HeadroomGuardrail(**defaults) + return HeadroomGuardrail(**{"gateway_turn": False, **defaults}) def _make_compress_response( @@ -610,12 +613,17 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( @pytest.mark.asyncio -@pytest.mark.parametrize("context", [ - None, {"issued_hashes": ["deadbeef000000000000dead"]}, - HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0), -]) +@pytest.mark.parametrize( + "context", + [ + None, + {"issued_hashes": ["deadbeef000000000000dead"]}, + HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0), + ], +) async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, context: object, + guardrail: HeadroomGuardrail, + context: object, ) -> None: """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 @@ -921,7 +929,9 @@ HASH_SHAPED_HISTORY = [ async def _apply( - guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None, + guardrail: HeadroomGuardrail, + messages: list, + ccr_hashes: list | None = None, request_data: dict[str, object] | None = None, ) -> dict: data: Final = request_data if request_data is not None else {"model": "claude-sonnet-5"} @@ -3149,3 +3159,320 @@ def test_affinity_requires_trusted_identity_and_caller_session( affinity: Final = _session_affinity(candidate) assert (affinity is None) is (session is None or not isinstance(auth, UserAPIKeyAuth)) assert (affinity == _session_affinity(baseline)) is (candidate == baseline) + + +def _gateway_answer(body: dict, **extra: object) -> MagicMock: + mock = MagicMock() + mock.status_code = 200 + mock.text = "{}" + mock.json.return_value = { + "messages": body.get("messages"), + "body": body, + "tokens_before": 100, + "tokens_after": 40, + **extra, + } + return mock + + +def _raised_status(status: int, text: str) -> httpx.HTTPStatusError: + """What AsyncHTTPHandler.post does with any non-2xx answer.""" + request = httpx.Request("POST", "http://headroom/v1/compress") + return httpx.HTTPStatusError(text, request=request, response=httpx.Response(status, text=text, request=request)) + + +def _caller(api_key: str = "hashed-key") -> dict: + return {"metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key=api_key)}} + + +def _gateway_guardrail(**kwargs) -> HeadroomGuardrail: + return _make_guardrail(gateway_turn=True, **kwargs) + + +def _posting(guardrail: HeadroomGuardrail, **mock: object): + return patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, **mock) + + +async def _apply_turn(guardrail: HeadroomGuardrail, data: dict) -> object: + rows = data["messages"] if "messages" in data else data["input"] + return await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(structured_messages=rows), request_data=data, input_type="request" + ) + + +async def _session_of(guardrail: HeadroomGuardrail, data: dict) -> str | None: + with _posting(guardrail, return_value=_gateway_answer({"messages": data["messages"]})) as post: + await _apply_turn(guardrail, data) + return post.await_args.kwargs["json"]["config"].get("session_id") + + +ANTHROPIC_TOOL = {"name": "Read", "description": "long", "input_schema": {"type": "object"}} +SHORT_ANTHROPIC_TOOL = {"name": "Read", "description": "short", "input_schema": {"type": "object"}} + + +@pytest.mark.asyncio +async def test_gateway_turn_on_anthropic_route_keeps_the_wire_shape_and_applies_the_answer(): + tool_result = { + "type": "tool_result", + "tool_use_id": "t1", + "content": "text", + "cache_control": {"type": "ephemeral"}, + } + messages = [ + {"role": "user", "content": "task"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}]}, + {"role": "user", "content": [tool_result]}, + ] + folded = [ + *messages[:2], + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "f"}]}, + ] + data = { + "model": "claude-opus-5-5", + "system": "system", + "tools": [ANTHROPIC_TOOL], + "max_tokens": 1024, + "stream": True, + "temperature": 0.2, + "messages": messages, + **_caller(), + } + guardrail = _gateway_guardrail() + answer = _gateway_answer({"messages": folded, "tools": [SHORT_ANTHROPIC_TOOL], "system": "trimmed"}) + + with _posting(guardrail, return_value=answer) as post: + result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + payload = post.await_args.kwargs["json"] + assert set(payload) == {"model", "system", "tools", "max_tokens", "stream", "messages", "config", "gateway"} + assert payload["messages"][2]["content"][0]["type"] == "tool_result" + assert payload["tools"] == [ANTHROPIC_TOOL] + assert payload["config"] == {"mode": "lossless"} + assert result["messages"] == folded + assert result["tools"] == [SHORT_ANTHROPIC_TOOL] + assert result["system"] == "trimmed" + assert result["temperature"] == 0.2 + stats = _recorded_guardrail_response(result) + assert stats["tokens_saved"] == 60 + entry = result["metadata"]["standard_logging_guardrail_information"][0] + assert extract_compression_saved_tokens({"guardrail_information": [entry]}) == 60 + + +@pytest.mark.asyncio +async def test_gateway_turn_on_chat_route_applies_the_answer_and_drops_what_headroom_dropped(): + tool = {"type": "function", "function": {"name": "Read", "description": "long", "parameters": {}}} + short = {"type": "function", "function": {"name": "Read", "description": "short", "parameters": {}}} + call = {"id": "c", "type": "function", "function": {"name": "Read", "arguments": "{}"}} + messages = [ + {"role": "user", "content": "task"}, + {"role": "assistant", "content": None, "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "c", "content": "text"}, + ] + folded = [*messages[:2], {"role": "tool", "tool_call_id": "c", "content": "folded"}] + forced = {"type": "function", "function": {"name": "Read"}} + data = {"model": "gpt-5.5", "tools": [tool], "tool_choice": forced, "messages": messages, **_caller()} + guardrail = _gateway_guardrail() + + with _posting(guardrail, return_value=_gateway_answer({"messages": folded, "tools": [short]})) as post: + result = await OpenAIChatCompletionsHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert post.await_args.kwargs["json"]["tool_choice"] == forced + assert result["messages"] == folded + assert result["tools"] == [short] + assert "tool_choice" not in result + + +@pytest.mark.asyncio +async def test_gateway_turn_on_responses_route_applies_the_answer(): + items = [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "task"}]}, + {"type": "function_call", "call_id": "c", "name": "Read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c", "output": "text"}, + ] + folded = [*items[:2], {"type": "function_call_output", "call_id": "c", "output": "folded"}] + tool = {"type": "function", "name": "Read", "description": "long", "parameters": {}} + short = {"type": "function", "name": "Read", "description": "short", "parameters": {}} + data = {"model": "gpt-5.5", "instructions": "system", "tools": [tool], "input": items, **_caller()} + guardrail = _gateway_guardrail() + answer = _gateway_answer({"input": folded, "tools": [short], "instructions": "trimmed"}) + + with _posting(guardrail, return_value=answer) as post: + result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + payload = post.await_args.kwargs["json"] + assert payload["input"] == items and payload["instructions"] == "system" + assert result["input"] == folded + assert result["tools"] == [short] + assert result["instructions"] == "trimmed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answer", + [ + pytest.param(lambda: _raised_status(400, "Invalid config.mode: 'lossless'"), id="mode_rejected"), + pytest.param( + lambda: MagicMock(status_code=200, text="{}", **{"json.return_value": {"messages": [{"role": "user"}]}}), + id="legacy_shape", + ), + ], +) +async def test_older_headroom_falls_back_and_is_asked_again_later(answer): + """Upgrading LiteLLM before Headroom must not turn every request into a 502.""" + guardrail = _gateway_guardrail() + legacy = _make_compress_response([{"role": "user", "content": "earlier compressed"}]) + messages = [ + {"role": "user", "content": "earlier"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "now"}, + ] + + def turn() -> dict: + return {"model": "m", "messages": list(messages), **_caller()} + + with _posting(guardrail, side_effect=[answer(), legacy, legacy, _gateway_answer({"messages": messages})]) as post: + await _apply_turn(guardrail, turn()) + assert "gateway" not in post.await_args.kwargs["json"] + await _apply_turn(guardrail, turn()) + assert "gateway" not in post.await_args.kwargs["json"] + guardrail._gateway_turn_retry_at = 0.0 # the retry window has passed + await _apply_turn(guardrail, turn()) + + assert "gateway" in post.await_args.kwargs["json"] + + +@pytest.mark.asyncio +async def test_other_400s_fail_closed_and_keep_the_gateway_turn(): + guardrail = _gateway_guardrail() + data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()} + + with _posting(guardrail, side_effect=_raised_status(400, "config.session_id must be a string")): + with pytest.raises(HTTPException) as raised: + await _apply_turn(guardrail, data) + + assert raised.value.status_code == 502 + assert guardrail._gateway_turn_retry_at == 0.0 + + +@pytest.mark.asyncio +async def test_gateway_turn_retries_a_session_lock_timeout_once(): + guardrail = _gateway_guardrail() + folded = [{"role": "user", "content": "folded"}] + data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()} + + with _posting(guardrail, side_effect=[_raised_status(503, "busy"), _gateway_answer({"messages": folded})]) as post: + await _apply_turn(guardrail, data) + + assert post.await_count == 2 + assert data["messages"] == folded + + +@pytest.mark.asyncio +async def test_gateway_session_is_one_conversation_of_one_caller(): + guardrail = _gateway_guardrail() + first = {"role": "user", "content": "a"} + marked = {"role": "user", "content": [{"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}]} + image = {"type": "image", "source": {"type": "url", "url": "https://x.test/a.png"}} + + def turn(row: dict, *rest: dict, api_key: str = "one", session: str = "s", user: str | None = None) -> dict: + extra = {"user": user} if user else {} + return {"model": "m", "messages": [row, *rest], "session_id": session, **_caller(api_key), **extra} + + base = await _session_of(guardrail, turn(first)) + assert base is not None + assert base == await _session_of(guardrail, turn(first, {"role": "assistant", "content": "b"})) + assert base == await _session_of(guardrail, turn(marked)) + assert base != await _session_of(guardrail, turn(first, session="other")) + assert base != await _session_of(guardrail, turn(first, api_key="two")) + assert await _session_of(guardrail, turn(first, user="alice")) != await _session_of( + guardrail, turn(first, user="bob") + ) + # Claude Code's subagents share the parent's session id. + assert base != await _session_of(guardrail, turn({"role": "user", "content": "subagent"})) + with_image = {"role": "user", "content": [image]} + marked_image = {"role": "user", "content": [{**image, "cache_control": {"type": "ephemeral"}}]} + assert await _session_of(guardrail, turn(with_image)) == await _session_of(guardrail, turn(marked_image)) + + +@pytest.mark.asyncio +async def test_gateway_turn_without_a_caller_session_stays_stateless(): + """Many end users behind one key can open with the same row, so they never share a session.""" + guardrail = _gateway_guardrail() + data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **_caller()} + assert await _session_of(guardrail, data) is None + + +@pytest.mark.asyncio +async def test_auto_router_hop_gets_compressed_structured_messages(): + """The routing hop has no caller on it and reads only the returned messages.""" + guardrail = _gateway_guardrail() + legacy_rows = [{"role": "user", "content": "earlier compressed"}] + messages = [ + {"role": "user", "content": "earlier"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "now"}, + ] + + with _posting(guardrail, return_value=_make_compress_response(legacy_rows)) as post: + result = await _apply_turn(guardrail, {"messages": messages, "model": "router"}) + + assert "gateway" not in post.await_args.kwargs["json"] + assert result["structured_messages"][0]["content"] == "earlier compressed" + + +@pytest.mark.asyncio +async def test_embeddings_input_never_reaches_headroom(): + guardrail = _gateway_guardrail() + data = {"model": "m", "input": ["some text"], **_caller()} + inputs = GenericGuardrailAPIInputs(texts=["some text"]) + + with _posting(guardrail) as post: + assert await guardrail.apply_guardrail(inputs=inputs, request_data=data, input_type="request") is inputs + post.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_gateway_turn_forwards_and_merges_anthropic_betas(): + guardrail = _gateway_guardrail() + answer = _gateway_answer( + {"messages": [{"role": "user", "content": "x"}]}, headers={"anthropic-beta": "plugin-beta"} + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "x"}], + "extra_headers": {"anthropic-beta": "configured-beta", "x-other": "kept"}, + "proxy_server_request": {"headers": {"Anthropic-Beta": "client-beta"}}, + **_caller(), + } + + with _posting(guardrail, return_value=answer) as post: + await _apply_turn(guardrail, data) + + assert post.await_args.kwargs["json"]["gateway"]["request_headers"] == {"anthropic-beta": "client-beta"} + assert data["extra_headers"] == {"anthropic-beta": "configured-beta,plugin-beta", "x-other": "kept"} + + +@pytest.mark.asyncio +async def test_gateway_turn_fail_open_leaves_the_request_alone(): + data = {"model": "m", "messages": [{"role": "user", "content": "task"}], "tools": [{"name": "Read"}], **_caller()} + inputs = GenericGuardrailAPIInputs(structured_messages=data["messages"]) + guardrail = _gateway_guardrail(unreachable_fallback="fail_open") + + with _posting(guardrail, side_effect=httpx.ConnectError("down")) as post: + returned = await guardrail.apply_guardrail(inputs=inputs, request_data=data, input_type="request") + + assert post.await_args.kwargs["json"]["gateway"] + assert returned is inputs + assert data["messages"] == [{"role": "user", "content": "task"}] + assert data["tools"] == [{"name": "Read"}] + + +@pytest.mark.asyncio +async def test_gateway_turn_fail_closed_raises_when_headroom_is_down(): + data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()} + guardrail = _gateway_guardrail() + + with _posting(guardrail, side_effect=httpx.ConnectError("down")): + with pytest.raises(HTTPException) as raised: + await _apply_turn(guardrail, data) + assert raised.value.status_code == 502 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3aa7c888e98..3cb5df74fe2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37055,7 +37055,7 @@ export interface components { category_thresholds?: components["schemas"]["LakeraCategoryThresholds"] | null; /** * Ccr Retrieval - * @description Inject the Headroom retrieval tool for hashes declared by the compression service. + * @description Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false. * @default true */ ccr_retrieval: boolean; @@ -37080,6 +37080,12 @@ export interface components { * @description Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable. */ client_secret?: string | null; + /** + * Compression Mode + * @description Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; 'lossy_inline' saves more without a way back. Null uses the Headroom server's default. + * @default lossless + */ + compression_mode: ("lossless" | "lossy_inline" | "lossless_then_lossy") | null; /** * Confidence Threshold * @description Only block or mask when detection confidence >= this value; below threshold, allow or log_only. @@ -37237,6 +37243,12 @@ export interface components { * @description noma_v2 only: name of this gateway, used as the gateway_host label on Noma scans */ gateway_name?: string | null; + /** + * Gateway Turn + * @description Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail compresses a subset of rows instead. + * @default true + */ + gateway_turn: boolean; /** * Grounding Check * @description Enable grounding verification to ensure output is grounded in provided context.