Add unit tests for MCP security guardrail

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 == 403
+        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 == 403
+        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