scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(headroom): send the whole request on Headroom's gateway turn

Browse source 
The guardrail sent Headroom only a subset of chat rows, flattened to strings,
and withheld the newest results. Claude Code traffic got no compression, other
clients paid more than with no guardrail because history was rewritten and the
prompt cache busted, and Read results were compressed lossily.

Now the guardrail sends the provider request (messages or input, system, tools)
with Headroom's gateway block and config.mode "lossless", and writes back what
Headroom returns. History stays byte-stable, nothing is lost, and server-side
plugins can rewrite tools and skills. Requests with no authenticated caller and
older Headroom servers keep the row-subset path, and an older server is asked
again every 5 minutes.
This commit is contained in:
Tejas Chopra 2026-11-03 08:52:03 +05:30
parent 704e73d6fe
commit e120419dc5
6 changed files with 690 additions and 82 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -17877,7 +17877,7 @@
},
"ccr_retrieval": {
"default": true,
"description": "Inject the Headroom retrieval tool for hashes declared by the compression service.",
"description": "Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false.",
"title": "Ccr Retrieval",
"type": "boolean"
},
@ -17930,6 +17930,24 @@
"description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.",
"title": "Client Secret"
},
"compression_mode": {
"anyOf": [
{
"enum": [
"lossless",
"lossy_inline",
"lossless_then_lossy"
],
"type": "string"
},
{
"type": "null"
}
],
"default": "lossless",
"description": "Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; 'lossy_inline' saves more without a way back. Null uses the Headroom server's default.",
"title": "Compression Mode"
},
"confidence_threshold": {
"default": 0.5,
"default_value": 0.5,
@ -18271,6 +18289,12 @@
"description": "noma_v2 only: name of this gateway, used as the gateway_host label on Noma scans",
"title": "Gateway Name"
},
"gateway_turn": {
"default": true,
"description": "Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail compresses a subset of rows instead.",
"title": "Gateway Turn",
"type": "boolean"
},
"grounding_check": {
"anyOf": [
{

2
litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py
View file

@ -37,6 +37,8 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) ->
unreachable_fallback=litellm_params.unreachable_fallback,
timeout=litellm_params.timeout,
ccr_retrieval=litellm_params.ccr_retrieval,
gateway_turn=litellm_params.gateway_turn,
compression_mode=litellm_params.compression_mode,
)
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType]
_callback

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

