From 7dfda63bb7f1ce809b30e2fa1781d9968baf185d Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 10 Sep 2026 10:13:51 +0000 Subject: [PATCH] fix(guardrails): don't add post_call output scan for MCP-only Presidio modes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_initializers.py | 22 +++++++++- .../proxy/guardrails/test_init_guardrails.py | 41 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 16369abbfb0..52b2a7ae9e7 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/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") diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index ceb084b4a4d..68835746742 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/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)],