From 47117d880c723120c4596b66af2f6ca8853f52e6 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 10 Sep 2026 10:13:51 +0000 Subject: [PATCH 1/3] 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 7858adeb55d..d219d116992 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -87,12 +87,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 8377db57b6e..54fede364b8 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -202,6 +202,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)], From 7b3582aa6614ad370de1b05a822af3b79b782933 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 16:31:44 +0000 Subject: [PATCH 2/3] fix(guardrails): treat tag-based Mode as MCP-only when all hooks are MCP hooks 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 | 60 ++++++++++++------- 2 files changed, 56 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index d219d116992..98bc6b5451d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -1,6 +1,8 @@ # litellm/proxy/guardrails/guardrail_initializers.py from typing import Any, Final +from typing_extensions import assert_never + import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import CommonProxyErrors @@ -96,14 +98,26 @@ _MCP_EVENT_HOOKS: Final = frozenset( ) -def _is_mcp_only_mode(mode: str | list[str] | Mode) -> bool: +def _configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]: match mode: case str(): - return mode in _MCP_EVENT_HOOKS + return (mode,) case list(): - return bool(mode) and all(m in _MCP_EVENT_HOOKS for m in mode) + return tuple(mode) case Mode(): - return False + 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) + ) + case _: + assert_never(mode) + + +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, ...]: diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 54fede364b8..fc2fb949143 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -202,45 +202,61 @@ 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, expected_hooks", + "mode, filter_scope, expect_output_scanned", [ - ("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"}), + ("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), ], ) -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 +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.""" - import litellm - from litellm.proxy.guardrails.guardrail_hooks.presidio import ( - _OPTIONAL_PresidioPIIMasking, - ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import Choices, Message, ModelResponse - guardrail_name = f"test_presidio_mcp_scope_{id(mode)}_{filter_scope}" + 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 ", "items": []}, + "default_on": True, } 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} + 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], + ] - 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 + 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( From 8a43fed20cabb1c038cf2b5389efd89dfa0b7a24 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 16:48:15 +0000 Subject: [PATCH 3/3] fix(guardrails): use explicit returns in _configured_event_hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_initializers.py | 26 +++++++------------ 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 98bc6b5451d..b7ab215a2cd 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -1,8 +1,6 @@ # litellm/proxy/guardrails/guardrail_initializers.py from typing import Any, Final -from typing_extensions import assert_never - import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import CommonProxyErrors @@ -99,20 +97,16 @@ _MCP_EVENT_HOOKS: Final = frozenset( def _configured_event_hooks(mode: str | list[str] | Mode) -> tuple[str, ...]: - match mode: - case str(): - return (mode,) - case list(): - return tuple(mode) - case 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) - ) - case _: - assert_never(mode) + 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: