From 371cabfebd072c261140904bcf40df08d33c4f25 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 17 Feb 2026 21:19:13 -0800 Subject: [PATCH] Add MCP Security guardrail to block unregistered MCP servers (#21429) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add MCP_SECURITY enum to SupportedGuardrailIntegrations * Add MCP security guardrail initializer * Add MCPSecurityGuardrail implementation * Add MCP Security policy template * Add Type filter to policy templates UI * Add unit tests for MCP security guardrail * fix(lint): remove unused Dict import from mcp_security_guardrail Co-Authored-By: Claude Sonnet 4.6 * Add French language support for EU AI Act Article 5 guardrail (#21427) * Add French language support for EU AI Act Article 5 template - Create eu_ai_act_article5_fr.yaml with comprehensive French keywords - Includes identifier words: concevoir, créer, développer, noter, classer, etc. - Includes block words: crédit social, comportement social, émotion des employés, etc. - Includes always-block keywords for explicit prohibited practices - Includes exceptions for research, compliance, and legitimate use cases - Catches circumvention attempts with phrase variations * Add comprehensive tests for French EU AI Act guardrail - Test 3 critical scenarios: blocked query, circumvention attempt, safe query - Test edge cases: case-insensitive, mixed language, research exceptions - All 7 tests passing - Validates both blocking and allowing behavior * Fix content filter to support conditional matching without inherit_from - Enable conditional matching when identifier_words + additional_block_words are present - Previously required inherit_from, but EU AI Act templates are self-contained - Fixes Greptile feedback: conditional matching now works as documented * Add pure conditional matching test for French guardrail - Test identifier + block word combinations not in always_block_keywords - Verifies conditional matching works independently - Addresses Greptile feedback about test coverage gap * Fix exception word bypass risk in French template - Replace short words (film, jeu, juste) with context-specific phrases - Prevents substring matching bypasses (e.g., enjeu matching jeu) - Add tests for bypass prevention and legitimate game context - Addresses Greptile security feedback * Make conditional match assertion more robust - Use getattr to safely access exception detail field - Check if detail is dict before calling .get() - Addresses Greptile feedback about brittle string assertion * Add French EU AI Act Article 5 policy template to registry - Add eu-ai-act-article5-fr template for French language support - Includes French description and guardrail info - Matches structure of English template * Address greptile review feedback (greploop iteration 1) - Use status_code=400 instead of 403 to match guardrail logging convention - Use prefix stripping instead of split('/')[-1] for robust server name extraction * remove French EU AI Act template from policy_templates.json --------- Co-authored-by: Julio Quinteros Pro Co-authored-by: Claude Sonnet 4.6 --- .../guardrail_hooks/mcp_security/__init__.py | 45 +++++ .../mcp_security/mcp_security_guardrail.py | 114 +++++++++++ litellm/types/guardrails.py | 1 + policy_templates.json | 32 +++ .../guardrail_hooks/test_mcp_security.py | 184 ++++++++++++++++++ .../components/policies/policy_templates.tsx | 59 ++++-- 6 files changed, 419 insertions(+), 16 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py new file mode 100644 index 0000000000..6efcb68a9e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py @@ -0,0 +1,45 @@ +from typing import TYPE_CHECKING, Literal, Optional, cast + +import litellm +from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( + MCPSecurityGuardrail, +) +from litellm.types.guardrails import SupportedGuardrailIntegrations + +if TYPE_CHECKING: + from litellm import Router + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + llm_router: Optional["Router"] = None, +): + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("MCP Security: guardrail_name is required") + + on_violation: Literal["block", "alert"] = cast( + Literal["block", "alert"], + getattr(litellm_params, "on_violation", "block"), + ) + + mcp_security_guardrail = MCPSecurityGuardrail( + guardrail_name=guardrail_name, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on or False, + on_violation=on_violation, + ) + + litellm.logging_callback_manager.add_litellm_callback(mcp_security_guardrail) + return mcp_security_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.MCP_SECURITY.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.MCP_SECURITY.value: MCPSecurityGuardrail, +} 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 0000000000..385143abc0 --- /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, 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=400, + 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[len(LITELLM_PROXY_MCP_SERVER_URL_PREFIX):] + 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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 7b95514637..3e12ccc6a4 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -64,6 +64,7 @@ class SupportedGuardrailIntegrations(Enum): ENKRYPTAI = "enkryptai" IBM_GUARDRAILS = "ibm_guardrails" LITELLM_CONTENT_FILTER = "litellm_content_filter" + MCP_SECURITY = "mcp_security" ONYX = "onyx" PROMPT_SECURITY = "prompt_security" GENERIC_GUARDRAIL_API = "generic_guardrail_api" diff --git a/policy_templates.json b/policy_templates.json index 29ecd5f983..5d2e8b72e7 100644 --- a/policy_templates.json +++ b/policy_templates.json @@ -999,5 +999,37 @@ ], "guardrails_remove": [] } + }, + { + "id": "mcp-security-unregistered-server-block", + "title": "MCP Security: Block Unregistered Servers", + "description": "Blocks requests that reference MCP servers not registered on this LiteLLM gateway. Prevents unauthorized tool access via unregistered MCP endpoints.", + "type": "MCP Security", + "region": "Global", + "icon": "ShieldCheckIcon", + "iconColor": "text-red-500", + "iconBg": "bg-red-50", + "guardrails": ["mcp-security-block"], + "complexity": "Low", + "guardrailDefinitions": [ + { + "guardrail_name": "mcp-security-block", + "litellm_params": { + "guardrail": "mcp_security", + "mode": "pre_call", + "default_on": true, + "on_violation": "block" + }, + "guardrail_info": { + "description": "Blocks requests referencing MCP servers not in the gateway registry" + } + } + ], + "templateData": { + "policy_name": "mcp-security-unregistered-server-block", + "description": "Blocks requests referencing MCP servers not registered on this gateway.", + "guardrails_add": ["mcp-security-block"], + "guardrails_remove": [] + } } ] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py new file mode 100644 index 0000000000..4444cd693f --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -0,0 +1,184 @@ +""" +Tests for MCP Security Guardrail. + +Validates that the guardrail blocks requests referencing unregistered MCP servers +and allows requests with only registered servers. Covers both /chat/completions +and /responses API paths (same pre_call_hook logic, different call_type). +""" + +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( + MCPSecurityGuardrail, +) +from litellm.types.guardrails import GuardrailEventHooks + + +@pytest.fixture +def guardrail(): + return MCPSecurityGuardrail( + guardrail_name="test-mcp-security", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + on_violation="block", + ) + + +class TestExtractMCPServerNames: + def test_extracts_litellm_proxy_mcp_servers(self): + tools = [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/github"}, + {"type": "function", "function": {"name": "get_weather"}}, + ] + names = MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) + assert names == {"zapier", "github"} + + def test_ignores_non_mcp_tools(self): + tools = [ + {"type": "function", "function": {"name": "get_weather"}}, + ] + names = MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) + assert names == set() + + def test_ignores_external_mcp_servers(self): + tools = [ + {"type": "mcp", "server_url": "https://external-server.com/mcp"}, + ] + names = MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) + assert names == set() + + def test_empty_tools(self): + names = MCPSecurityGuardrail._extract_mcp_server_names_from_tools([]) + assert names == set() + + +class TestMCPSecurityGuardrailPreCall: + @pytest.mark.asyncio + @patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) + async def test_blocks_unregistered_server_chat_completions( + self, mock_manager, guardrail + ): + """Simulates /chat/completions with an unregistered MCP server.""" + mock_manager.get_registry.return_value = {"zapier": MagicMock()} + + data = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, + ], + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["test-mcp-security"], + } + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type="acompletion", + ) + assert exc_info.value.status_code == 400 + assert "evil_server" in str(exc_info.value.detail) + + @pytest.mark.asyncio + @patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) + async def test_blocks_unregistered_server_responses_api( + self, mock_manager, guardrail + ): + """Simulates /responses with an unregistered MCP server.""" + mock_manager.get_registry.return_value = {"github": MagicMock()} + + data = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/unknown_server"}, + ], + "model": "gpt-4o", + "input": "What can you do?", + "guardrails": ["test-mcp-security"], + } + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type="aresponses", + ) + assert exc_info.value.status_code == 400 + assert "unknown_server" in str(exc_info.value.detail) + + @pytest.mark.asyncio + @patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) + async def test_allows_registered_servers(self, mock_manager, guardrail): + """All MCP servers are registered - request passes through.""" + mock_manager.get_registry.return_value = { + "zapier": MagicMock(), + "github": MagicMock(), + } + + data = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/github"}, + ], + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["test-mcp-security"], + } + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type="acompletion", + ) + assert result == data + + @pytest.mark.asyncio + async def test_passthrough_no_mcp_tools(self, guardrail): + """Request with no MCP tools passes through without checking registry.""" + data = { + "tools": [ + {"type": "function", "function": {"name": "get_weather"}}, + ], + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["test-mcp-security"], + } + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type="acompletion", + ) + assert result == data + + @pytest.mark.asyncio + async def test_passthrough_no_tools(self, guardrail): + """Request with no tools at all passes through.""" + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["test-mcp-security"], + } + + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type="acompletion", + ) + assert result == data diff --git a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx index f8b3f14124..478293f19e 100644 --- a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx +++ b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx @@ -117,16 +117,25 @@ const PolicyTemplates: React.FC = ({ onUseTemplate, access const [templates, setTemplates] = useState([]); const [isLoading, setIsLoading] = useState(false); const [selectedRegion, setSelectedRegion] = useState("All"); + const [selectedType, setSelectedType] = useState("All"); const availableRegions = useMemo(() => { const regions = new Set(templates.map(t => t.region || "Global")); return ["All", ...Array.from(regions).sort()]; }, [templates]); + const availableTypes = useMemo(() => { + const types = new Set(templates.map(t => t.type || "General")); + return ["All", ...Array.from(types).sort()]; + }, [templates]); + const filteredTemplates = useMemo(() => { - if (selectedRegion === "All") return templates; - return templates.filter(t => (t.region || "Global") === selectedRegion); - }, [templates, selectedRegion]); + return templates.filter(t => { + const regionMatch = selectedRegion === "All" || (t.region || "Global") === selectedRegion; + const typeMatch = selectedType === "All" || (t.type || "General") === selectedType; + return regionMatch && typeMatch; + }); + }, [templates, selectedRegion, selectedType]); useEffect(() => { const fetchTemplates = async () => { @@ -169,19 +178,37 @@ const PolicyTemplates: React.FC = ({ onUseTemplate, access -
- Region: - setSelectedRegion(e.target.value)} - buttonStyle="solid" - > - {availableRegions.map(region => ( - - {region} - - ))} - +
+
+ Region: + setSelectedRegion(e.target.value)} + buttonStyle="solid" + > + {availableRegions.map(region => ( + + {region} + + ))} + +
+ {availableTypes.length > 2 && ( +
+ Type: + setSelectedType(e.target.value)} + buttonStyle="solid" + > + {availableTypes.map(type => ( + + {type} + + ))} + +
+ )}