diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index dfc45417814..2e14170c114 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -2618,8 +2618,9 @@ async def _run_post_mcp_call_guardrails( user_api_key_dict=post_auth, ) from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.utils import mcp_guardrail_gate_data - base_request_data: Final[Mapping[str, object]] = ( + base_request_data: Final[Mapping[str, object]] = mcp_guardrail_gate_data( litellm_logging_obj.model_call_details if litellm_logging_obj is not None else request_data ) request_metadata: Final[object | None] = base_request_data.get("metadata") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a83e7d62616..58a4d8d69b1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1852,6 +1852,27 @@ def _failure_fields_to_lift( return MappingProxyType({key: value for key, value in _entries if value is not None}) +def mcp_guardrail_gate_data(request_data: Mapping[str, object]) -> Mapping[str, object]: + """``request_data`` reshaped so guardrail gating can read the request's metadata. + + MCP tool calls carry the request's ``Logging.model_call_details``, which nests + the request metadata under ``litellm_params``. ``CustomGuardrail.should_run_guardrail`` + only reads the top level, so without this lift the guardrail list a key, team + or policy resolved onto the request is invisible and only ``default_on`` + guardrails ever run, and ``opted_out_global_guardrails`` never takes one off. + A bucket the caller already set at the top level wins, because the MCP call + site that sets it has merged the guardrail auth's guardrails into it. + """ + params: Final = request_data.get("litellm_params") + if not isinstance(params, Mapping): + return request_data + nested: Final = cast("Mapping[str, object]", params) # cast-ok: litellm_params is always a str-keyed dict + lifted: Final = MappingProxyType({key: nested[key] for key in ("metadata", "litellm_metadata") if key in nested}) + if not lifted: + return request_data + return MappingProxyType({**lifted, **request_data}) + + @dataclass(frozen=True) class _CallbackCapabilities: """Cached per-hook capability flags derived from ``litellm.callbacks``. @@ -4605,15 +4626,13 @@ class ProxyLogging: verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook") return response + gate_data: Final = mcp_guardrail_gate_data(request_data) for callback in caps.resolved_callbacks: if not isinstance(callback, CustomGuardrail): continue if "apply_guardrail" not in type(callback).__dict__ or callback.use_native_lifecycle_hooks: continue - if ( - callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call) - is not True - ): + if callback.should_run_guardrail(data=gate_data, event_type=GuardrailEventHooks.post_mcp_call) is not True: continue response = await self._run_guardrail_with_metrics( callback, diff --git a/tests/unit/proxy/mcp_server/test_operations.py b/tests/unit/proxy/mcp_server/test_operations.py index c482f588521..a3353690046 100644 --- a/tests/unit/proxy/mcp_server/test_operations.py +++ b/tests/unit/proxy/mcp_server/test_operations.py @@ -1,11 +1,11 @@ import asyncio from collections.abc import Callable, Sequence from typing import Final, Literal -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from mcp.server.context import ServerRequestContext -from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult +from mcp.types import CallToolResult, GetPromptRequest, GetPromptRequestParams, GetPromptResult, TextContent from mcp.types import Tool as MCPTool import litellm @@ -2050,3 +2050,32 @@ async def test_discovery_optional_probes_use_remaining_listing_budget(delay, mon else: assert entered.is_set() assert cleaned.is_set() + + +@pytest.mark.asyncio +async def test_post_mcp_call_guardrails_keep_the_request_metadata_on_the_guardrail_auth_path(): + """Guardrails and opt-outs resolved onto the request must reach post_mcp_call gating. + + Protocol tool calls pass a guardrail auth that is a separate object from the + caller's, and the logging payload nests the request metadata under litellm_params. + """ + request_metadata: Final = { + "guardrails": ["policy-guardrail"], + "user_api_key_metadata": {"opted_out_global_guardrails": ["global-guardrail"]}, + } + logging_obj: Final = MagicMock(model_call_details={"litellm_params": {"metadata": request_metadata}}) + proxy_logging: Final = MagicMock(post_mcp_call_hook=AsyncMock(side_effect=lambda response, **_: response)) + result: Final = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging): + await operations._run_post_mcp_call_guardrails( + result=result, + litellm_logging_obj=logging_obj, + original_auth=UserAPIKeyAuth(api_key="sk-caller", user_id="u1"), + guardrail_auth=UserAPIKeyAuth(api_key="sk-caller", user_id="u1"), + request_data={}, + ) + + gated: Final = proxy_logging.post_mcp_call_hook.await_args.kwargs["request_data"]["metadata"] + assert gated["guardrails"] == ["policy-guardrail"] + assert gated["user_api_key_metadata"] == request_metadata["user_api_key_metadata"] diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index f74ff9b1fca..2801361eabe 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -2166,8 +2166,8 @@ class TestSendEmailStartTls: class _RecordingMCPGuardrail(CustomGuardrail): """Unified guardrail that masks every text it is handed.""" - def __init__(self, event_hook, masked_text="", raises=None): - super().__init__(guardrail_name="mcp-output-guardrail", event_hook=event_hook, default_on=True) + def __init__(self, event_hook, masked_text="", raises=None, default_on=True): + super().__init__(guardrail_name="mcp-output-guardrail", event_hook=event_hook, default_on=default_on) self.masked_text = masked_text self.raises = raises self.call_count = 0 @@ -2201,26 +2201,32 @@ class _NativeMCPGuardrail(CustomGuardrail): @pytest.fixture -def restore_callbacks(): - """Restore the process-wide callback state post_mcp_call_hook reads. +def install_callbacks(): + """Install callbacks on the process-wide list post_mcp_call_hook reads, and restore it afterwards. ProxyLogging caches callback capabilities keyed on id()s of litellm.callbacks, - so a restored-but-different list can collide with a stale entry after GC and - leak a has_guardrail verdict into unrelated tests in the same worker. + so a replaced list can collide with a stale entry after GC and leak a + has_guardrail verdict into unrelated tests in the same worker. The cache is + cleared on both install and restore for that reason. """ original = list(litellm.callbacks) - yield + + def install(*callbacks): + litellm.callbacks = list(callbacks) + ProxyLogging._callback_capabilities_cache.clear() + + yield install litellm.callbacks = original ProxyLogging._callback_capabilities_cache.clear() @pytest.mark.asyncio -async def test_post_mcp_call_hook_masks_tool_result(restore_callbacks): +async def test_post_mcp_call_hook_masks_tool_result(install_callbacks): """A post_mcp_call guardrail must see the tool result text and mask it in the returned result.""" from mcp.types import CallToolResult, TextContent guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) - litellm.callbacks = [guardrail] + install_callbacks(guardrail) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) @@ -2236,12 +2242,12 @@ async def test_post_mcp_call_hook_masks_tool_result(restore_callbacks): @pytest.mark.asyncio -async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(restore_callbacks): +async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(install_callbacks): """A guardrail not configured for post_mcp_call must not scan MCP tool results.""" from mcp.types import CallToolResult, TextContent guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_call) - litellm.callbacks = [guardrail] + install_callbacks(guardrail) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) @@ -2256,13 +2262,13 @@ async def test_post_mcp_call_hook_skips_guardrail_configured_for_other_hooks(res @pytest.mark.asyncio -async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(restore_callbacks): +async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(install_callbacks): """Guardrails that implement async_post_mcp_tool_call_hook are dispatched by the logging object, so this hook must not run them a second time.""" from mcp.types import CallToolResult, TextContent guardrail = _NativeMCPGuardrail() - litellm.callbacks = [guardrail] + install_callbacks(guardrail) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) @@ -2277,7 +2283,7 @@ async def test_post_mcp_call_hook_skips_guardrail_without_apply_guardrail(restor @pytest.mark.asyncio -async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks): +async def test_post_mcp_call_hook_propagates_guardrail_block(install_callbacks): """A guardrail rejecting the tool result must raise out of the hook.""" from mcp.types import CallToolResult, TextContent @@ -2287,7 +2293,7 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks): event_hook=GuardrailEventHooks.post_mcp_call, raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="mcp-output-guardrail"), ) - litellm.callbacks = [guardrail] + install_callbacks(guardrail) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) @@ -2299,6 +2305,158 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks): ) +def _mcp_logging_payload(metadata, metadata_key="metadata"): + """The request payload every MCP call site hands ``post_mcp_call_hook``. + + ``Logging.model_call_details`` nests the request metadata under + ``litellm_params``; see ``litellm.utils.function_setup``. The bucket is + named ``litellm_metadata`` on the routes that reserve ``metadata`` for the + provider, which is how MCP tools invoked from /responses arrive here. + """ + return { + "model": "MCP: echo", + "call_type": "call_mcp_tool", + "litellm_params": {"api_base": "", metadata_key: metadata}, + } + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_runs_guardrail_requested_through_request_metadata(install_callbacks): + """A guardrail attached by key, team or policy (not default_on) must mask MCP tool output. + + The MCP paths pass the logging payload, where the resolved guardrail list + sits under litellm_params.metadata, so gating on the top level alone left + masking working only for default_on guardrails. + """ + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=_mcp_logging_payload({"guardrails": [guardrail.guardrail_name]}), + user_api_key_dict=None, + ) + + assert guardrail.call_count == 1 + assert [item.text for item in returned.content] == [""] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_skips_guardrail_absent_from_request_metadata(install_callbacks): + """A non-default_on guardrail the request never asked for must leave tool output alone.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=_mcp_logging_payload({"guardrails": ["some-other-guardrail"]}), + user_api_key_dict=None, + ) + + assert guardrail.call_count == 0 + assert [item.text for item in returned.content] == ["jane@example.com"] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_runs_guardrail_requested_through_litellm_metadata(install_callbacks): + """The same lift must work for routes whose request metadata bucket is litellm_metadata.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=_mcp_logging_payload({"guardrails": [guardrail.guardrail_name]}, metadata_key="litellm_metadata"), + user_api_key_dict=None, + ) + + assert guardrail.call_count == 1 + assert [item.text for item in returned.content] == [""] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_honors_opt_out_in_request_metadata(install_callbacks): + """A key that opted out of a global guardrail must not have MCP tool output scanned by it.""" + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data=_mcp_logging_payload( + {"user_api_key_metadata": {"opted_out_global_guardrails": [guardrail.guardrail_name]}} + ), + user_api_key_dict=None, + ) + + assert guardrail.call_count == 0 + assert [item.text for item in returned.content] == ["jane@example.com"] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_runs_default_on_guardrail_when_payload_has_no_metadata_bucket(install_callbacks): + """A logging payload carrying no metadata bucket must leave the Default On path working. + + There is nothing to lift here, and that no-op must not change the verdict. + """ + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"model": "MCP: echo", "litellm_params": {"api_base": ""}}, + user_api_key_dict=None, + ) + + assert guardrail.call_count == 1 + assert [item.text for item in returned.content] == [""] + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_keeps_top_level_metadata_the_caller_shaped(install_callbacks): + """A top-level metadata bucket the call site already shaped must drive gating over the nested one. + + The MCP call site merges the guardrail auth's guardrails, such as a granting + team's, into the top level, so lifting litellm_params over it would drop them. + """ + from mcp.types import CallToolResult, TextContent + + guardrail = _RecordingMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call, default_on=False) + install_callbacks(guardrail) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={ + "metadata": {"guardrails": [guardrail.guardrail_name]}, + **_mcp_logging_payload({"guardrails": ["some-other-guardrail"]}), + }, + user_api_key_dict=None, + ) + + assert guardrail.call_count == 1 + assert [item.text for item in returned.content] == [""] + + @pytest.mark.asyncio async def test_prisma_health_check_failure_names_itself_at_operator_visible_level(caplog): """A failing DB health check has to name the check that failed, at a level @@ -2406,7 +2564,7 @@ async def test_update_data_key_branch_stamps_settings_updated_at(): @pytest.mark.asyncio -async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks): +async def test_post_mcp_call_hook_skips_opted_out_guardrail(install_callbacks): """A guardrail that keeps its native lifecycle hooks must not have MCP tool results scanned through the unified path, even though it implements apply_guardrail.""" from mcp.types import CallToolResult, TextContent @@ -2421,7 +2579,7 @@ async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks): return await super().apply_guardrail(inputs, request_data, input_type, **kwargs) guardrail = _OptedOutMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) - litellm.callbacks = [guardrail] + install_callbacks(guardrail) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)