fix(guardrails): align the max_messages window and stop identity being taken from the request body

Three findings from review of the previous commit.

max_messages applied the same number to two lists that count different things:
texts is fragment-based (a multimodal turn contributes several entries, a
tool-call-only turn contributes none) while structured_messages is message-based,
so the two views of the payload could cover different parts of the conversation,
sending text from an omitted turn or dropping text from a retained one. The
message window is now the source of truth and the text window is the number of
fragments those retained messages contain, counted with a role-agnostic helper so
an unenumerated role cannot silently shrink it.

/guardrails/apply_guardrail built request_data straight from the caller's
metadata, so a caller could claim a user_api_key_alias or user_api_key_team_id
that a guardrail exempts and get its text back unscanned. The authenticated
identity is now layered on top of the caller's metadata via the same sanitizer
the in-line path uses, so what authentication established wins. Caller metadata
that is not identity still passes through untouched, and every guardrail on that
route now sees the caller information it already sees in-line.

The per_session dedup cache was keyed only on the caller-supplied session id, so
whoever claimed an id first suppressed another tenant's entry for the same id.
The key is now namespaced by the authenticated key hash, falling back to the team
id.
This commit is contained in:
Itay Ovadia 2026-08-15 22:06:49 +03:00
parent c3ad8ac34e
commit 155e4080c4
7 changed files with 248 additions and 38 deletions

View file

