scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

fix(headroom): forward tenant-scoped session affinity (#49255)

Browse source
This commit is contained in:
tin-berri 2026-10-30 12:32:03 -07:00 • committed by GitHub
parent b2466f8d7b
commit 0d5cf9b2f7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 209 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/litellm_core_utils/provider_affinity.py
View file

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

91
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -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"},

1
litellm/proxy/litellm_pre_call_utils.py
View file

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

2
litellm/types/litellm_params.py
View file

@ -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(

12
litellm/types/proxy/guardrails/guardrail_hooks/headroom.py
View file

@ -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(

16
tests/unit/litellm_core_utils/test_provider_affinity.py
View file

@ -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: {

174
tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py
View file

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

3
tests/unit/types/test_litellm_params.py
View file

@ -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,
}