From 155e4080c4ba956d198d247fe7717c19a4f69712 Mon Sep 17 00:00:00 2001 From: Itay Ovadia Date: Sat, 15 Aug 2026 22:06:49 +0300 Subject: [PATCH] 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. --- litellm/proxy/guardrails/_content_utils.py | 14 +++ .../proxy/guardrails/guardrail_endpoints.py | 10 +- .../generic_guardrail_api.py | 12 +- .../generic_guardrail_api/payload_policy.py | 52 ++++++-- .../generic_guardrail_api/record_scope.py | 11 +- .../test_generic_guardrail_api.py | 113 ++++++++++++++++++ .../guardrails/test_guardrail_endpoints.py | 74 ++++++++---- 7 files changed, 248 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 635034cecb7..d0b4e1d9158 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -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``. diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b68d4a68b79..164e6b14b59 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 797c252bdaf..a69a17ae678 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py index 7b3e1c014dc..f7192579087 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/payload_policy.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py index 28a9430bb9a..802dddb3e3c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/record_scope.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index a417be5e69b..6a5e3f8a74e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 71ff9111b60..d98046eacbc 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -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