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