@ -112,6 +112,20 @@ def iter_message_text(data: dict[str, Any]) -> Iterator[str]:
yield from _iter_text_parts_in_content(message.get("content"))
def iter_messages_text(messages: Sequence[Mapping[str, object]]) -> Iterator[str]:
"""Yield every text fragment carried by a message list, whatever the role.
Role-agnostic on purpose: a caller counting what a set of messages
contributes (payload windowing, size accounting) must not miss a turn whose
role it did not think to enumerate.
"""
for message in messages:
# The annotation is the contract; the values arrive from provider
# translation handlers, so a stray non-message entry must not raise.
if isinstance(message, Mapping): # pyright: ignore[reportUnnecessaryIsInstance] # untyped upstream data
yield from _iter_text_parts_in_content(message.get("content"))
def iter_role_text(messages: Sequence[Mapping[str, object]], roles: Collection[str]) -> Iterator[str]:
"""Yield every text fragment carried by messages whose role is in ``roles``.

View file

@ -2288,6 +2288,7 @@ async def apply_guardrail(
executor as thread_pool_executor,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.proxy_server import (
general_settings,
proxy_config,
@ -2332,9 +2333,16 @@ async def apply_guardrail(
if litellm_logging_obj is not None:
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
# The caller supplies this metadata, so the authenticated identity is
# layered on top of it: a guardrail keyed on user_api_key_alias or
# user_api_key_team_id (exemptions, per-tenant behavior) must not be
# steerable by what the request body claims to be.
authenticated_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
user_api_key_dict
)
request_data: Final[dict] = {
**({"messages": request.messages} if request.messages is not None else {}),
**({"metadata": request.metadata} if request.metadata is not None else {}),
"metadata": {**(request.metadata or {}), **(authenticated_metadata or {})},
}
_input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(

View file

@ -605,7 +605,17 @@ class GenericGuardrailAPI(CustomGuardrail):
input_type=input_type,
logging_obj=logging_obj,
)
if self._record_scope.should_suppress(get_session_id_from_request_data(request_data or {})):
# The session id is caller-supplied, so the dedup key is namespaced by the
# authenticated caller: the key hash is unique per virtual key, with the
# team id as the fallback. Both come from auth, not from the body.
identity: Final = (
self._extract_user_api_key_metadata(request_data) if request_data else GenericGuardrailAPIMetadata()
)
session_id: Final = get_session_id_from_request_data(request_data) if request_data else None
if self._record_scope.should_suppress(
session_id,
tenant=identity.get("user_api_key_hash") or identity.get("user_api_key_team_id"),
):
suppress_guardrail_information_record()
return result

View file

@ -21,6 +21,7 @@ from pydantic import JsonValue
from litellm._logging import verbose_proxy_logger
from litellm.proxy.guardrails._content_utils import (
iter_messages_text,
map_messages_image_urls,
map_messages_text,
)
@ -121,7 +122,39 @@ def _shape_text(text: str, policy: PayloadPolicy) -> str:
return stripped[: policy.max_text_chars] if policy.max_text_chars is not None else stripped
def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue, frozenset[int]]:
def _retained_messages(messages: JsonValue, policy: PayloadPolicy) -> list[JsonValue] | None:
"""The message window ``max_messages`` keeps, or None when there is nothing to window."""
if not isinstance(messages, list) or not messages:
return None
if policy.max_messages is None:
return messages
return messages[-policy.max_messages :]
def _text_window_size(retained: list[JsonValue] | None, policy: PayloadPolicy, total_texts: int) -> int:
"""How many trailing ``texts`` entries belong to the retained messages.
``texts`` is fragment-based (a multimodal turn contributes several entries, a
tool-call-only turn contributes none) while ``max_messages`` counts messages,
so applying the same number to both would leave the two views of the payload
covering different parts of the conversation. Counting the fragments inside
the retained messages keeps them aligned. Without structured messages to
count (a response payload, an embedding call) the message count is the only
bound available, so it is applied to the fragments directly.
"""
if policy.max_messages is None:
return total_texts
if retained is None:
return min(policy.max_messages, total_texts)
fragments: Final = sum(1 for _ in iter_messages_text(retained))
return min(fragments, total_texts)
def _shape_texts(
texts: Sequence[str],
policy: PayloadPolicy,
window_size: int,
) -> tuple[JsonValue, frozenset[int]]:
"""Window and rewrite the flat text list, and report which indices changed.
``max_messages`` has to bound ``texts`` as well as ``structured_messages``:
@ -131,19 +164,18 @@ def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue
shifts every position, so the caller treats the whole list as unusable for
write-back (see ``merge_guardrailed_texts``).
"""
windowed: Final = texts[-policy.max_messages :] if policy.max_messages is not None else texts
windowed: Final = texts[len(texts) - window_size :] if window_size < len(texts) else texts
shaped: Final = [_shape_text(text, policy) for text in windowed] # mutable-ok: JSON texts is an array
if len(windowed) != len(texts):
return shaped, frozenset(range(len(texts)))
return shaped, frozenset(index for index, text in enumerate(texts) if text != shaped[index])
def _shape_messages(messages: JsonValue, policy: PayloadPolicy) -> JsonValue:
if not isinstance(messages, list) or not messages:
def _shape_messages(retained: list[JsonValue] | None, policy: PayloadPolicy) -> JsonValue:
if retained is None:
return None
windowed: Final = messages[-policy.max_messages :] if policy.max_messages is not None else messages
texted: Final = (
map_messages_text(windowed, lambda text: _shape_text(text, policy)) if policy.shapes_text else windowed
map_messages_text(retained, lambda text: _shape_text(text, policy)) if policy.shapes_text else retained
)
# send_images=False has to cover the inline image parts too, otherwise the
# base64 payload still leaves LiteLLM inside structured_messages. The part is
@ -170,8 +202,12 @@ def shape_payload(
texts: Final = request.texts
texts_sent: Final = bool(texts) and "texts" not in exclude
shaped_texts, altered_indices = _shape_texts(texts or (), policy) if texts_sent else (None, frozenset[int]())
shaped_messages: Final = _shape_messages(dumped.get("structured_messages"), policy)
retained: Final = _retained_messages(dumped.get("structured_messages"), policy)
window_size: Final = _text_window_size(retained, policy, len(texts or ()))
shaped_texts, altered_indices = (
_shape_texts(texts or (), policy, window_size) if texts_sent else (None, frozenset[int]())
)
shaped_messages: Final = _shape_messages(retained, policy)
payload: Final = { # mutable-ok: the POST body is a JSON object
**dumped,

View file

@ -44,13 +44,17 @@ class RecordScope:
def scope(self) -> GuardrailInformationScope:
return self._scope
def should_suppress(self, session_id: str | None) -> bool:
def should_suppress(self, session_id: str | None, *, tenant: str | None = None) -> bool:
"""Whether this successful call should skip its logging entry.
``per_session`` claims the session on its first call, so the calls that
follow it suppress. A session id is required: without one there is
nothing to dedup against, so the call records as it would under
``per_call`` rather than silently dropping every entry.
The session id is caller-supplied, so the key is namespaced by the
authenticated caller: otherwise whoever claims an id first would
suppress another tenant's telemetry for the same id.
"""
if self._scope == "per_call":
return False
@ -58,7 +62,8 @@ class RecordScope:
return True
if session_id is None:
return False
if self._recorded_sessions.get_cache(session_id) is True:
key: Final = f"{tenant or ''}::{session_id}"
if self._recorded_sessions.get_cache(key) is True:
return True
self._recorded_sessions.set_cache(session_id, True, ttl=_SESSION_CACHE_TTL_SECONDS)
self._recorded_sessions.set_cache(key, True, ttl=_SESSION_CACHE_TTL_SECONDS)
return False

View file

@ -3190,3 +3190,116 @@ class TestConfigValidationWarnings:
assert "tools" not in handler.payloads[0]
assert result["tools"] == original_tools
class TestMaxMessagesWindowAlignment:
"""The texts window must cover the same turns as the message window.
`texts` is fragment-based and `max_messages` counts messages, so applying the
same number to both would let the two views of the payload describe different
parts of the conversation.
"""
@pytest.mark.asyncio
async def test_multipart_turn_does_not_leak_earlier_messages(self):
"""A retained turn with two text parts must not push an omitted turn's text in."""
handler = _RecordingHandler()
guardrail = _make_guardrail(handler, max_messages=1)
messages = [
{"role": "user", "content": "old turn that must not be sent"},
{
"role": "assistant",
"content": [
{"type": "text", "text": "kept part one"},
{"type": "text", "text": "kept part two"},
],
},
]
await guardrail.apply_guardrail(
inputs={
"texts": ["old turn that must not be sent", "kept part one", "kept part two"],
"structured_messages": messages,
},
request_data={},
input_type="request",
)
payload = handler.payloads[0]
assert payload["texts"] == ["kept part one", "kept part two"]
assert [m["role"] for m in payload["structured_messages"]] == ["assistant"]
@pytest.mark.asyncio
async def test_textless_turn_does_not_drop_retained_text(self):
"""A retained tool-call-only turn contributes no fragment, so the window shrinks."""
handler = _RecordingHandler()
guardrail = _make_guardrail(handler, max_messages=2)
messages = [
{"role": "user", "content": "first"},
{"role": "user", "content": "second"},
{
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "run", "arguments": "{}"}}],
},
]
await guardrail.apply_guardrail(
inputs={"texts": ["first", "second"], "structured_messages": messages},
request_data={},
input_type="request",
)
payload = handler.payloads[0]
# The window keeps the last two messages, which carry exactly one fragment.
assert payload["texts"] == ["second"]
assert [m["role"] for m in payload["structured_messages"]] == ["user", "assistant"]
@pytest.mark.asyncio
async def test_response_payload_without_messages_still_bounded(self):
"""Nothing to align against, so the message count bounds the fragments."""
handler = _RecordingHandler()
guardrail = _make_guardrail(handler, max_messages=2)
await guardrail.apply_guardrail(
inputs={"texts": ["a", "b", "c", "d"]},
request_data={},
input_type="response",
)
assert handler.payloads[0]["texts"] == ["c", "d"]
class TestSessionScopeIsolation:
"""per_session dedup must not let one caller suppress another's telemetry."""
@staticmethod
def _request_data(session_id: str, key_hash: str) -> dict:
return {
"litellm_session_id": session_id,
"metadata": {"user_api_key_hash": key_hash},
}
@pytest.mark.asyncio
async def test_same_session_id_from_another_key_still_records(self):
handler = _RecordingHandler()
guardrail = _make_guardrail(handler, guardrail_information_scope="per_session")
first = self._request_data("shared-id", "hash-tenant-a")
second = self._request_data("shared-id", "hash-tenant-b")
await guardrail.apply_guardrail(inputs={"texts": ["a"]}, request_data=first, input_type="request")
await guardrail.apply_guardrail(inputs={"texts": ["b"]}, request_data=second, input_type="request")
assert len(_recorded_entries(first)) == 1
assert len(_recorded_entries(second)) == 1
@pytest.mark.asyncio
async def test_same_key_and_session_still_dedups(self):
handler = _RecordingHandler()
guardrail = _make_guardrail(handler, guardrail_information_scope="per_session")
data = self._request_data("shared-id", "hash-tenant-a")
for _ in range(3):
await guardrail.apply_guardrail(inputs={"texts": ["a"]}, request_data=data, input_type="request")
assert len(_recorded_entries(data)) == 1

