From 21d7870cf52b8e901ad627315e4ebc996ba6c1a1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Feb 2026 16:08:08 -0800 Subject: [PATCH] Add MCPSecurityGuardrail implementation --- .../mcp_security/mcp_security_guardrail.py | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py new file mode 100644 index 00000000000..2740f2caecd --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -0,0 +1,114 @@ +""" +MCP Security Guardrail for LiteLLM. + +Validates that MCP servers referenced in request tools are registered +on the LiteLLM gateway. Blocks or alerts when unregistered servers are found. +""" + +from typing import Any, Dict, List, Literal, Optional, Set, Union + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX, +) +from litellm.types.guardrails import GuardrailEventHooks + + +class MCPSecurityGuardrail(CustomGuardrail): + def __init__( + self, + on_violation: Literal["block", "alert"] = "block", + **kwargs, + ): + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call] + super().__init__(**kwargs) + self.on_violation = on_violation + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: Any, + data: dict, + call_type: str, + ) -> Optional[Union[Exception, str, dict]]: + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + is not True + ): + return data + + unregistered = self._find_unregistered_mcp_servers(data) + if not unregistered: + return data + + message = ( + f"MCP Security: request references unregistered MCP server(s): " + f"{', '.join(sorted(unregistered))}. " + f"Only servers registered on this gateway are allowed." + ) + + if self.on_violation == "block": + raise HTTPException( + status_code=403, + detail={ + "error": "Violated guardrail policy", + "guardrail": "mcp_security", + "unregistered_servers": sorted(unregistered), + "detection_message": message, + }, + ) + else: + verbose_proxy_logger.warning(message) + + return data + + @staticmethod + def _extract_mcp_server_names_from_tools(tools: List[dict]) -> Set[str]: + """Extract MCP server names from tools with type=mcp and litellm_proxy server_url.""" + server_names: Set[str] = set() + for tool in tools: + if not isinstance(tool, dict): + continue + if tool.get("type") != "mcp": + continue + server_url = tool.get("server_url", "") + if not isinstance(server_url, str): + continue + if server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX): + name = server_url.split("/")[-1] + if name: + server_names.add(name) + return server_names + + @staticmethod + def _find_unregistered_mcp_servers(data: dict) -> Set[str]: + """Check tools in data against the MCP server registry. Returns set of unregistered server names.""" + tools = data.get("tools") + if not tools or not isinstance(tools, list): + return set() + + requested_servers = ( + MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) + ) + if not requested_servers: + return set() + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + registry = global_mcp_server_manager.get_registry() + registered_names = set(registry.keys()) + + return requested_servers - registered_names