fix(realtime): screen one joined text per item and pass the session's tenant context to guardrails

Multi-part items reach each guardrail once as the space-joined text, as on main, and only an item a
guardrail rewrites gets its parts screened again to keep per-part masks. Guardrails now receive the
session request data with its metadata, so key and team policies apply, and a ModifyResponseException
block is reported as a block. Nova scopes every message item as the user turn it sends to Bedrock.

litellm/litellm_core_utils/realtime_guardrails.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias
 from pydantic import TypeAdapter, ValidationError
 
 from litellm._logging import verbose_logger
+from litellm.exceptions import ModifyResponseException
 
 if TYPE_CHECKING:
     from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -132,7 +133,7 @@ def _violation_text(error: Exception) -> str:
 
 
 def _is_guardrail_block(error: Exception) -> bool:
-    return hasattr(error, "status_code") or isinstance(error, ValueError)
+    return hasattr(error, "status_code") or isinstance(error, (ValueError, ModifyResponseException))
 
 
 def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]:
@@ -145,12 +146,12 @@ def _screened_texts(screened: object, texts: tuple[str, ...]) -> tuple[str, ...]
 
 
 async def _guardrail_outcome(
-    guardrail: "CustomGuardrail", texts: tuple[str, ...], user_api_key_dict: object | None
+    guardrail: "CustomGuardrail", texts: tuple[str, ...], request_data: Mapping[str, object]
 ) -> RealtimeGuardrailBlock | tuple[str, ...]:
     try:
         screened: Final = await guardrail.apply_guardrail(
             inputs={"texts": list(texts), "images": []},
-            request_data={"user_api_key_dict": user_api_key_dict},
+            request_data=dict(request_data),
             input_type="request",
         )
     except Exception as error:
@@ -166,42 +167,27 @@ async def _guardrail_outcome(
 
 
 async def _screen(
-    guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], user_api_key_dict: object | None
+    guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object]
 ) -> RealtimeGuardrailBlock | tuple[str, ...]:
     if not guardrails:
         return texts
-    outcome: Final = await _guardrail_outcome(guardrails[0], texts, user_api_key_dict)
+    outcome: Final = await _guardrail_outcome(guardrails[0], texts, request_data)
     if isinstance(outcome, RealtimeGuardrailBlock):
         return outcome
-    return await _screen(guardrails[1:], outcome, user_api_key_dict)
+    return await _screen(guardrails[1:], outcome, request_data)
 
 
-def _joins(texts: tuple[str, ...]) -> tuple[str, str]:
-    return "".join(texts), " ".join(texts)
-
-
-def _screening_texts(texts: tuple[str, ...]) -> tuple[str, ...]:
-    return texts if len(texts) < 2 else (*texts, *_joins(texts))
-
-
-async def _rewritten_texts(
-    guardrails: Sequence["CustomGuardrail"],
-    texts: tuple[str, ...],
-    screened: tuple[str, ...],
-    user_api_key_dict: object | None,
+async def _rewritten_parts(
+    guardrails: Sequence["CustomGuardrail"], texts: tuple[str, ...], request_data: Mapping[str, object]
 ) -> RealtimeGuardrailBlock | tuple[str, ...]:
-    parts: Final = screened[: len(texts)]
-    if len(texts) < 2:
+    parts: Final = await _screen(guardrails, texts, request_data)
+    if isinstance(parts, RealtimeGuardrailBlock):
         return parts
-    joins: Final = _joins(parts)
-    screened_joins: Final = (
-        screened[len(texts) :] if parts == texts else await _screen(guardrails, joins, user_api_key_dict)
-    )
-    if isinstance(screened_joins, RealtimeGuardrailBlock):
-        return screened_joins
-    if screened_joins == joins:
-        return parts
-    return (screened_joins[0],) if screened_joins[0] != joins[0] else (screened_joins[1],)
+    rejoined: Final = (" ".join(parts),)
+    screened_rejoined: Final = await _screen(guardrails, rejoined, request_data)
+    if isinstance(screened_rejoined, RealtimeGuardrailBlock):
+        return screened_rejoined
+    return parts if screened_rejoined == rejoined else screened_rejoined
 
 
 def _in_role_scope(guardrail: "CustomGuardrail", role: str) -> bool:
@@ -251,10 +237,14 @@ async def realtime_guardrail_verdict(
         if guardrail is not None and _in_role_scope(guardrail, role)
     }
     in_scope: Final = tuple(guardrails.values())
-    screened: Final = await _screen(in_scope, _screening_texts(texts), user_api_key_dict)
+    guardrail_request_data: Final = {**request_data, "user_api_key_dict": user_api_key_dict}
+    joined: Final = (" ".join(texts),)
+    screened: Final = await _screen(in_scope, joined, guardrail_request_data)
     if isinstance(screened, RealtimeGuardrailBlock):
         return screened
-    rewritten: Final = await _rewritten_texts(in_scope, texts, screened, user_api_key_dict)
+    if screened == joined:
+        return None
+    rewritten: Final = screened if len(texts) < 2 else await _rewritten_parts(in_scope, texts, guardrail_request_data)
     if isinstance(rewritten, RealtimeGuardrailBlock):
         return rewritten
     return None if rewritten == texts else RealtimeGuardrailRewrite(texts=rewritten)

litellm/llms/bedrock/realtime/handler.py

@@ -452,7 +452,7 @@ class BedrockRealtimeGuardrails:
         guarded_item: Final = realtime_guarded_item(client_event)
         if guarded_item is None:
             return _ClientEventRoute.FORWARD
-        role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else guarded_item.role
+        role: Final = "tool" if isinstance(guarded_item, RealtimeToolOutput) else "user"
         match await self._verdict(guarded_item.texts, role, GuardrailEventHooks.pre_call):
             case None:
                 return _ClientEventRoute.FORWARD

tests/unit/litellm_core_utils/test_realtime_guardrails.py

@@ -7,6 +7,7 @@ from fastapi import HTTPException
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
+from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail
 from litellm.litellm_core_utils.realtime_guardrails import (
     RealtimeGuardrailBlock,
     RealtimeGuardrailRewrite,
@@ -142,7 +143,19 @@ async def test_masked_texts_come_back_as_a_rewrite_and_reach_later_guardrails_ma
     )
 
     assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED]", "and [MASKED]", "hi"))
