mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 08:21:53 +00:00
Add MCP Security guardrail to block unregistered MCP servers (#21429)
* 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 <noreply@anthropic.com> * 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 <jquinter@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Ishaan Jaff
co-authored by
Julio Quinteros Pro
Claude Sonnet 4.6
Julio Quinteros Pro
Claude Sonnet 4.6
parent
5b5306a540
commit
371cabfebd
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -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"
|
||||
|
||||
@@ -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": []
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -117,16 +117,25 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, access
|
||||
const [templates, setTemplates] = useState<any[]>([]);
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
const [selectedRegion, setSelectedRegion] = useState<string>("All");
|
||||
const [selectedType, setSelectedType] = useState<string>("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<PolicyTemplatesProps> = ({ onUseTemplate, access
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-3 mb-4">
|
||||
<span className="text-sm font-medium text-gray-700">Region:</span>
|
||||
<Radio.Group
|
||||
value={selectedRegion}
|
||||
onChange={(e) => setSelectedRegion(e.target.value)}
|
||||
buttonStyle="solid"
|
||||
>
|
||||
{availableRegions.map(region => (
|
||||
<Radio.Button key={region} value={region}>
|
||||
{region}
|
||||
</Radio.Button>
|
||||
))}
|
||||
</Radio.Group>
|
||||
<div className="flex items-center gap-6 mb-4">
|
||||
<div className="flex items-center gap-3">
|
||||
<span className="text-sm font-medium text-gray-700">Region:</span>
|
||||
<Radio.Group
|
||||
value={selectedRegion}
|
||||
onChange={(e) => setSelectedRegion(e.target.value)}
|
||||
buttonStyle="solid"
|
||||
>
|
||||
{availableRegions.map(region => (
|
||||
<Radio.Button key={region} value={region}>
|
||||
{region}
|
||||
</Radio.Button>
|
||||
))}
|
||||
</Radio.Group>
|
||||
</div>
|
||||
{availableTypes.length > 2 && (
|
||||
<div className="flex items-center gap-3">
|
||||
<span className="text-sm font-medium text-gray-700">Type:</span>
|
||||
<Radio.Group
|
||||
value={selectedType}
|
||||
onChange={(e) => setSelectedType(e.target.value)}
|
||||
buttonStyle="solid"
|
||||
>
|
||||
{availableTypes.map(type => (
|
||||
<Radio.Button key={type} value={type}>
|
||||
{type}
|
||||
</Radio.Button>
|
||||
))}
|
||||
</Radio.Group>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className="grid grid-cols-1 md:grid-cols-2 xl:grid-cols-3 gap-6">
|
||||
|
||||
Reference in New Issue
Block a user