@ -54,7 +54,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.headroom import (
HEADROOM_REQUEST_CONTEXT_KEY,
HeadroomRequestContext,
)
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, GuardrailStatus
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -72,6 +72,21 @@ _COMPRESS_TIMEOUT_SECONDS: Final = 60.0
HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve"
_HASH_PATTERN: Final = re.compile(r"[a-f0-9]{12,24}")
_HASH_CACHE_TTL_SECONDS: Final = 15 * 60
# The provider request Headroom's gateway turn reads, and the part it may rewrite.
# Everything else in request_data is LiteLLM's own or the client's to keep.
_GATEWAY_REWRITABLE_KEYS: Final = ("messages", "input", "system", "instructions", "tools", "tool_choice")
_GATEWAY_BODY_KEYS: Final = (*_GATEWAY_REWRITABLE_KEYS, "model", "max_tokens", "max_output_tokens", "stream")
# How long a Headroom without the gateway turn is left alone before it's asked again,
# so a rolling upgrade or one old replica doesn't switch the feature off for good.
_GATEWAY_UNSUPPORTED_RETRY_SECONDS: Final = 300.0
_COMPRESS_STAT_KEYS: Final = (
"tokens_before",
"tokens_after",
"tokens_saved",
"tool_tokens_saved",
"compression_ratio",
"transforms_applied",
)
# Narrows the base class's bare-dict ``request_data`` at the boundary so its
# untranslated messages can be read with concrete types (values pass through by
# reference, so this is a shallow top-level reconstruction).
@ -337,6 +352,37 @@ class _CompressResult:
ccr_hashes: frozenset[str] = frozenset()
@dataclass(frozen=True, slots=True)
class _TurnAnswered:
answer: Mapping[str, object]
body: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class _TurnUnsupported:
status_code: int
@dataclass(frozen=True, slots=True)
class _TurnFailed:
"""Headroom failed, and fail_open lets the request through as it was."""
def _is_count(value: object) -> TypeGuard[int | float]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
return isinstance(value, (int, float)) and not isinstance(value, bool)
def _compression_stats(body: Mapping[str, object]) -> Mapping[str, object]:
"""The savings Headroom reported. Spend tracking (extract_compression_saved_tokens)
reads only tokens_saved, which the service can omit, so it's derived then."""
stats: Final = {key: body[key] for key in _COMPRESS_STAT_KEYS if key in body}
before: Final = stats.get("tokens_before")
after: Final = stats.get("tokens_after")
if "tokens_saved" in stats or not _is_count(before) or not _is_count(after):
return stats
return {**stats, "tokens_saved": before - after}
def _build_headroom_retrieve_tool() -> dict[str, object]:
return {
"type": "function",
@ -364,20 +410,110 @@ def _build_headroom_retrieve_tool() -> dict[str, object]:
}
def _session_affinity(request_data: Mapping[str, object]) -> str | None:
def _request_tenant(request_data: Mapping[str, object]) -> tuple[object, ...] | None:
"""Who is calling, or None for a request with no caller on it."""
auth: Final = request_user_api_key_auth(request_data)
session_id: Final = get_caller_session_id(request_data)
if not isinstance(auth, UserAPIKeyAuth) or not session_id:
if not isinstance(auth, UserAPIKeyAuth):
return None
tenant: Final = (
("key", auth.api_key) if auth.api_key else ("principal", auth.org_id, auth.team_id, auth.user_id, auth.agent_id)
)
if not any(tenant[1:]):
return tenant if any(tenant[1:]) else None
def _session_affinity(request_data: Mapping[str, object]) -> str | None:
session_id: Final = get_caller_session_id(request_data)
tenant: Final = _request_tenant(request_data)
if not session_id or tenant is None:
return None
identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":"))
return hashlib.sha256(identity.encode("utf-8")).hexdigest()
def _without_cache_control(part: object) -> object:
return {key: value for key, value in part.items() if key != "cache_control"} if _is_str_object_dict(part) else part
def _conversation_identity(row: Mapping[str, object]) -> object:
"""A row as a conversation key. Clients move cache_control between turns, and
with it can resend the same text as a part list instead of a string."""
content: Final = row.get("content")
if isinstance(content, str) or is_all_text_parts(content):
return content_to_text(content)
return [_without_cache_control(part) for part in content] if _is_object_list(content) else content
def _gateway_session_id(
request_data: Mapping[str, object], rows: Sequence[object], model: object, tenant: tuple[object, ...]
) -> str | None:
"""One Headroom session per conversation, or None to stay stateless.
Headroom replays what it already sent in a session, which keeps the provider's
prompt cache warm, and turns in one session queue on one lock. So only the
caller's own session id joins requests into a session: one app key can serve
many end users whose conversations all start the same way. Claude Code's
subagents share their parent's session id, and the first user row tells them apart.
"""
caller_session: Final = get_caller_session_id(request_data)
if not caller_session:
return None
first_user_row: Final = next((row for row in rows if _is_str_object_dict(row) and row.get("role") == "user"), None)
identity: Final = json.dumps(
[
"headroom-session-v2",
tenant,
request_data.get("user"),
caller_session,
model,
_conversation_identity(first_user_row) if first_user_row is not None else None,
],
sort_keys=True,
default=str,
)
return f"litellm-{hashlib.sha256(identity.encode('utf-8')).hexdigest()[:32]}"
def _client_anthropic_beta(request_data: Mapping[str, object]) -> str | None:
proxy_request: Final = request_data.get("proxy_server_request")
headers: Final = proxy_request.get("headers") if _is_str_object_dict(proxy_request) else None
if not _is_str_object_dict(headers):
return None
return next(
(value for name, value in headers.items() if name.lower() == "anthropic-beta" and isinstance(value, str)),
None,
)
def _is_gateway_turn_unsupported(response: HttpxResponse, body: object) -> bool:
"""An older Headroom: it rejects the mode, or answers in the legacy messages-only shape."""
if response.status_code == 400:
return "config.mode" in response.text
return _is_str_object_dict(body) and "messages" in body and "body" not in body
def _json_or_none(response: HttpxResponse) -> object:
"""A 200 answer's JSON body, or None for any other status or a body that isn't JSON."""
if response.status_code != 200:
return None
try:
return _json_body(response)
except ValueError:
return None
def _merged_extra_headers(existing: object, added: object) -> Mapping[str, object] | None:
"""The request's extra headers plus the ones Headroom asks for, anthropic-beta values combined."""
if not _is_str_object_dict(added) or not added:
return None
current: Final = existing if _is_str_object_dict(existing) else {}
if "anthropic-beta" not in current or "anthropic-beta" not in added:
return {**current, **added}
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
beta: Final = merge_anthropic_beta_headers(str(current["anthropic-beta"]), str(added["anthropic-beta"]))
return {**current, **added, "anthropic-beta": beta}
def has_headroom_retrieve_tool(tools: object) -> bool:
return has_tool_with_name(tools, HEADROOM_RETRIEVE_TOOL_NAME)
@ -515,6 +651,8 @@ class HeadroomGuardrail(CustomGuardrail):
unreachable_fallback: str | None = None,
timeout: float | None = None,
ccr_retrieval: bool = True,
gateway_turn: bool = True,
compression_mode: Literal["lossless", "lossy_inline", "lossless_then_lossy"] | None = "lossless",
):
self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/")
if not self.headroom_api_base:
@ -527,6 +665,9 @@ class HeadroomGuardrail(CustomGuardrail):
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
"fail_open" if unreachable_fallback == "fail_open" else "fail_closed"
)
self.gateway_turn = gateway_turn
self._gateway_turn_retry_at = 0.0
self.compression_mode = compression_mode
self.ccr_retrieval = ccr_retrieval
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
@ -607,34 +748,17 @@ class HeadroomGuardrail(CustomGuardrail):
if self.ccr_retrieval:
payload["config"] = {"mode": "ccr"}
try:
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
url=f"{self.headroom_api_base}/v1/compress",
json=payload,
headers=self._request_headers(affinity),
timeout=self.timeout,
)
except httpx.HTTPStatusError as e:
return _CompressResult(
self._handle_compress_failure(
messages,
"Headroom compression service returned an error",
_build_compress_failure_detail(e.response.status_code, e.response.text),
),
False,
{},
)
except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e:
response: Final = await self._post_compress(payload, affinity)
if isinstance(response, str):
return _CompressResult(
self._handle_compress_failure(
messages,
"Headroom compression service unreachable",
{"detail": str(e)},
{"detail": response},
),
False,
{},
)
response: Final[HttpxResponse] = raw_response
if response.status_code != 200:
return _CompressResult(
@ -714,32 +838,146 @@ class HeadroomGuardrail(CustomGuardrail):
body.get("compression_ratio", 0),
)
stats: Final = {
key: body[key]
for key in (
"tokens_before",
"tokens_after",
"tokens_saved",
"compression_ratio",
"transforms_applied",
)
if key in body
}
tokens_before: Final = stats.get("tokens_before")
tokens_after: Final = stats.get("tokens_after")
if (
"tokens_saved" not in stats
and isinstance(tokens_before, (int, float))
and not isinstance(tokens_before, bool)
and isinstance(tokens_after, (int, float))
and not isinstance(tokens_after, bool)
):
# Spend tracking (extract_compression_saved_tokens) reads only
# tokens_saved, which the live compression service omits; derive it
# so savings are counted, but let a service-sent value win.
stats["tokens_saved"] = tokens_before - tokens_after
stats: Final = dict(_compression_stats(body))
return _CompressResult(filtered, True, stats, _read_ccr_hashes(body))
async def _post_compress(self, payload: Mapping[str, object], affinity: str | None) -> HttpxResponse | str:
"""Headroom's HTTP answer, or why it could not be reached."""
try:
return await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
url=f"{self.headroom_api_base}/v1/compress",
json=dict(payload),
headers=self._request_headers(affinity),
timeout=self.timeout,
)
except httpx.HTTPStatusError as e:
# AsyncHTTPHandler raises on any non-2xx, and the status decides what happens next.
return e.response
except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e:
return str(e)
async def _call_gateway_turn(
self, payload: Mapping[str, object], affinity: str | None
) -> _TurnAnswered | _TurnUnsupported | _TurnFailed:
"""A failed call goes through the configured fallback, so fail_closed raises here."""
first: Final = await self._post_compress(payload, affinity)
# A 503 means a session's turn lock timed out, and Headroom asks the gateway to retry.
retry: Final = not isinstance(first, str) and first.status_code == 503
sent: Final = await self._post_compress(payload, affinity) if retry else first
if isinstance(sent, str):
_ = self._handle_compress_failure([], "Headroom compression service unreachable", {"detail": sent})
return _TurnFailed()
answer: Final = _json_or_none(sent)
if _is_gateway_turn_unsupported(sent, answer):
return _TurnUnsupported(sent.status_code)
body: Final = answer.get("body") if _is_str_object_dict(answer) else None
if not _is_str_object_dict(answer) or not _is_str_object_dict(body):
_ = self._handle_compress_failure(
[],
"Headroom compression service returned an error",
_build_compress_failure_detail(sent.status_code, sent.text),
)
return _TurnFailed()
return _TurnAnswered(answer, body)
async def _apply_gateway_turn(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object], # mutable-ok: the rewritten request is written back into request_data
) -> GenericGuardrailAPIInputs | None:
"""Send Headroom the whole provider request and take back what it rewrites.
Reads and writes request_data, not the translated inputs: on the Anthropic
route those are OpenAI-shaped copies, while Headroom's folds and hooks need
the wire shape. The inputs handed back carry no texts or tools and the same
structured_messages object, so the translation handlers write nothing over
the rewritten request.
None means this request takes the row-subset path: it isn't a caller's live
request (the auto-router's routing hop passes a throwaway request_data and
reads the result's structured_messages), or Headroom is too old.
"""
tenant: Final = _request_tenant(request_data)
rows: Final = request_data["messages"] if "messages" in request_data else request_data.get("input")
if tenant is None or "structured_messages" not in inputs or not _is_object_list(rows) or not rows:
return None
model: Final = self.headroom_model or request_data.get("model")
session_id: Final = _gateway_session_id(request_data, rows, model, tenant)
config: Final = {
**({"session_id": session_id} if session_id else {}),
**({"mode": self.compression_mode} if self.compression_mode else {}),
}
client_beta: Final = _client_anthropic_beta(request_data)
payload: Final = {
**{key: request_data[key] for key in _GATEWAY_BODY_KEYS if key in request_data},
**({"model": model} if model else {}),
"config": config,
"gateway": {
"can_redrive": False,
"can_relay_response": False,
"request_headers": {"anthropic-beta": client_beta} if client_beta else {},
},
}
affinity: Final = _session_affinity(request_data) or session_id
start_time: Final = time.time()
outcome: Final = await self._call_gateway_turn(payload, affinity)
end_time: Final = time.time()
match outcome:
case _TurnUnsupported(status_code=status_code):
self._gateway_turn_retry_at = time.monotonic() + _GATEWAY_UNSUPPORTED_RETRY_SECONDS
verbose_proxy_logger.warning(
"Headroom at %s has no lossless gateway turn (HTTP %s); using row-subset compression and "
"asking again in %d seconds. Upgrade headroom-ai to get it.",
self.headroom_api_base,
status_code,
_GATEWAY_UNSUPPORTED_RETRY_SECONDS,
)
return None
case _TurnFailed():
self._record_outcome(
request_data,
"guardrail_failed_to_respond",
{"error": "headroom compression unavailable; request forwarded uncompressed"},
start_time,
end_time,
)
return inputs
case _TurnAnswered(answer=answer, body=body):
# A key Headroom sent back without was dropped on purpose (a tool_choice naming
# a tool it deferred, say), so it goes from the request too.
dropped: Final = tuple(key for key in _GATEWAY_REWRITABLE_KEYS if key in payload and key not in body)
for key in dropped:
_ = request_data.pop(key, None)
rewrites: Final = {key: body[key] for key in _GATEWAY_REWRITABLE_KEYS if key in body}
extra_headers: Final = _merged_extra_headers(request_data.get("extra_headers"), answer.get("headers"))
request_data.update({**rewrites, **({"extra_headers": extra_headers} if extra_headers else {})})
self._record_outcome(request_data, "success", _compression_stats(answer), start_time, end_time)
return GenericGuardrailAPIInputs(structured_messages=inputs["structured_messages"])
def _record_outcome(
self,
request_data: dict[str, object], # mutable-ok: LiteLLM's guardrail logging writes into request_data
status: GuardrailStatus,
response: Mapping[str, object],
start_time: float,
end_time: float,
) -> None:
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=response,
request_data=request_data,
guardrail_status=status,
guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER,
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
async def _call_retrieve(self, hash_value: str, affinity: str | None, query: str | None = None) -> str:
params: Final[dict[str, str]] = {}
if query:
@ -799,6 +1037,11 @@ class HeadroomGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Headroom: background request; skipping compression")
return inputs
if self.gateway_turn and time.monotonic() >= self._gateway_turn_retry_at:
turned: Final = await self._apply_gateway_turn(inputs, request_data)
if turned is not None:
return turned
structured_messages: Final = inputs.get("structured_messages")
if not _is_object_list(structured_messages) or not structured_messages:
return inputs
@ -833,21 +1076,14 @@ class HeadroomGuardrail(CustomGuardrail):
)
end_time: Final = time.time()
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
if not result.succeeded:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"},
request_data=request_data,
guardrail_status="guardrail_failed_to_respond",
guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER,
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
self._record_outcome(
request_data,
"guardrail_failed_to_respond",
{"error": "headroom compression unavailable; request forwarded uncompressed"},
start_time,
end_time,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
# Hand back the caller's own inputs object. Translation handlers
# detect "the guardrail rewrote the messages" by identity, so
# returning a rebuilt copy sends an unchanged request through the
@ -860,16 +1096,7 @@ class HeadroomGuardrail(CustomGuardrail):
protected_indices=protected_indices,
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=result.stats,
request_data=request_data,
guardrail_status="success",
guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER,
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
self._record_outcome(request_data, "success", result.stats, start_time, end_time)
hashes: Final = result.ccr_hashes if self.ccr_retrieval else frozenset()
if not hashes:

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

@ -38,7 +38,23 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]):
)
ccr_retrieval: bool = Field(
default=True,
description="Inject the Headroom retrieval tool for hashes declared by the compression service.",
description="Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false.",
)
gateway_turn: bool = Field(
default=True,
description=(
"Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway "
"contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools "
"and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail "
"compresses a subset of rows instead."
),
)
compression_mode: Literal["lossless", "lossy_inline", "lossless_then_lossy"] | None = Field(
default="lossless",
description=(
"Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; "
"'lossy_inline' saves more without a way back. Null uses the Headroom server's default."
),
)
@staticmethod

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

