Merge pull request #40571 from BerriAI/litellm_presidio_mcp_mode_no_post_call_scan

fix(guardrails): don't add post_call output scan for MCP-only Presidio modes

litellm/proxy/guardrails/guardrail_initializers.py

@@ -87,12 +87,40 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail):
     return _lakera_v2_callback
 
 
+_MCP_EVENT_HOOKS: Final = frozenset(
+    {
+        GuardrailEventHooks.pre_mcp_call.value,
+        GuardrailEventHooks.during_mcp_call.value,
+        GuardrailEventHooks.post_mcp_call.value,
+    }
+)
+
+
+def _configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]:
+    if isinstance(mode, str):
+        return (mode,)
+    if isinstance(mode, list):
+        return tuple(mode)
+    return tuple(
+        hook
+        for value in (*mode.tags.values(), mode.default)
+        if value is not None
+        for hook in ((value,) if isinstance(value, str) else value)
+    )
+
+
+def _is_mcp_only_mode(mode: str | list[str] | Mode) -> bool:
+    hooks: Final = _configured_event_hooks(mode)
+    return bool(hooks) and all(hook in _MCP_EVENT_HOOKS for hook in hooks)
+
+
 def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> tuple[CustomGuardrail, ...]:
     from litellm.proxy.guardrails.guardrail_hooks.presidio import (
         _OPTIONAL_PresidioPIIMasking,
     )
 
-    filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None) or "both"
+    explicit_filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None)
+    filter_scope: Final = explicit_filter_scope or ("input" if _is_mcp_only_mode(litellm_params.mode) else "both")
     run_input: Final = filter_scope in ("input", "both")
     run_output: Final = filter_scope in ("output", "both")
 

tests/test_litellm/proxy/guardrails/test_init_guardrails.py

@@ -202,6 +202,63 @@ def test_initialize_presidio_forwards_analyze_chunk_size_bytes():
     assert initialized[-1].presidio_analyze_chunk_size_bytes == 250_000
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "mode, filter_scope, expect_output_scanned",
+    [
+        ("pre_mcp_call", None, False),
+        (["pre_mcp_call", "post_mcp_call"], None, False),
+        ({"tags": {"team:mcp": "pre_mcp_call"}, "default": ["pre_mcp_call", "post_mcp_call"]}, None, False),
+        ({"tags": {"team:mcp": ["pre_mcp_call"]}, "default": "pre_call"}, None, True),
+        ({"tags": {}}, None, True),
+        ("pre_mcp_call", "both", True),
+        ("pre_mcp_call", "output", True),
+        ("pre_call", None, True),
+    ],
+)
+async def test_initialize_presidio_mcp_only_mode_skips_post_call_output_scan(mode, filter_scope, expect_output_scanned):
+    """Regression: an MCP-only Presidio guardrail used to also scan the LLM
+    response on post_call, so a blocked MCP tool call that the model repeated in
+    its answer turned the whole request into an HTTP 400 instead of a 200."""
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.guardrails import GuardrailEventHooks
+    from litellm.types.utils import Choices, Message, ModelResponse
+
+    llm_answer = "Call me at 415-555-2671"
+    litellm_params = {
+        "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
+        "mode": mode,
+        "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
+        "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
+        "mock_redacted_text": {"text": "Call me at <PHONE_NUMBER>", "items": []},
+        "default_on": True,
+    }
+    if filter_scope is not None:
+        litellm_params["presidio_filter_scope"] = filter_scope
+
+    guardrail_handler = InMemoryGuardrailHandler()
+    result = guardrail_handler.initialize_guardrail(
+        guardrail={"guardrail_name": "test_presidio_mcp_scope", "litellm_params": litellm_params}
+    )
+    guardrail_id = result["guardrail_id"]
+    callbacks = [
+        guardrail_handler.guardrail_id_to_custom_guardrail[guardrail_id],
+        *guardrail_handler.guardrail_id_to_sibling_callbacks[guardrail_id],
+    ]
+
+    request_data = {"metadata": {}}
+    response = ModelResponse(
+        choices=[Choices(message=Message(role="assistant", content=llm_answer), index=0, finish_reason="stop")]
+    )
+    for callback in callbacks:
+        if callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call):
+            await callback.async_post_call_success_hook(
+                data=request_data, user_api_key_dict=UserAPIKeyAuth(), response=response
+            )
+
+    assert (response.choices[0].message.content != llm_answer) is expect_output_scanned
+
+
 @pytest.mark.parametrize(
     "config_value, expected",
     [(True, True), (False, False), (None, False)],