-    assert second.screened[0][:3] == ("my [MASKED]", "and swordfish", "hi")
+    assert second.screened[0] == ("my [MASKED] and swordfish hi",)
+
+
+@pytest.mark.asyncio
+async def test_a_clean_multi_part_item_reaches_each_guardrail_once_as_the_space_joined_text() -> None:
+    guardrail: Final = MaskingGuardrail("hunter2")
+
+    verdict: Final = await realtime_guardrail_verdict(
+        ("Compare this photo", "with that one"), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,)
+    )
+
+    assert verdict is None
+    assert guardrail.screened == [("Compare this photo with that one",)]
 
 
 @pytest.mark.asyncio
@@ -155,23 +168,20 @@ async def test_texts_a_guardrail_leaves_unchanged_need_no_rewrite() -> None:
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize(
-    ("keyword", "parts"),
-    [
-        pytest.param("secret plan", ("tell me the secret", "plan"), id="split at a space"),
-        pytest.param("secret-project", ("Tell me about secret-", "project"), id="split mid-word"),
-    ],
-)
-async def test_a_phrase_split_across_typed_text_parts_is_still_blocked(keyword: str, parts: tuple[str, ...]) -> None:
+async def test_a_phrase_split_across_typed_text_parts_is_still_blocked() -> None:
     client_event: Final = _item_create(
-        {"type": "message", "role": "user", "content": [{"type": "input_text", "text": part} for part in parts]}
+        {
+            "type": "message",
+            "role": "user",
+            "content": [{"type": "input_text", "text": "tell me the secret"}, {"type": "input_text", "text": "plan"}],
+        }
     )
     guarded_item: Final = realtime_guarded_item(client_event)
     assert isinstance(guarded_item, RealtimeMessageText)
 
     phrase_filter: Final = ContentFilterGuardrail(
         guardrail_name="phrase-filter",
-        blocked_words=[BlockedWord(keyword=keyword, action=ContentFilterAction.BLOCK)],
+        blocked_words=[BlockedWord(keyword="secret plan", action=ContentFilterAction.BLOCK)],
         event_hook=GuardrailEventHooks.pre_call,
         default_on=True,
     )
@@ -218,7 +228,12 @@ async def test_a_mask_inside_one_part_keeps_every_part_in_place() -> None:
 @pytest.mark.asyncio
 async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_would_have_read() -> None:
     verdict: Final = await realtime_guardrail_verdict(
-        ("my hunt", "er2 is here"), "user", (GuardrailEventHooks.pre_call,), {}, None, (MaskingGuardrail("hunter2"),)
+        ("my secret", "plan is here"),
+        "user",
+        (GuardrailEventHooks.pre_call,),
+        {},
+        None,
+        (MaskingGuardrail("secret plan"),),
     )
 
     assert verdict == RealtimeGuardrailRewrite(texts=("my [MASKED] is here",))
