diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql new file mode 100644 index 0000000000..8d3e02bd05 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allow_all_keys" BOOLEAN NOT NULL DEFAULT false; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index ea47b6ed03..e565135bbc 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + allow_all_keys Boolean @default(false) } // Generate Tokens for Proxy @@ -748,4 +749,4 @@ model LiteLLM_SkillsTable { created_by String? updated_at DateTime @default(now()) @updatedAt updated_by String? -} \ No newline at end of file +} diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d193fc27fb..a5ac966062 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -260,6 +260,7 @@ class MCPServerManager: allowed_params=server_config.get("allowed_params", None), access_groups=server_config.get("access_groups", None), static_headers=server_config.get("static_headers", None), + allow_all_keys=bool(server_config.get("allow_all_keys", False)), ) self.config_mcp_servers[server_id] = new_server @@ -549,6 +550,7 @@ class MCPServerManager: access_groups=getattr(mcp_server, "mcp_access_groups", None), allowed_tools=getattr(mcp_server, "allowed_tools", None), disallowed_tools=getattr(mcp_server, "disallowed_tools", None), + allow_all_keys=mcp_server.allow_all_keys, ) return new_server @@ -581,6 +583,14 @@ class MCPServerManager: all_servers = list(self.get_registry().values()) return {server.server_id for server in all_servers} + def get_allow_all_keys_server_ids(self) -> List[str]: + """Return server IDs that bypass per-key restrictions.""" + return [ + server.server_id + for server in self.get_registry().values() + if server.allow_all_keys + ] + async def get_allowed_mcp_servers( self, user_api_key_auth: Optional[UserAPIKeyAuth] = None ) -> List[str]: @@ -593,6 +603,8 @@ class MCPServerManager: if user_api_key_auth and _user_has_admin_view(user_api_key_auth): return list(self.get_registry().keys()) + allow_all_server_ids = self.get_allow_all_keys_server_ids() + try: allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers( user_api_key_auth @@ -600,14 +612,17 @@ class MCPServerManager: verbose_logger.debug( f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}" ) - if len(allowed_mcp_servers) == 0: + combined_servers = set(allowed_mcp_servers) + combined_servers.update(allow_all_server_ids) + + if len(combined_servers) == 0: verbose_logger.debug( "No allowed MCP Servers found for user api key auth." ) - return allowed_mcp_servers + return list(combined_servers) except Exception as e: verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}.") - return [] + return allow_all_server_ids async def get_tools_for_server(self, server_id: str) -> List[MCPTool]: """ @@ -2238,6 +2253,7 @@ class MCPServerManager: authorization_url=server.authorization_url, token_url=server.token_url, registration_url=server.registration_url, + allow_all_keys=server.allow_all_keys, ) async def get_all_mcp_servers_with_health_and_teams( @@ -2331,6 +2347,7 @@ class MCPServerManager: authorization_url=server.authorization_url, token_url=server.token_url, registration_url=server.registration_url, + allow_all_keys=server.allow_all_keys, ) list_mcp_servers.append(mcp_server_table) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7abfe7a96b..fa32f60c07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1037,6 +1037,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + allow_all_keys: bool = False @model_validator(mode="before") @classmethod @@ -1097,6 +1098,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + allow_all_keys: bool = False @model_validator(mode="before") @classmethod @@ -1149,6 +1151,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + allow_all_keys: bool = False class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -3821,4 +3824,4 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): - vector_store: LiteLLM_ManagedVectorStoresTable \ No newline at end of file + vector_store: LiteLLM_ManagedVectorStoresTable diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9111f53a51..a871a6637a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -212,6 +212,7 @@ if MCP_AVAILABLE: authorization_url=payload.authorization_url, token_url=payload.token_url, registration_url=payload.registration_url, + allow_all_keys=payload.allow_all_keys, ) def get_prisma_client_or_throw(message: str): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index ea47b6ed03..e565135bbc 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + allow_all_keys Boolean @default(false) } // Generate Tokens for Proxy @@ -748,4 +749,4 @@ model LiteLLM_SkillsTable { created_by String? updated_at DateTime @default(now()) @updatedAt updated_by String? -} \ No newline at end of file +} diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 869037546c..892a5146a7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -47,4 +47,5 @@ class MCPServer(BaseModel): args: Optional[List[str]] = None env: Optional[Dict[str, str]] = None access_groups: Optional[List[str]] = None + allow_all_keys: bool = False model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/schema.prisma b/schema.prisma index ea47b6ed03..e565135bbc 100644 --- a/schema.prisma +++ b/schema.prisma @@ -211,6 +211,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + allow_all_keys Boolean @default(false) } // Generate Tokens for Proxy @@ -748,4 +749,4 @@ model LiteLLM_SkillsTable { created_by String? updated_at DateTime @default(now()) @updatedAt updated_by String? -} \ No newline at end of file +} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index a1fbddec58..8062243dfd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1057,6 +1057,110 @@ async def test_list_tools_multiple_servers_prefixed_names(): assert names == ["jira-toolA", "zapier-toolA"] +@pytest.mark.asyncio +async def test_mcp_manager_allows_public_servers_without_permissions(): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + public_server = MCPServer( + server_id="public", + name="public", + transport=MCPTransport.http, + allow_all_keys=True, + ) + manager.registry = {public_server.server_id: public_server} + + with patch( + "litellm.proxy.management_endpoints.common_utils._user_has_admin_view", + return_value=False, + ), patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers", + AsyncMock(return_value=[]), + ): + allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth()) + + assert allowed == ["public"] + + +@pytest.mark.asyncio +async def test_mcp_manager_returns_public_when_permission_lookup_fails(): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + public_server = MCPServer( + server_id="public", + name="public", + transport=MCPTransport.http, + allow_all_keys=True, + ) + manager.registry = {public_server.server_id: public_server} + + with patch( + "litellm.proxy.management_endpoints.common_utils._user_has_admin_view", + return_value=False, + ), patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers", + AsyncMock(side_effect=Exception("boom")), + ): + allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth()) + + assert allowed == ["public"] + + +@pytest.mark.asyncio +async def test_mcp_manager_merges_public_and_restricted_servers(): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + public_server = MCPServer( + server_id="public", + name="public", + transport=MCPTransport.http, + allow_all_keys=True, + ) + scoped_server = MCPServer( + server_id="restricted", + name="restricted", + transport=MCPTransport.http, + ) + manager.registry = { + public_server.server_id: public_server, + scoped_server.server_id: scoped_server, + } + + with patch( + "litellm.proxy.management_endpoints.common_utils._user_has_admin_view", + return_value=False, + ), patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers", + AsyncMock(return_value=["restricted"]), + ): + allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth()) + + assert set(allowed) == {"public", "restricted"} + + @pytest.mark.asyncio async def test_call_mcp_tool_user_unauthorized_access(): """Test that a user cannot call a tool from a server they don't have access to"""