diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d1b49039e8e..f33317a10d5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2574,7 +2574,14 @@ class MCPServerManager: GuardrailRaisedException, HTTPException, ) as e: - # Re-raise guardrail exceptions to properly fail the MCP call + # Re-raise guardrail exceptions to properly fail the MCP call. + # Attach guardrail logging info so _execute_tool_calls can surface it + # in observability (UI logs, Langfuse, DataDog, etc.). + guardrail_info = synthetic_llm_data.get("metadata", {}).get( + "standard_logging_guardrail_information" + ) + if guardrail_info: + e._guardrail_logging_info = guardrail_info # type: ignore[attr-defined] verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}") raise e diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 37be832d350..fcbeb52737e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -56,6 +56,8 @@ class ToolPermissionGuardrail(CustomGuardrail): kwargs["supported_event_hooks"] = [ GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.during_mcp_call, ] super().__init__(**kwargs) @@ -636,9 +638,30 @@ class ToolPermissionGuardrail(CustomGuardrail): add_guardrail_to_applied_guardrails_header, ) - event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call + # pre_mcp_call path: synthetic data has mcp_tool_name but no tools list + mcp_tool_name: Optional[str] = data.get("mcp_tool_name") + event_type: GuardrailEventHooks = ( + GuardrailEventHooks.pre_mcp_call + if mcp_tool_name is not None + else GuardrailEventHooks.pre_call + ) if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data + if mcp_tool_name is not None: + is_allowed, _, message = self._check_tool_permission(mcp_tool_name) + if not is_allowed and message is not None: + verbose_proxy_logger.warning(f"Tool Permission Guardrail: {message}") + raise HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "detection_message": message, + }, + ) + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + return data new_tools = self._collect_request_tools(data) if not new_tools: @@ -676,6 +699,43 @@ class ToolPermissionGuardrail(CustomGuardrail): ) return data + @log_guardrail_information + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> None: + """ + Enforce tool permission rules during MCP tool execution (during_mcp_call). + + The data dict is the synthetic MCP payload built by _convert_mcp_to_llm_format, + which carries mcp_tool_name as the namespaced '{server}-{tool}' string. + """ + if self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.during_mcp_call + ) is not True: + return + + mcp_tool_name: Optional[str] = data.get("mcp_tool_name") + if mcp_tool_name is None: + return + + is_allowed, _, message = self._check_tool_permission(mcp_tool_name) + if not is_allowed and message is not None: + verbose_proxy_logger.warning(f"Tool Permission Guardrail: {message}") + raise HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "detection_message": message, + }, + ) + + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + @log_guardrail_information async def async_post_call_success_hook( self, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32c887f17b2..d6e6720c4d2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -608,6 +608,17 @@ class ProxyLogging: role="user", content=tool_call_content ) + # Build namespaced tool name ({server_name}-{tool_name}) so guardrail rules + # written against the prefixed form (e.g. "exa-.*") match correctly. + server_name = kwargs.get("server_name") or getattr( + request_obj, "server_name", None + ) + namespaced_tool_name = ( + f"{server_name}-{request_obj.tool_name}" + if server_name + else request_obj.tool_name + ) + # Create synthetic LLM data that guardrails can process synthetic_data = { "messages": [synthetic_message], @@ -617,7 +628,7 @@ class ProxyLogging: "user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"), "user_api_key_hash": kwargs.get("user_api_key_hash"), "user_api_key_request_route": kwargs.get("user_api_key_request_route"), - "mcp_tool_name": request_obj.tool_name, # Keep original for reference + "mcp_tool_name": namespaced_tool_name, # namespaced {server}-{tool} for rule matching "mcp_arguments": request_obj.arguments, # Keep original for reference # Raw Bearer token from the original HTTP request — allows guardrails # (e.g. MCPJWTSigner) to independently verify the caller's identity diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 94cff6922b5..482672b8780 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -853,6 +853,11 @@ class LiteLLM_Proxy_MCP_Handler: ) except BlockedPiiEntityError as e: + _guardrail_info = getattr(e, "_guardrail_logging_info", None) + if _guardrail_info: + logging_request_data.setdefault("metadata", {})[ + "standard_logging_guardrail_information" + ] = _guardrail_info await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( proxy_logging_obj=proxy_logging_obj, user_api_key_auth=user_api_key_auth, @@ -871,6 +876,11 @@ class LiteLLM_Proxy_MCP_Handler: } ) except GuardrailRaisedException as e: + _guardrail_info = getattr(e, "_guardrail_logging_info", None) + if _guardrail_info: + logging_request_data.setdefault("metadata", {})[ + "standard_logging_guardrail_information" + ] = _guardrail_info await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( proxy_logging_obj=proxy_logging_obj, user_api_key_auth=user_api_key_auth, @@ -889,6 +899,11 @@ class LiteLLM_Proxy_MCP_Handler: } ) except HTTPException as e: + _guardrail_info = getattr(e, "_guardrail_logging_info", None) + if _guardrail_info: + logging_request_data.setdefault("metadata", {})[ + "standard_logging_guardrail_information" + ] = _guardrail_info await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( proxy_logging_obj=proxy_logging_obj, user_api_key_auth=user_api_key_auth, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 716b4470d25..2dfb700460f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -742,6 +742,202 @@ class TestToolPermissionGuardrail: assert "Permission denied" in choice.message.content +class TestToolPermissionGuardrailMCPPreCall: + """Tests for the pre_mcp_call fast-path in async_pre_call_hook.""" + + def _make_guardrail( + self, + *, + decision: str, + pattern: str, + default_action: str = "allow", + default_on: bool = True, + ) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name="test-mcp-guardrail", + event_hook=["pre_mcp_call", "during_mcp_call"], + rules=[{"id": "rule_1", "tool_name": pattern, "decision": decision}], + default_action=default_action, + on_disallowed_action="block", + default_on=default_on, + ) + + @pytest.mark.asyncio + async def test_mcp_tool_name_denied_raises_http_exception(self): + """pre_mcp_call: denied tool raises HTTPException(400).""" + guardrail = self._make_guardrail(decision="deny", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa", "mcp_arguments": {"query": "test"}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_mcp_tool_name_allowed_passes(self): + """pre_mcp_call: allowed tool returns data without raising.""" + guardrail = self._make_guardrail(decision="allow", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa", "mcp_arguments": {"query": "test"}} + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + assert isinstance(result, dict) + + @pytest.mark.asyncio + async def test_mcp_namespaced_tool_name_matches_prefix_rule(self): + """Rule written as 'exa-.*' must match the namespaced form produced by + _convert_mcp_to_llm_format, i.e. '{server}-{tool}' = 'exa-web_search_exa'.""" + guardrail = self._make_guardrail(decision="deny", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa"} + + with pytest.raises(HTTPException): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + + @pytest.mark.asyncio + async def test_mcp_tool_name_not_matching_rule_uses_default_allow(self): + """When no rule matches, default_action='allow' should pass.""" + guardrail = self._make_guardrail( + decision="deny", pattern=r"exa-.*", default_action="allow" + ) + data = {"mcp_tool_name": "github-search_repos"} + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + assert isinstance(result, dict) + + @pytest.mark.asyncio + async def test_mcp_tool_name_not_matching_rule_uses_default_deny(self): + """When no rule matches, default_action='deny' should block.""" + guardrail = self._make_guardrail( + decision="allow", pattern=r"exa-.*", default_action="deny" + ) + data = {"mcp_tool_name": "github-search_repos"} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_no_mcp_tool_name_falls_through_to_tools_list(self): + """When mcp_tool_name is absent, hook falls through to the tools-list path.""" + guardrail = self._make_guardrail(decision="deny", pattern=r"exa-.*", default_on=False) + data = { + "tools": [{"type": "function", "function": {"name": "Bash"}}], + "guardrails": ["test-mcp-guardrail"], + } + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + assert isinstance(result, dict) + assert "tools" in result + + @pytest.mark.asyncio + async def test_mcp_allowed_tool_sets_applied_guardrails_header(self): + """Allowed MCP tool must call add_guardrail_to_applied_guardrails_header + so the guardrail appears in x-litellm-applied-guardrails.""" + guardrail = self._make_guardrail(decision="allow", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa"} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="call_mcp_tool", + ) + + applied = (data.get("metadata") or {}).get("applied_guardrails", []) + assert "test-mcp-guardrail" in applied + + def test_supported_event_hooks_includes_mcp(self): + """ToolPermissionGuardrail must declare pre_mcp_call and during_mcp_call + so that mode=['pre_mcp_call','during_mcp_call'] passes _validate_event_hook.""" + guardrail = self._make_guardrail(decision="deny", pattern=r".*") + assert GuardrailEventHooks.pre_mcp_call in (guardrail.supported_event_hooks or []) + assert GuardrailEventHooks.during_mcp_call in (guardrail.supported_event_hooks or []) + + @pytest.mark.asyncio + async def test_during_mcp_call_denied_raises_http_exception(self): + """during_mcp_call: denied tool raises HTTPException(400) via async_moderation_hook.""" + guardrail = self._make_guardrail(decision="deny", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa", "mcp_arguments": {"query": "test"}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="call_mcp_tool", + ) + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_during_mcp_call_allowed_passes(self): + """during_mcp_call: allowed tool does not raise.""" + guardrail = self._make_guardrail(decision="allow", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa", "mcp_arguments": {"query": "test"}} + + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="call_mcp_tool", + ) + + @pytest.mark.asyncio + async def test_during_mcp_call_allowed_sets_applied_guardrails_header(self): + """during_mcp_call: allowed tool must appear in applied-guardrails header.""" + guardrail = self._make_guardrail(decision="allow", pattern=r"exa-.*") + data = {"mcp_tool_name": "exa-web_search_exa"} + + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="call_mcp_tool", + ) + + applied = (data.get("metadata") or {}).get("applied_guardrails", []) + assert "test-mcp-guardrail" in applied + + @pytest.mark.asyncio + async def test_during_mcp_call_no_mcp_tool_name_is_noop(self): + """during_mcp_call: missing mcp_tool_name is a no-op (does not raise).""" + guardrail = self._make_guardrail(decision="deny", pattern=r".*") + data = {"mcp_arguments": {"query": "test"}} + + await guardrail.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="call_mcp_tool", + ) + + class TestToolPermissionGuardrailIntegration: """Integration tests for Tool Permission Guardrail"""