@@ -228,8 +243,8 @@ async def test_a_mask_across_parts_replaces_them_with_the_masked_text_the_model_
 @pytest.mark.parametrize(
     "parts",
     [
-        pytest.param(("abc", "123-45-6789 hunter2"), id="masks inside one part"),
-        pytest.param(("abc", "123-45-6789 hunt", "er2"), id="a mask inside a part and one across parts"),
+        pytest.param(("abc", "123-45-6789 hunter2 rocks"), id="masks inside one part"),
+        pytest.param(("abc", "123-45-6789 hunter2", "rocks"), id="a mask inside a part and one across parts"),
     ],
 )
 async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_differently(
@@ -240,7 +255,7 @@ async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_dif
         patterns=[
             ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)
         ],
-        blocked_words=[BlockedWord(keyword="hunter2", action=ContentFilterAction.MASK)],
+        blocked_words=[BlockedWord(keyword="hunter2 rocks", action=ContentFilterAction.MASK)],
         event_hook=GuardrailEventHooks.pre_call,
         default_on=True,
     )
@@ -250,9 +265,9 @@ async def test_no_mask_one_text_part_gets_is_lost_when_its_joined_text_masks_dif
     )
 
     assert isinstance(verdict, RealtimeGuardrailRewrite)
-    forwarded: Final = "".join(verdict.texts)
+    forwarded: Final = " ".join(verdict.texts)
     assert "123-45-6789" not in forwarded
-    assert "hunter2" not in forwarded
+    assert "hunter2 rocks" not in forwarded
     assert forwarded.startswith("abc")
 
 
@@ -328,3 +343,33 @@ async def test_a_guardrail_that_crashes_is_not_treated_as_a_block() -> None:
 
     with pytest.raises(RuntimeError, match="guardrail backend down"):
         await realtime_guardrail_verdict(("hello",), "user", (GuardrailEventHooks.pre_call,), {}, None, (guardrail,))
+
+
+_TEAM_POLICY_CODE: Final = (
+    "def apply_guardrail(inputs, request_data, input_type):\n"
+    '    if request_data["team_id"] == "team-strict":\n'
+    '        return block("blocked for this team")\n'
+    "    return allow()\n"
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("team_id", "blocked"),
+    [pytest.param("team-strict", True, id="policy team"), pytest.param("team-other", False, id="other team")],
+)
+async def test_a_team_scoped_guardrail_sees_the_session_team(team_id: str, blocked: bool) -> None:
+    team_policy = CustomCodeGuardrail(
+        custom_code=_TEAM_POLICY_CODE, guardrail_name="team-policy", event_hook="pre_call", default_on=True
+    )
+
+    verdict = await realtime_guardrail_verdict(
+        ("hello",),
+        "user",
+        (GuardrailEventHooks.pre_call,),
+        {"litellm_metadata": {"user_api_key_team_id": team_id}},
+        None,
+        (team_policy,),
+    )
+
+    assert isinstance(verdict, RealtimeGuardrailBlock) is blocked

tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py

@@ -1937,7 +1937,9 @@ class TestBedrockRealtimeGuardrails:
         assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION]
 
     @pytest.mark.asyncio
-    async def test_a_system_message_the_guardrail_skips_reaches_bedrock(self, stub_aws_models: None) -> None:
+    async def test_a_system_message_nova_receives_as_a_user_turn_is_screened_by_a_system_skipping_guardrail(
+        self, stub_aws_models: None
+    ) -> None:
         guardrail: Final = SystemSkippingGuardrail(
             "hunter2",
             guardrail_name="no_secrets",
@@ -1946,13 +1948,10 @@ class TestBedrockRealtimeGuardrails:
             realtime_violation_message=self.VIOLATION,
         )
 
-        stream, client_ws = await self._forward_client(
-            [self._user_text(self.BLOCKED_TEXT, role="system"), self._user_text(self.BLOCKED_TEXT)], guardrail
-        )
+        stream, client_ws = await self._forward_client([self._user_text(self.BLOCKED_TEXT, role="system")], guardrail)
 
-        assert self.BLOCKED_TEXT in json.dumps(self._bedrock_events(stream))
+        assert "hunter2" not in json.dumps(self._bedrock_events(stream))
         assert [error["message"] for error in self._guardrail_errors(client_ws)] == [self.VIOLATION]
-        assert guardrail.checked == [((self.BLOCKED_TEXT,), None)]
 
     @pytest.mark.asyncio
     async def test_the_next_clean_turn_is_answered_after_a_blocked_item_with_no_response_create(