From b3a214bce707e08141ea9784958e8a847440264d Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Thu, 17 Jul 2025 02:27:43 +0530 Subject: [PATCH] [MCP Gateway] Allow MCP access groups to be added via the config LIT-312 (#12654) * allow mcp access groups to be added via the config * fix code and doc quality * fix mypy * create helpers: --- .../mcp_server/auth/user_api_key_auth_mcp.py | 65 +++++++++++++------ .../mcp_server/mcp_server_manager.py | 1 + .../mcp_management_endpoints.py | 45 ++++++------- .../types/mcp_server/mcp_server_manager.py | 1 + tests/mcp_tests/test_mcp_server.py | 36 ++++++++++ 5 files changed, 105 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 54428fafc23..708adf66db5 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -273,34 +273,61 @@ class MCPRequestHandler: all_servers = direct_mcp_servers + access_group_servers return list(set(all_servers)) + @staticmethod + def _get_config_server_ids_for_access_groups(config_mcp_servers, access_groups: List[str]) -> set: + """ + Helper to get server_ids from config-loaded servers that match any of the given access groups. + """ + server_ids = set() + for server_id, server in config_mcp_servers.items(): + if server.access_groups: + if any(group in server.access_groups for group in access_groups): + server_ids.add(server_id) + return server_ids + + @staticmethod + async def _get_db_server_ids_for_access_groups(prisma_client, access_groups: List[str]) -> set: + """ + Helper to get server_ids from DB servers that match any of the given access groups. + """ + server_ids = set() + if access_groups and prisma_client is not None: + try: + mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + where={ + "mcp_access_groups": { + "hasSome": access_groups + } + } + ) + for server in mcp_servers: + server_ids.add(server.server_id) + except Exception as e: + verbose_logger.debug(f"Error getting MCP servers from access groups: {e}") + return server_ids + @staticmethod async def _get_mcp_servers_from_access_groups( access_groups: List[str] ) -> List[str]: """ - Resolve MCP access groups to server IDs by querying the MCP server table + Resolve MCP access groups to server IDs by querying BOTH the MCP server table (DB) AND config-loaded servers """ from litellm.proxy.proxy_server import prisma_client + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - if not access_groups or prisma_client is None: - return [] + # Use the new helper for config-loaded servers + server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups( + global_mcp_server_manager.config_mcp_servers, access_groups + ) - try: - # Find all MCP servers that have any of the specified access groups - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( - where={ - "mcp_access_groups": { - "hasSome": access_groups - } - } - ) - - # Extract server IDs - server_ids = [server.server_id for server in mcp_servers] - return server_ids - except Exception as e: - verbose_logger.debug(f"Error getting MCP servers from access groups: {e}") - return [] + # Use the new helper for DB servers + db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups( + prisma_client, access_groups + ) + server_ids.update(db_server_ids) + + return list(server_ids) @staticmethod async def get_mcp_access_groups( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 012ee1569de..fedf9070961 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -132,6 +132,7 @@ class MCPServerManager: spec_version=server_config.get("spec_version", MCPSpecVersion.mar_2025), auth_type=server_config.get("auth_type", None), mcp_info=mcp_info, + access_groups=server_config.get("access_groups", None), ) self.config_mcp_servers[server_id] = new_server verbose_logger.debug( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 257c48d1ad7..f3a06e8723b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -123,34 +123,31 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Get all available MCP access groups from the database + Get all available MCP access groups from the database AND config """ from litellm.proxy.proxy_server import prisma_client + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - if prisma_client is None: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={ - "error": "Database not connected. Connect a database to your proxy" - }, - ) + access_groups = set() - try: - # Get all MCP servers and extract their access groups - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() - # Extract all unique access groups - access_groups = set() - for server in mcp_servers: - if server.mcp_access_groups: - access_groups.update(server.mcp_access_groups) - - # Convert to sorted list - access_groups_list = sorted(list(access_groups)) - - return {"access_groups": access_groups_list} - except Exception as e: - verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}") - return {"access_groups": []} + # Get from config-loaded servers + for server in global_mcp_server_manager.config_mcp_servers.values(): + if server.access_groups: + access_groups.update(server.access_groups) + + # Get from DB + if prisma_client is not None: + try: + mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + for server in mcp_servers: + if hasattr(server, 'mcp_access_groups') and server.mcp_access_groups: + access_groups.update(server.mcp_access_groups) + except Exception as e: + verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}") + + # Convert to sorted list + access_groups_list = sorted(list(access_groups)) + return {"access_groups": access_groups_list} ## FastAPI Routes @router.get( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 8ed95b83949..cc822826e95 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -27,4 +27,5 @@ class MCPServer(BaseModel): command: Optional[str] = None args: Optional[List[str]] = None env: Optional[Dict[str, str]] = None + access_groups: Optional[List[str]] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 6021c9c88fe..e5268ab82f7 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -719,3 +719,39 @@ async def test_get_tools_from_mcp_servers(): pytest.fail(f"Unexpected error in tests: {str(e)}") +def test_mcp_server_manager_access_groups_from_config(): + """ + Test that access_groups are loaded from config and can be resolved. + """ + test_manager = MCPServerManager() + test_manager.load_servers_from_config({ + "config_server": { + "url": "https://config-mcp-server.com/mcp", + "transport": MCPTransport.http, + "access_groups": ["group-a", "group-b"] + }, + "other_server": { + "url": "https://other-mcp-server.com/mcp", + "transport": MCPTransport.http, + "access_groups": ["group-b", "group-c"] + } + }) + # Check that access_groups are loaded + config_server = next((s for s in test_manager.config_mcp_servers.values() if s.name == "config_server"), None) + assert config_server is not None + assert set(config_server.access_groups) == {"group-a", "group-b"} + # Check that the lookup logic finds the correct server ids + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + # Patch global_mcp_server_manager for this test + import litellm.proxy._experimental.mcp_server.mcp_server_manager as mcp_server_manager_mod + mcp_server_manager_mod.global_mcp_server_manager = test_manager + # Should find config_server for group-a, both for group-b, other_server for group-c + import asyncio + server_ids_a = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-a"])) + server_ids_b = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-b"])) + server_ids_c = asyncio.run(MCPRequestHandler._get_mcp_servers_from_access_groups(["group-c"])) + assert any(config_server.server_id == sid for sid in server_ids_a) + assert set(server_ids_b) == set([s.server_id for s in test_manager.config_mcp_servers.values() if "group-b" in s.access_groups]) + assert any(s.name == "other_server" and s.server_id in server_ids_c for s in test_manager.config_mcp_servers.values()) + +