View file

@ -1552,11 +1552,41 @@ async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker):
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["What are tax loopholes?"]},
request_data={"metadata": {"forbidden_topics": ["tax"]}},
input_type="request",
call = mock_guardrail.apply_guardrail.await_args.kwargs
assert call["inputs"] == {"texts": ["What are tax loopholes?"]}
assert call["input_type"] == "request"
assert call["request_data"]["metadata"]["forbidden_topics"] == ["tax"]
@pytest.mark.asyncio
async def test_apply_guardrail_authenticated_identity_overrides_client_metadata(mocker):
"""A caller must not be able to claim another identity in the request body.
Guardrails key exemptions and per-tenant behavior off user_api_key_* metadata,
so the authenticated values have to win over whatever the body says.
"""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
request = ApplyGuardrailRequest(
guardrail_name="test-guardrail",
text="hello",
metadata={
"user_api_key_alias": "exempt-batch-worker",
"user_api_key_team_id": "exempt-team",
"forbidden_topics": ["tax"],
},
)
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(key_alias="real-caller", team_id="real-team"),
)
metadata = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
assert metadata["user_api_key_alias"] == "real-caller"
assert metadata["user_api_key_team_id"] == "real-team"
# Non-identity metadata the caller sent is still forwarded untouched.
assert metadata["forbidden_topics"] == ["tax"]
@pytest.mark.asyncio
@ -1578,33 +1608,28 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["What are tax loopholes?"]},
request_data={
"messages": messages,
"metadata": {"forbidden_topics": ["tax"]},
},
input_type="request",
)
call = mock_guardrail.apply_guardrail.await_args.kwargs
assert call["request_data"]["messages"] == messages
assert call["request_data"]["metadata"]["forbidden_topics"] == ["tax"]
@pytest.mark.asyncio
async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
"""Without metadata, request_data stays empty (backward-compatible)."""
async def test_apply_guardrail_carries_authenticated_metadata_when_none_sent(mocker):
"""request_data always carries the authenticated identity, even with no client
metadata, so a guardrail sees the same caller information it sees in-line."""
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
await apply_guardrail(
fastapi_request=mocker.Mock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
user_api_key_dict=UserAPIKeyAuth(key_alias="known-caller"),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["hello"]},
request_data={},
input_type="request",
)
call = mock_guardrail.apply_guardrail.await_args.kwargs
assert call["inputs"] == {"texts": ["hello"]}
assert "messages" not in call["request_data"]
assert call["request_data"]["metadata"]["user_api_key_alias"] == "known-caller"
@pytest.mark.asyncio
@ -1625,11 +1650,10 @@ async def test_apply_guardrail_forwards_explicit_empty_messages_and_metadata(moc
user_api_key_dict=UserAPIKeyAuth(),
)
mock_guardrail.apply_guardrail.assert_awaited_once_with(
inputs={"texts": ["hello"]},
request_data={"messages": [], "metadata": {}},
input_type="request",
)
call = mock_guardrail.apply_guardrail.await_args.kwargs
assert call["request_data"]["messages"] == []
# The explicitly-empty metadata is still forwarded, now carrying identity.
assert "metadata" in call["request_data"]
@pytest.mark.asyncio