diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5cb63c98d56..d0d011b25f2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18439,7 +18439,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" }, @@ -18492,6 +18492,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, @@ -18833,6 +18851,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 2272b4bea43..235b4d0d201 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37834,7 +37834,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; @@ -37859,6 +37859,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. @@ -38016,6 +38022,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.