Merge 683f8cce88 into 4bafebbcdc

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": [
                 {

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))

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:

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

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'.