Merge 7dfda63bb7 into 0c98afa780

litellm/proxy/guardrails/guardrail_initializers.py

@@ -86,12 +86,32 @@ 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 _is_mcp_only_mode(mode: str | list[str] | Mode) -> bool:
+    match mode:
+        case str():
+            return mode in _MCP_EVENT_HOOKS
+        case list():
+            return bool(mode) and all(m in _MCP_EVENT_HOOKS for m in mode)
+        case Mode():
+            return False
+
+
 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

@@ -156,6 +156,47 @@ def test_initialize_presidio_forwards_analyze_chunk_size_bytes():
     assert initialized[-1].presidio_analyze_chunk_size_bytes == 250_000
 
 
+@pytest.mark.parametrize(
+    "mode, filter_scope, expected_hooks",
+    [
+        ("pre_mcp_call", None, {"pre_mcp_call"}),
+        (["pre_mcp_call", "post_mcp_call"], None, {"pre_mcp_call", "post_mcp_call"}),
+        ("pre_mcp_call", "both", {"pre_mcp_call", "post_call"}),
+        ("pre_call", None, {"pre_call", "post_call"}),
+    ],
+)
+def test_initialize_presidio_mcp_mode_does_not_add_post_call_scan(mode, filter_scope, expected_hooks):
+    """Regression: a `pre_mcp_call` Presidio guardrail used to also register a
+    `post_call` output scanner, so a blocked tool call that the model mentioned in
+    its answer turned the whole request into an HTTP 400 instead of a 200."""
+    import litellm
+    from litellm.proxy.guardrails.guardrail_hooks.presidio import (
+        _OPTIONAL_PresidioPIIMasking,
+    )
+
+    guardrail_name = f"test_presidio_mcp_scope_{id(mode)}_{filter_scope}"
+    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",
+    }
+    if filter_scope is not None:
+        litellm_params["presidio_filter_scope"] = filter_scope
+
+    InMemoryGuardrailHandler().initialize_guardrail(
+        guardrail={"guardrail_name": guardrail_name, "litellm_params": litellm_params}
+    )
+
+    registered_hooks = {
+        hook
+        for callback in litellm.callbacks
+        if isinstance(callback, _OPTIONAL_PresidioPIIMasking) and callback.guardrail_name == guardrail_name
+        for hook in ([callback.event_hook] if isinstance(callback.event_hook, str) else callback.event_hook)
+    }
+    assert registered_hooks == expected_hooks
+
+
 @pytest.mark.parametrize(
     "config_value, expected",
     [(True, True), (False, False), (None, False)],