From e2b03fd8bd7c403e433c4cd2316c031ef4c60efb Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 14:23:56 -0700 Subject: [PATCH] fix(proxy): check an MCP call's skip against the guardrail's MCP pre-call event --- litellm/proxy/utils.py | 8 +++- tests/unit/proxy/test_proxy_utils.py | 62 +++++++++++++++++++++------- 2 files changed, 54 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4400da70d45..ea82b7f4ac0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1687,13 +1687,17 @@ def _is_custom_pre_call_hook(callback: CustomLogger) -> bool: async def _skipped_for_unroutable_model( guardrail: CustomGuardrail, data: Mapping[str, object], + call_type: str, probe: "UnroutableModelProbe | None", caps: _CallbackCapabilities, ) -> bool: + event_type: Final = ( + GuardrailEventHooks.pre_mcp_call if call_type in MCP_GUARDRAIL_CALL_TYPES else GuardrailEventHooks.pre_call + ) return ( probe is not None and caps.resolved_callbacks.index(guardrail) > caps.last_custom_pre_call_hook - and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + and guardrail.should_run_guardrail(data=data, event_type=event_type) # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict and await probe.refuses(data) ) @@ -3147,6 +3151,7 @@ class ProxyLogging: if await _skipped_for_unroutable_model( _callback, data, # pyright: ignore[reportUnknownArgumentType] # data is the untyped request dict + call_type, probe, caps, ): @@ -3226,6 +3231,7 @@ class ProxyLogging: if not await _skipped_for_unroutable_model( guardrail, data, # pyright: ignore[reportUnknownArgumentType] # data is the untyped request dict + call_type, probe, caps, ) diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 11e7335a309..be22efaed2d 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -3724,17 +3724,23 @@ class _ShippedGuardrail(CustomGuardrail): __module__ = "litellm.proxy.guardrails.guardrail_hooks.shipped_double" - def __init__(self, name: str = "shipped", run_in_parallel: bool = False, default_on: bool = True) -> None: + def __init__( + self, + name: str = "shipped", + run_in_parallel: bool = False, + default_on: bool = True, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + ) -> None: super().__init__( guardrail_name=name, - event_hook=GuardrailEventHooks.pre_call, + event_hook=event_hook, default_on=default_on, run_in_parallel=run_in_parallel, ) - self.seen: Final[list[object]] = [] + self.seen: tuple[object, ...] = () async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): - self.seen.append(data.get("messages")) + self.seen = (*self.seen, data.get("messages")) return data @@ -3751,10 +3757,10 @@ class _ShippedHook(CustomLogger): def __init__(self) -> None: super().__init__() - self.seen: Final[list[object]] = [] + self.seen: tuple[object, ...] = () async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): - self.seen.append(data.get("messages")) + self.seen = (*self.seen, data.get("messages")) return data @@ -3779,7 +3785,7 @@ def _unroutable_model_probe(): ) -async def _walk(monkeypatch, callbacks, model: str, probe) -> dict: +async def _walk(monkeypatch, callbacks, model: str, probe, call_type: str = "completion") -> dict: from litellm.caching.caching import DualCache from litellm.proxy.utils import ProxyLogging @@ -3788,7 +3794,7 @@ async def _walk(monkeypatch, callbacks, model: str, probe) -> dict: return await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), data={"model": model, "messages": [dict(message) for message in _SECRET_PROMPT]}, - call_type="completion", + call_type=call_type, routing_probe=probe, ) @@ -3804,9 +3810,9 @@ async def test_pre_call_hook_skips_a_shipped_guardrail_only_for_a_model_routing_ known_probe: Final = _unroutable_model_probe() await _walk(monkeypatch, [known_guardrail], "gpt-4o-mini", known_probe) - assert unknown_guardrail.seen == [] + assert unknown_guardrail.seen == () assert isinstance(unknown_probe.refusal, ProxyModelNotFoundError) - assert known_guardrail.seen == [_SECRET_PROMPT] + assert known_guardrail.seen == (_SECRET_PROMPT,) assert known_probe.refusal is None @@ -3815,7 +3821,7 @@ async def test_pre_call_hook_without_a_probe_runs_the_shipped_guardrail_for_an_u result: Final = await _walk(monkeypatch, [guardrail], "no-such-model", None) - assert guardrail.seen == [_SECRET_PROMPT] + assert guardrail.seen == (_SECRET_PROMPT,) assert result["messages"] == _SECRET_PROMPT @@ -3830,7 +3836,7 @@ async def test_pre_call_hook_runs_an_operator_guardrail_and_every_guardrail_befo await _walk(monkeypatch, [before, operator, after], "no-such-model", _unroutable_model_probe()) - assert (len(before.seen), len(operator.seen), after.seen) == (1, 1, []) + assert (len(before.seen), len(operator.seen), after.seen) == (1, 1, ()) async def test_pre_call_hook_runs_a_custom_code_guardrail_for_a_model_routing_refuses(monkeypatch): @@ -3858,7 +3864,7 @@ async def test_pre_call_hook_runs_every_guardrail_when_one_routes_on_sensitive_d await _walk(monkeypatch, [guardrail, sensitive_router], "no-such-model", _unroutable_model_probe()) - assert guardrail.seen == [_SECRET_PROMPT] + assert guardrail.seen == (_SECRET_PROMPT,) async def test_pre_call_hook_hides_the_prompt_from_hooks_after_a_skipped_guardrail(monkeypatch): @@ -3869,17 +3875,43 @@ async def test_pre_call_hook_hides_the_prompt_from_hooks_after_a_skipped_guardra result: Final = await _walk(monkeypatch, [_ShippedGuardrail(), hook], "no-such-model", _unroutable_model_probe()) redacted: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] - assert hook.seen == [redacted] + assert hook.seen == (redacted,) assert result["messages"] == redacted +@pytest.mark.parametrize( + ("event_hook", "guardrail_skipped"), + [ + pytest.param(GuardrailEventHooks.pre_mcp_call, True, id="mcp-guardrail"), + pytest.param(GuardrailEventHooks.pre_call, False, id="chat-only-guardrail"), + ], +) +async def test_pre_call_hook_checks_an_mcp_call_against_the_mcp_guardrail_event( + monkeypatch, event_hook, guardrail_skipped +): + from litellm.constants import REDACTED_BY_LITELLM + + guardrail: Final = _ShippedGuardrail(event_hook=event_hook) + hook: Final = _ShippedHook() + probe: Final = _unroutable_model_probe() + + await _walk(monkeypatch, [guardrail, hook], "no-such-model", probe, call_type="call_mcp_tool") + + redacted: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] + assert (guardrail.seen, hook.seen, probe.refusal is not None) == ( + (), + ((redacted if guardrail_skipped else _SECRET_PROMPT),), + guardrail_skipped, + ) + + async def test_pre_call_hook_leaves_the_request_alone_when_no_guardrail_applies(monkeypatch): hook: Final = _ShippedHook() probe: Final = _unroutable_model_probe() result: Final = await _walk(monkeypatch, [_ShippedGuardrail(default_on=False), hook], "no-such-model", probe) - assert hook.seen == [_SECRET_PROMPT] + assert hook.seen == (_SECRET_PROMPT,) assert result["messages"] == _SECRET_PROMPT assert probe.refusal is None