@ -33,6 +33,9 @@ from fastapi import HTTPException
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
HEADROOM_RETRIEVE_TOOL_NAME,
HeadroomGuardrail,
@ -95,7 +98,7 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail:
default_on=True,
)
defaults.update(kwargs)
return HeadroomGuardrail(**defaults)
return HeadroomGuardrail(**{"gateway_turn": False, **defaults})
def _make_compress_response(
@ -610,12 +613,17 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404(
@pytest.mark.asyncio
@pytest.mark.parametrize("context", [
None, {"issued_hashes": ["deadbeef000000000000dead"]},
HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0),
])
@pytest.mark.parametrize(
"context",
[
None,
{"issued_hashes": ["deadbeef000000000000dead"]},
HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0),
],
)
async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call(
guardrail: HeadroomGuardrail, context: object,
guardrail: HeadroomGuardrail,
context: object,
) -> None:
"""A hash-shaped string planted in message text must not be honored when
this guardrail has no record of ever issuing it, even if it's echoed back
@ -921,7 +929,9 @@ HASH_SHAPED_HISTORY = [
async def _apply(
guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None,
guardrail: HeadroomGuardrail,
messages: list,
ccr_hashes: list | None = None,
request_data: dict[str, object] | None = None,
) -> dict:
data: Final = request_data if request_data is not None else {"model": "claude-sonnet-5"}
@ -3149,3 +3159,320 @@ def test_affinity_requires_trusted_identity_and_caller_session(
affinity: Final = _session_affinity(candidate)
assert (affinity is None) is (session is None or not isinstance(auth, UserAPIKeyAuth))
assert (affinity == _session_affinity(baseline)) is (candidate == baseline)
def _gateway_answer(body: dict, **extra: object) -> MagicMock:
mock = MagicMock()
mock.status_code = 200
mock.text = "{}"
mock.json.return_value = {
"messages": body.get("messages"),
"body": body,
"tokens_before": 100,
"tokens_after": 40,
**extra,
}
return mock
def _raised_status(status: int, text: str) -> httpx.HTTPStatusError:
"""What AsyncHTTPHandler.post does with any non-2xx answer."""
request = httpx.Request("POST", "http://headroom/v1/compress")
return httpx.HTTPStatusError(text, request=request, response=httpx.Response(status, text=text, request=request))
def _caller(api_key: str = "hashed-key") -> dict:
return {"metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key=api_key)}}
def _gateway_guardrail(**kwargs) -> HeadroomGuardrail:
return _make_guardrail(gateway_turn=True, **kwargs)
def _posting(guardrail: HeadroomGuardrail, **mock: object):
return patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, **mock)
async def _apply_turn(guardrail: HeadroomGuardrail, data: dict) -> object:
rows = data["messages"] if "messages" in data else data["input"]
return await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(structured_messages=rows), request_data=data, input_type="request"
)
async def _session_of(guardrail: HeadroomGuardrail, data: dict) -> str | None:
with _posting(guardrail, return_value=_gateway_answer({"messages": data["messages"]})) as post:
await _apply_turn(guardrail, data)
return post.await_args.kwargs["json"]["config"].get("session_id")
ANTHROPIC_TOOL = {"name": "Read", "description": "long", "input_schema": {"type": "object"}}
SHORT_ANTHROPIC_TOOL = {"name": "Read", "description": "short", "input_schema": {"type": "object"}}
@pytest.mark.asyncio
async def test_gateway_turn_on_anthropic_route_keeps_the_wire_shape_and_applies_the_answer():
tool_result = {
"type": "tool_result",
"tool_use_id": "t1",
"content": "text",
"cache_control": {"type": "ephemeral"},
}
messages = [
{"role": "user", "content": "task"},
{"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}]},
{"role": "user", "content": [tool_result]},
]
folded = [
*messages[:2],
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "f"}]},
]
data = {
"model": "claude-opus-5-5",
"system": "system",
"tools": [ANTHROPIC_TOOL],
"max_tokens": 1024,
"stream": True,
"temperature": 0.2,
"messages": messages,
**_caller(),
}
guardrail = _gateway_guardrail()
answer = _gateway_answer({"messages": folded, "tools": [SHORT_ANTHROPIC_TOOL], "system": "trimmed"})
with _posting(guardrail, return_value=answer) as post:
result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
payload = post.await_args.kwargs["json"]
assert set(payload) == {"model", "system", "tools", "max_tokens", "stream", "messages", "config", "gateway"}
assert payload["messages"][2]["content"][0]["type"] == "tool_result"
assert payload["tools"] == [ANTHROPIC_TOOL]
assert payload["config"] == {"mode": "lossless"}
assert result["messages"] == folded
assert result["tools"] == [SHORT_ANTHROPIC_TOOL]
assert result["system"] == "trimmed"
assert result["temperature"] == 0.2
stats = _recorded_guardrail_response(result)
assert stats["tokens_saved"] == 60
entry = result["metadata"]["standard_logging_guardrail_information"][0]
assert extract_compression_saved_tokens({"guardrail_information": [entry]}) == 60
@pytest.mark.asyncio
async def test_gateway_turn_on_chat_route_applies_the_answer_and_drops_what_headroom_dropped():
tool = {"type": "function", "function": {"name": "Read", "description": "long", "parameters": {}}}
short = {"type": "function", "function": {"name": "Read", "description": "short", "parameters": {}}}
call = {"id": "c", "type": "function", "function": {"name": "Read", "arguments": "{}"}}
messages = [
{"role": "user", "content": "task"},
{"role": "assistant", "content": None, "tool_calls": [call]},
{"role": "tool", "tool_call_id": "c", "content": "text"},
]
folded = [*messages[:2], {"role": "tool", "tool_call_id": "c", "content": "folded"}]
forced = {"type": "function", "function": {"name": "Read"}}
data = {"model": "gpt-5.5", "tools": [tool], "tool_choice": forced, "messages": messages, **_caller()}
guardrail = _gateway_guardrail()
with _posting(guardrail, return_value=_gateway_answer({"messages": folded, "tools": [short]})) as post:
result = await OpenAIChatCompletionsHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
assert post.await_args.kwargs["json"]["tool_choice"] == forced
assert result["messages"] == folded
assert result["tools"] == [short]
assert "tool_choice" not in result
@pytest.mark.asyncio
async def test_gateway_turn_on_responses_route_applies_the_answer():
items = [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "task"}]},
{"type": "function_call", "call_id": "c", "name": "Read", "arguments": "{}"},
{"type": "function_call_output", "call_id": "c", "output": "text"},
]
folded = [*items[:2], {"type": "function_call_output", "call_id": "c", "output": "folded"}]
tool = {"type": "function", "name": "Read", "description": "long", "parameters": {}}
short = {"type": "function", "name": "Read", "description": "short", "parameters": {}}
data = {"model": "gpt-5.5", "instructions": "system", "tools": [tool], "input": items, **_caller()}
guardrail = _gateway_guardrail()
answer = _gateway_answer({"input": folded, "tools": [short], "instructions": "trimmed"})
with _posting(guardrail, return_value=answer) as post:
result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
payload = post.await_args.kwargs["json"]
assert payload["input"] == items and payload["instructions"] == "system"
assert result["input"] == folded
assert result["tools"] == [short]
assert result["instructions"] == "trimmed"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"answer",
[
pytest.param(lambda: _raised_status(400, "Invalid config.mode: 'lossless'"), id="mode_rejected"),
pytest.param(
lambda: MagicMock(status_code=200, text="{}", **{"json.return_value": {"messages": [{"role": "user"}]}}),
id="legacy_shape",
),
],
)
async def test_older_headroom_falls_back_and_is_asked_again_later(answer):
"""Upgrading LiteLLM before Headroom must not turn every request into a 502."""
guardrail = _gateway_guardrail()
legacy = _make_compress_response([{"role": "user", "content": "earlier compressed"}])
messages = [
{"role": "user", "content": "earlier"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "now"},
]
def turn() -> dict:
return {"model": "m", "messages": list(messages), **_caller()}
with _posting(guardrail, side_effect=[answer(), legacy, legacy, _gateway_answer({"messages": messages})]) as post:
await _apply_turn(guardrail, turn())
assert "gateway" not in post.await_args.kwargs["json"]
await _apply_turn(guardrail, turn())
assert "gateway" not in post.await_args.kwargs["json"]
guardrail._gateway_turn_retry_at = 0.0 # the retry window has passed
await _apply_turn(guardrail, turn())
assert "gateway" in post.await_args.kwargs["json"]
@pytest.mark.asyncio
async def test_other_400s_fail_closed_and_keep_the_gateway_turn():
guardrail = _gateway_guardrail()
data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()}
with _posting(guardrail, side_effect=_raised_status(400, "config.session_id must be a string")):
with pytest.raises(HTTPException) as raised:
await _apply_turn(guardrail, data)
assert raised.value.status_code == 502
assert guardrail._gateway_turn_retry_at == 0.0
@pytest.mark.asyncio
async def test_gateway_turn_retries_a_session_lock_timeout_once():
guardrail = _gateway_guardrail()
folded = [{"role": "user", "content": "folded"}]
data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()}
with _posting(guardrail, side_effect=[_raised_status(503, "busy"), _gateway_answer({"messages": folded})]) as post:
await _apply_turn(guardrail, data)
assert post.await_count == 2
assert data["messages"] == folded
@pytest.mark.asyncio
async def test_gateway_session_is_one_conversation_of_one_caller():
guardrail = _gateway_guardrail()
first = {"role": "user", "content": "a"}
marked = {"role": "user", "content": [{"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}]}
image = {"type": "image", "source": {"type": "url", "url": "https://x.test/a.png"}}
def turn(row: dict, *rest: dict, api_key: str = "one", session: str = "s", user: str | None = None) -> dict:
extra = {"user": user} if user else {}
return {"model": "m", "messages": [row, *rest], "session_id": session, **_caller(api_key), **extra}
base = await _session_of(guardrail, turn(first))
assert base is not None
assert base == await _session_of(guardrail, turn(first, {"role": "assistant", "content": "b"}))
assert base == await _session_of(guardrail, turn(marked))
assert base != await _session_of(guardrail, turn(first, session="other"))
assert base != await _session_of(guardrail, turn(first, api_key="two"))
assert await _session_of(guardrail, turn(first, user="alice")) != await _session_of(
guardrail, turn(first, user="bob")
)
# Claude Code's subagents share the parent's session id.
assert base != await _session_of(guardrail, turn({"role": "user", "content": "subagent"}))
with_image = {"role": "user", "content": [image]}
marked_image = {"role": "user", "content": [{**image, "cache_control": {"type": "ephemeral"}}]}
assert await _session_of(guardrail, turn(with_image)) == await _session_of(guardrail, turn(marked_image))
@pytest.mark.asyncio
async def test_gateway_turn_without_a_caller_session_stays_stateless():
"""Many end users behind one key can open with the same row, so they never share a session."""
guardrail = _gateway_guardrail()
data = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **_caller()}
assert await _session_of(guardrail, data) is None
@pytest.mark.asyncio
async def test_auto_router_hop_gets_compressed_structured_messages():
"""The routing hop has no caller on it and reads only the returned messages."""
guardrail = _gateway_guardrail()
legacy_rows = [{"role": "user", "content": "earlier compressed"}]
messages = [
{"role": "user", "content": "earlier"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "now"},
]
with _posting(guardrail, return_value=_make_compress_response(legacy_rows)) as post:
result = await _apply_turn(guardrail, {"messages": messages, "model": "router"})
assert "gateway" not in post.await_args.kwargs["json"]
assert result["structured_messages"][0]["content"] == "earlier compressed"
@pytest.mark.asyncio
async def test_embeddings_input_never_reaches_headroom():
guardrail = _gateway_guardrail()
data = {"model": "m", "input": ["some text"], **_caller()}
inputs = GenericGuardrailAPIInputs(texts=["some text"])
with _posting(guardrail) as post:
assert await guardrail.apply_guardrail(inputs=inputs, request_data=data, input_type="request") is inputs
post.assert_not_awaited()
@pytest.mark.asyncio
async def test_gateway_turn_forwards_and_merges_anthropic_betas():
guardrail = _gateway_guardrail()
answer = _gateway_answer(
{"messages": [{"role": "user", "content": "x"}]}, headers={"anthropic-beta": "plugin-beta"}
)
data = {
"model": "m",
"messages": [{"role": "user", "content": "x"}],
"extra_headers": {"anthropic-beta": "configured-beta", "x-other": "kept"},
"proxy_server_request": {"headers": {"Anthropic-Beta": "client-beta"}},
**_caller(),
}
with _posting(guardrail, return_value=answer) as post:
await _apply_turn(guardrail, data)
assert post.await_args.kwargs["json"]["gateway"]["request_headers"] == {"anthropic-beta": "client-beta"}
assert data["extra_headers"] == {"anthropic-beta": "configured-beta,plugin-beta", "x-other": "kept"}
@pytest.mark.asyncio
async def test_gateway_turn_fail_open_leaves_the_request_alone():
data = {"model": "m", "messages": [{"role": "user", "content": "task"}], "tools": [{"name": "Read"}], **_caller()}
inputs = GenericGuardrailAPIInputs(structured_messages=data["messages"])
guardrail = _gateway_guardrail(unreachable_fallback="fail_open")
with _posting(guardrail, side_effect=httpx.ConnectError("down")) as post:
returned = await guardrail.apply_guardrail(inputs=inputs, request_data=data, input_type="request")
assert post.await_args.kwargs["json"]["gateway"]
assert returned is inputs
assert data["messages"] == [{"role": "user", "content": "task"}]
assert data["tools"] == [{"name": "Read"}]
@pytest.mark.asyncio
async def test_gateway_turn_fail_closed_raises_when_headroom_is_down():
data = {"model": "m", "messages": [{"role": "user", "content": "task"}], **_caller()}
guardrail = _gateway_guardrail()
with _posting(guardrail, side_effect=httpx.ConnectError("down")):
with pytest.raises(HTTPException) as raised:
await _apply_turn(guardrail, data)
assert raised.value.status_code == 502

14
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -37055,7 +37055,7 @@ export interface components {
category_thresholds?: components["schemas"]["LakeraCategoryThresholds"] | null;
/**
* Ccr Retrieval
* @description Inject the Headroom retrieval tool for hashes declared by the compression service.
* @description Inject the Headroom retrieval tool for hashes declared by the compression service. Applies only when gateway_turn is false.
* @default true
*/
ccr_retrieval: boolean;
@ -37080,6 +37080,12 @@ export interface components {
* @description Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.
*/
client_secret?: string | null;
/**
* Compression Mode
* @description Headroom only. config.mode for a gateway turn. 'lossless' loses nothing and needs no retrieval tool; 'lossy_inline' saves more without a way back. Null uses the Headroom server's default.
* @default lossless
*/
compression_mode: ("lossless" | "lossy_inline" | "lossless_then_lossy") | null;
/**
* Confidence Threshold
* @description Only block or mask when detection confidence >= this value; below threshold, allow or log_only.
@ -37237,6 +37243,12 @@ export interface components {
* @description noma_v2 only: name of this gateway, used as the gateway_host label on Noma scans
*/
gateway_name?: string | null;
/**
* Gateway Turn
* @description Headroom only. Send Headroom the whole provider request (messages, system, tools) on its gateway contract, so history stays byte-stable for the prompt cache and server-side plugins can rewrite tools and skills. With false, or against an older headroom-ai (asked again every 5 minutes), the guardrail compresses a subset of rows instead.
* @default true
*/
gateway_turn: boolean;
/**
* Grounding Check
* @description Enable grounding verification to ensure output is grounded in provided context.