From 88fab5d2805c49c978d5ada229d7a584741ed6ec Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 7 Apr 2026 07:09:29 -0400 Subject: [PATCH] fix mcp tool guardrails --- .../mcp_server/mcp_server_manager.py | 9 +++++++- .../guardrail_hooks/tool_permission.py | 22 ++++++++++++++++++- litellm/proxy/utils.py | 13 ++++++++++- .../mcp/litellm_proxy_mcp_handler.py | 15 +++++++++++++ 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7b87e7e7e61..8d31b9fb98e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2174,7 +2174,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 6dd0288cb09..79b5bd7adbe 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) @@ -515,9 +517,27 @@ 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, + }, + ) + return data new_tools: Optional[List[ChatCompletionToolParam]] = data.get("tools") if new_tools is None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1e..75aaa5720c7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -546,6 +546,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], @@ -555,7 +566,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 b729cdb92f2..6ffcc43b088 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -849,6 +849,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, @@ -867,6 +872,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, @@ -885,6 +895,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,