diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7198d74909c..eaa125a176d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15790,6 +15790,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 dac802a1ddc..3cf6193a3e7 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, @@ -56,6 +60,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) ) @@ -87,6 +105,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. @@ -505,6 +593,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: @@ -518,6 +607,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, ) @@ -540,10 +634,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 @@ -596,6 +692,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: @@ -605,7 +702,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: @@ -734,7 +831,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 @@ -743,7 +840,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: @@ -823,6 +920,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() @@ -939,6 +1037,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 "" @@ -959,6 +1062,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 5bdcdddf4b7..c1c8c63d715 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -19,6 +19,9 @@ Tests cover: - 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 @@ -34,6 +37,7 @@ 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, @@ -3075,3 +3079,242 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): 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 f56fa87e93b..83a22766b0e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36148,6 +36148,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'.