From 0d5cf9b2f7d8cba315cd6c550bd98aa9a2fa88e1 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 12:32:03 -0700 Subject: [PATCH] fix(headroom): forward tenant-scoped session affinity (#49255) --- .../litellm_core_utils/provider_affinity.py | 42 ++++- .../guardrail_hooks/headroom/headroom.py | 91 ++++----- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/litellm_params.py | 2 + .../guardrails/guardrail_hooks/headroom.py | 12 +- .../test_provider_affinity.py | 16 +- .../guardrail_hooks/test_headroom.py | 174 +++++++++--------- tests/unit/types/test_litellm_params.py | 3 + 8 files changed, 209 insertions(+), 132 deletions(-) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index c831e92c09d..767eab962c8 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,12 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + SESSION_ID_FROM_TRACE_ID_PARAM_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -98,6 +103,41 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) +def get_caller_session_id(litellm_params: object | None) -> str | None: + """Resolve a caller conversation without accepting a generated session or a promoted trace id.""" + if litellm_params is None: + return None + metadata_values: Final = _metadata_values(litellm_params) + preserved_session: Final = next( + ( + value + for metadata in metadata_values + if isinstance(value := _get_value(metadata, CALLER_SESSION_ID_METADATA_KEY), str) and value + ), + None, + ) + if preserved_session: + return preserved_session + direct_session: Final = _get_value(litellm_params, "session_id") + if isinstance(direct_session, str) and direct_session: + return direct_session + + proxy_request: Final = _get_value(litellm_params, "proxy_server_request") + headers: Final = as_str_mapping(_get_value(proxy_request, "headers")) + trace_header: Final = next( + (value for key, value in (headers or {}).items() if key.lower() == "x-litellm-trace-id"), None + ) + candidates: Final = ( + get_caller_litellm_session_id(litellm_params), + *( + _get_value(metadata, "session_id") + for metadata in metadata_values + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) + ), + ) + return next((value for value in candidates if isinstance(value, str) and value and value != trace_header), None) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index dac802a1ddc..7f48766bf26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -1,10 +1,10 @@ from __future__ import annotations +import hashlib import json import math import re import time -import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, TypeGuard @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( group_tool_exchanges, has_tool_with_name, ) +from litellm.litellm_core_utils.provider_affinity import get_caller_session_id from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, @@ -39,6 +40,7 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( merge_rewritten_text_parts, ) from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER +from litellm.router_utils.common_utils import request_user_api_key_auth from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.integrations.custom_logger import ( @@ -47,6 +49,11 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, as_converted_stream, ) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -357,15 +364,18 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: } -def _resolve_call_id(logging_obj: object, request_state: dict[str, object]) -> str | None: - """Resolve the litellm_call_id shared by a request's pre-call hook and its - agentic-loop hooks, so CCR hash validation can be scoped per call instead - of trusting any hash-shaped string that shows up in message text.""" - logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) - if isinstance(logging_call_id, str) and logging_call_id: - return logging_call_id - kwargs_call_id: Final = request_state.get("litellm_call_id") - return kwargs_call_id if isinstance(kwargs_call_id, str) else None +def _session_affinity(request_data: Mapping[str, object]) -> str | None: + 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: + 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 None + identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":")) + return hashlib.sha256(identity.encode("utf-8")).hexdigest() def has_headroom_retrieve_tool(tools: object) -> bool: @@ -521,7 +531,6 @@ class HeadroomGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - self._issued_hashes_by_call_id: dict[str, tuple[frozenset[str], float]] = {} super().__init__( # pyright: ignore[reportUnknownMemberType] guardrail_name=guardrail_name, event_hook=event_hook, @@ -540,10 +549,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) -> 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 is not None: + headers["x-headroom-affinity"] = affinity return headers @staticmethod @@ -569,14 +580,6 @@ class HeadroomGuardrail(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) - def _prune_expired_hashes(self) -> None: - now: Final = time.monotonic() - self._issued_hashes_by_call_id = { - call_id: (hashes, expiry) - for call_id, (hashes, expiry) in self._issued_hashes_by_call_id.items() - if expiry > now - } - def _handle_compress_failure( self, messages: list[dict[str, object]], @@ -596,16 +599,19 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, + affinity: str | None, ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model + 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(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except httpx.HTTPStatusError as e: @@ -734,7 +740,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, affinity: str | None, query: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: params["query"] = query @@ -743,7 +749,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: @@ -777,13 +783,14 @@ class HeadroomGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: if input_type != "request": return inputs + request_data.pop(HEADROOM_REQUEST_CONTEXT_KEY, None) if self._should_bypass(request_data): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs @@ -809,9 +816,7 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - # request_data is a bare dict on the base signature; narrow it before - # reading the untranslated messages so long tool names can be recovered. - raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_messages: Final = request_data.get("messages") raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] @@ -820,9 +825,11 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() + affinity: Final = _session_affinity(request_data) result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, + affinity=affinity, ) end_time: Final = time.time() @@ -868,12 +875,9 @@ class HeadroomGuardrail(CustomGuardrail): if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] - self._prune_expired_hashes() - call_id = _resolve_call_id(logging_obj, request_data) - if not call_id: - call_id = str(uuid.uuid4()) - request_data["litellm_call_id"] = call_id - self._issued_hashes_by_call_id[call_id] = (frozenset(hashes), time.monotonic() + _HASH_CACHE_TTL_SECONDS) + request_data[HEADROOM_REQUEST_CONTEXT_KEY] = HeadroomRequestContext( + frozenset(hashes), affinity, time.monotonic() + _HASH_CACHE_TTL_SECONDS + ) existing_tools: Final = inputs.get("tools") retrieve_tool: Final = _build_headroom_retrieve_tool() @@ -934,9 +938,11 @@ class HeadroomGuardrail(CustomGuardrail): ) -> AgenticLoopPlan: tool_calls: Final[list[dict[str, object]]] = tools.get("tool_calls", []) - self._prune_expired_hashes() - call_id: Final = _resolve_call_id(logging_obj, kwargs) - valid_hashes = self._issued_hashes_by_call_id.get(call_id, (frozenset(), 0.0))[0] if call_id else frozenset() + context: Final = _REQUEST_DATA_ADAPTER.validate_python(kwargs).get(HEADROOM_REQUEST_CONTEXT_KEY) + valid_context: Final = ( + context if isinstance(context, HeadroomRequestContext) and context.expires_at > time.monotonic() else None + ) + valid_hashes: Final = valid_context.issued_hashes if valid_context is not None else frozenset() retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: @@ -944,11 +950,6 @@ class HeadroomGuardrail(CustomGuardrail): raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None - # A hash is only honored if it was issued by *this request's own* - # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by - # message text alone is forgeable -- an attacker can plant a - # hash-shaped string in their own prompt, and a hash issued for one - # request would validate for any other request that echoes it back. if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", @@ -958,6 +959,7 @@ class HeadroomGuardrail(CustomGuardrail): else: content = await self._call_retrieve( hash_value=hash_value, + affinity=valid_context.affinity if valid_context is not None else None, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) @@ -996,7 +998,12 @@ class HeadroomGuardrail(CustomGuardrail): max_tokens=max_tokens, optional_params=optional_params_without_max_tokens, kwargs={ - k: v for k, v in kwargs.items() if not k.startswith("_headroom") and k != "litellm_logging_obj" + **{ + k: v + for k, v in _REQUEST_DATA_ADAPTER.validate_python(kwargs).items() + if not k.startswith("_headroom") and k != "litellm_logging_obj" + }, + **({HEADROOM_REQUEST_CONTEXT_KEY: valid_context} if valid_context is not None else {}), }, ), metadata={"tool_type": "headroom_ccr"}, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9b52f2f8a8d..ae9a6d4b959 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2566,6 +2566,7 @@ def refresh_proxy_server_request_body_snapshot( VECTOR_STORE_CALLER_METADATA_KWARG, STREAM_USAGE_INJECTED_FIELD, "_litellm_strip_stream_usage", + "_headroom_request_context", } ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 2065093f295..3c2c636b5ed 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets from litellm.types.caching import DynamicCacheControl from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage + from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ConfigurableClientsideParamsCustomAuth, DeploymentTypedDict, RetryPolicy from litellm.types.router_weights import RouterWeights @@ -352,6 +353,7 @@ class CallState: @dataclass(frozen=True, slots=True, kw_only=True) class AgenticLoopState: + headroom_context: "HeadroomRequestContext | None" = field(default=None, metadata=wire("_headroom_request_context")) depth: int | None = field(default=None, metadata=wire("_agentic_loop_depth")) fingerprints: Sequence[str] | None = field(default=None, metadata=wire("_agentic_loop_fingerprints")) api_surface: Literal["chat_completions", "responses"] | None = field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..76e806e3899 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -1,9 +1,19 @@ -from typing import Literal +from dataclasses import dataclass +from typing import Final, Literal from pydantic import BaseModel, Field from .base import GuardrailConfigModel +HEADROOM_REQUEST_CONTEXT_KEY: Final = "_headroom_request_context" + + +@dataclass(frozen=True, slots=True) +class HeadroomRequestContext: + issued_hashes: frozenset[str] + affinity: str | None + expires_at: float + class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): api_base: str | None = Field( diff --git a/tests/unit/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py index edf4f5169b5..6ccd6b17c01 100644 --- a/tests/unit/litellm_core_utils/test_provider_affinity.py +++ b/tests/unit/litellm_core_utils/test_provider_affinity.py @@ -1,12 +1,16 @@ +from collections.abc import Callable + import pytest from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.provider_affinity import ( add_provider_affinity_header, + get_caller_session_id, get_stable_session_id, ) +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize( ("litellm_params", "expected"), [ @@ -16,18 +20,22 @@ from litellm.litellm_core_utils.provider_affinity import ( ({"litellm_metadata": {"session_id": "litellm-metadata-session"}}, "litellm-metadata-session"), ], ) -def test_get_stable_session_id_uses_explicit_session_sources(litellm_params: dict, expected: str): - assert get_stable_session_id(litellm_params) == expected +def test_session_ids(litellm_params: dict[str, object], expected: str, resolve: Callable[[object], str | None]) -> None: + assert resolve(litellm_params) == expected def test_get_stable_session_id_does_not_use_trace_id(): assert get_stable_session_id({"litellm_trace_id": "per-request-trace"}) is None + assert get_caller_session_id( + {"litellm_session_id": "trace", "proxy_server_request": {"headers": {"x-litellm-trace-id": "trace"}}} + ) is None +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_get_stable_session_id_ignores_proxy_generated_session(metadata_key: str): +def test_generated_sessions_are_ignored(metadata_key: str, resolve: Callable[[object], str | None]) -> None: assert ( - get_stable_session_id( + resolve( { "litellm_session_id": "generated-session", metadata_key: { diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 741fea5f2b9..0cdcb2e2d07 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -23,7 +23,6 @@ Tests cover: import datetime import json -import time from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -33,17 +32,22 @@ import respx from fastapi import HTTPException import litellm - from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( - HeadroomGuardrail, - has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + _session_affinity, + has_headroom_retrieve_tool, ) from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import ( CallTypes, GenericGuardrailAPIInputs, @@ -148,6 +152,10 @@ def guardrail() -> HeadroomGuardrail: return _make_guardrail() +def _issued_context(*hashes: str, affinity: str | None = None) -> HeadroomRequestContext: + return HeadroomRequestContext(frozenset(hashes), affinity, float("inf")) + + def _recorded_guardrail_entries(request_data: dict) -> list: for container_key in ("metadata", "litellm_metadata"): container = request_data.get(container_key) @@ -175,20 +183,22 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( structured_messages=ORIGINAL_MESSAGES, ) mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data = {"model": "gpt-4o"} + auth: Final = UserAPIKeyAuth(api_key="tenant") + request_data: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": auth}} with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response, - ): + ) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) + assert mock_post.call_args.kwargs["headers"]["x-headroom-affinity"] == _session_affinity(request_data) assert result.get("structured_messages") == EXPECTED_MESSAGES entries = _recorded_guardrail_entries(request_data) @@ -510,10 +520,6 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -530,11 +536,12 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac", affinity="original")}, ) assert plan.run_agentic_loop is True assert plan.request_patch is not None + assert HEADROOM_REQUEST_CONTEXT_KEY in plan.request_patch.kwargs follow_up = plan.request_patch.messages assert follow_up is not None @@ -545,6 +552,7 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( assert tool_result_message["tool_call_id"] == "call_abc123" mock_get.assert_called_once() + assert mock_get.call_args.kwargs["headers"]["x-headroom-affinity"] == "original" call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] assert "b573993006976af767214fac" in call_url @@ -576,10 +584,6 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( "content": "Retrieve more: hash=deadbeef000000000000dead", } ] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"deadbeef000000000000dead"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -596,7 +600,7 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -606,9 +610,13 @@ 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), +]) async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, -): + 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 in the current request's own messages (e.g. via prompt injection).""" @@ -625,8 +633,6 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( arguments={"hash": "deadbeef000000000000dead"}, tool_id="call_xyz", ) - assert not guardrail._issued_hashes_by_call_id - with patch.object( guardrail.async_handler, "get", @@ -641,7 +647,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-unknown"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: context}, ) mock_get.assert_not_called() @@ -653,17 +659,11 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( @pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( +async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_request( guardrail: HeadroomGuardrail, ): - """A hash issued for one request must not be retrievable by a different - request just because the second request echoes that hash-shaped string - back in its own messages -- retrieval must be scoped per litellm_call_id, - not derived by re-scanning attacker-controlled message text.""" - guardrail._issued_hashes_by_call_id["call-A"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) + issuer_data: Final[dict[str, object]] = {"model": "gpt-4o", "litellm_call_id": "call-A"} + await _apply(guardrail, ORIGINAL_MESSAGES, ["b573993006976af767214fac"], request_data=issuer_data) tool_calls = [ { @@ -693,7 +693,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_c anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-B"}, + kwargs={"litellm_call_id": "call-A"}, ) mock_get.assert_not_called() @@ -736,10 +736,6 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -756,7 +752,7 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -806,7 +802,6 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste "arguments": json.dumps({"hash": retrieve_hash}), } ] - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({retrieve_hash}), time.monotonic() + 999) with patch.object( guardrail.async_handler, @@ -832,7 +827,7 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(retrieve_hash)}, ) assert plan.request_patch is not None @@ -876,10 +871,6 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -896,7 +887,7 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -929,16 +920,20 @@ HASH_SHAPED_HISTORY = [ ] -async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: - request_data = {"model": "claude-sonnet-5"} +async def _apply( + 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"} def _echo(**kwargs): + assert kwargs["json"].get("config") == ({"mode": "ccr"} if guardrail.ccr_retrieval else None) return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): return await guardrail.apply_guardrail( inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), - request_data=request_data, + request_data=data, input_type="request", ) @@ -950,10 +945,11 @@ async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail tool must not appear and no hash may be registered as issued. Covers a service that omits ccr_hashes, returns it empty, or returns a non-list.""" for ccr_hashes in (None, [], "b573993006976af767214fac"): - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + request_data: Final[dict[str, object]] = {} + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes, request_data=request_data) assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -961,15 +957,18 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr """Only the hashes the service reports in ccr_hashes are honored, in the service's own 12 to 24 hex grammar; anything else is dropped because each entry is interpolated into the /v1/retrieve URL.""" + request_data: Final[dict[str, object]] = {} result = await _apply( guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + request_data=request_data, ) assert has_headroom_retrieve_tool(result.get("tools") or []) - ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() - assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + context: Final = request_data[HEADROOM_REQUEST_CONTEXT_KEY] + assert isinstance(context, HeadroomRequestContext) + assert context.issued_hashes == frozenset({"98ca69107318", "b573993006976af767214fac"}) @pytest.mark.asyncio @@ -987,7 +986,6 @@ async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatc assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id @pytest.mark.asyncio @@ -1498,6 +1496,7 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ] guardrail = _make_guardrail(unreachable_fallback="fail_open") + request_data: Final = {HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")} inputs = GenericGuardrailAPIInputs( texts=["hello"], structured_messages=messages_with_fake_hash, @@ -1511,13 +1510,13 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ): result = await guardrail.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data, input_type="request", ) assert result["structured_messages"] == messages_with_fake_hash assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -2164,18 +2163,18 @@ def _openai_completion_payload(message: dict, finish_reason: str) -> dict: } -def _openai_tool_call_payload() -> dict: +def _openai_tool_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_completion_payload( { "role": "assistant", "content": None, "tool_calls": [ { - "id": "call_ccr", + "id": f"call_{hash_value}", "type": "function", "function": { "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), }, } ], @@ -2215,14 +2214,14 @@ def _openai_responses_payload(output_item: dict) -> dict: } -def _openai_responses_retrieve_call_payload() -> dict: +def _openai_responses_retrieve_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_responses_payload( { "type": "function_call", "id": "fc_ccr", - "call_id": "call_ccr", + "call_id": f"call_{hash_value}", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), "status": "completed", } ) @@ -2381,10 +2380,6 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( to a client that never declared the tool.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2392,6 +2387,7 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_tool_call_payload("a" * 24)), httpx.Response(200, json=_openai_text_payload(final_answer)), ] ) @@ -2408,20 +2404,22 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( tools=[_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["messages"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) @pytest.mark.asyncio @@ -2435,10 +2433,6 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( instead of streaming a headroom_retrieve function_call to the client.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2446,6 +2440,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( side_effect=[ httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_retrieve_call_payload("a" * 24)), httpx.Response(200, json=_openai_responses_text_payload(final_answer)), ] ) @@ -2462,6 +2457,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) events = [event async for event in response] @@ -2471,14 +2467,15 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( assert streamed_text == final_answer assert not any("function_call" in str(getattr(event, "type", "")) for event in events) assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["input"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( @@ -2491,10 +2488,6 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( completed response object.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -2518,6 +2511,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH)}, ) events = list(response) @@ -2917,7 +2911,6 @@ def _anthropic_response_with_text_and_tool_call() -> dict: async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) logging_obj = MagicMock() logging_obj.litellm_call_id = "call-1" logging_obj.model_call_details = {} @@ -2938,7 +2931,7 @@ async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("h" * 24)}, ) @@ -3033,7 +3026,7 @@ async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail new_callable=AsyncMock, return_value=_make_retrieve_response("original"), ) as mock_get: - result = await guardrail._call_retrieve("a" * 24) + result = await guardrail._call_retrieve("a" * 24, affinity=None) assert result == "original" timeout = _timeout_of(mock_get.call_args) @@ -3143,3 +3136,16 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): assert guardrail.timeout.read == 60.0 assert guardrail.timeout.connect == 5.0 + + +@pytest.mark.parametrize("session", [None, "conversation", "other"]) +@pytest.mark.parametrize("scope", [None, {}, {"api_key": "key"}, {"api_key": "other"}, {"user_id": "user"}]) +def test_affinity_requires_trusted_identity_and_caller_session( + scope: dict[str, str] | None, session: str | None +) -> None: + baseline: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="key")}} + auth: Final = UserAPIKeyAuth(**scope) if scope else scope + candidate: Final = {"session_id": session, "metadata": {"user_api_key_auth": auth}} + 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) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e369b7e93ab..22eb9929e7f 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -33,6 +33,7 @@ from litellm.types.litellm_params import ( wire_names, ) from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ( ConfigurableClientsideParamsCustomAuth, @@ -231,6 +232,7 @@ OPTION_NAMES: Final = ( ) AGENTIC_LOOP_STATE_NAMES: Final = ( + "_headroom_request_context", "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", @@ -529,6 +531,7 @@ TYPE_HINT_NAMESPACE: Final[Mapping[str, object]] = { "Logging": Logging, "SecretFields": SecretFields, "CompactionState": CompactionState, + "HeadroomRequestContext": HeadroomRequestContext, "RouterWeights": RouterWeights, "AttemptedFallbackTargets": AttemptedFallbackTargets, }