diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 79e3b0f762..2269548574 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,7 +1,7 @@ -from litellm._uuid import uuid from typing import Any, Dict, Iterable, List, Optional, Set, Union from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, @@ -30,7 +30,7 @@ def _prepare_mcp_server_data( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Convert model to dict - data_dict = data.model_dump() + data_dict = data.model_dump(exclude_none=True) # Ensure alias is always present in the dict (even if None) if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 9172568f30..e16b5192e5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -285,6 +285,13 @@ class MCPServerManager: self.registry[mcp_server.server_id] = new_server verbose_logger.debug(f"Added MCP Server: {name_for_prefix}") + def get_all_mcp_server_ids(self) -> List[str]: + """ + Get all MCP server IDs + """ + all_servers = self.get_registry().values() + return [server.server_id for server in all_servers] + async def get_allowed_mcp_servers( self, user_api_key_auth: Optional[UserAPIKeyAuth] = None ) -> List[str]: @@ -1176,6 +1183,13 @@ class MCPServerManager: } ) + ## filter out invalid servers + valid_server_ids = self.get_all_mcp_server_ids() + filtered_list_mcp_servers = [] + for server in list_mcp_servers: + if server.server_id not in valid_server_ids: + filtered_list_mcp_servers.append(server) + # Map servers to their teams and return with health data from typing import cast @@ -1198,9 +1212,7 @@ class MCPServerManager: else [] ), allowed_tools=( - server.allowed_tools - if server.allowed_tools is not None - else [] + server.allowed_tools if server.allowed_tools is not None else [] ), mcp_info=server.mcp_info, teams=cast( @@ -1212,7 +1224,7 @@ class MCPServerManager: args=getattr(server, "args", None) or [], env=getattr(server, "env", None) or {}, ) - for server in list_mcp_servers + for server in filtered_list_mcp_servers ] async def reload_servers_from_database(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 146f9434b8..a5eaa0f6f5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -909,6 +909,72 @@ class TestMCPServerManager: assert "tool_1" in tool_names assert "tool_2" in tool_names + @pytest.mark.asyncio + async def test_get_all_mcp_servers_filters_invalid_server_ids(self): + """Test that get_all_mcp_servers_with_health_and_teams filters out servers with invalid server_ids""" + manager = MCPServerManager() + + # Mock get_allowed_mcp_servers to return some server IDs + manager.get_allowed_mcp_servers = AsyncMock(return_value=["valid-server-1", "invalid-server-1"]) + + # Mock get_all_mcp_server_ids to return only valid server IDs + manager.get_all_mcp_server_ids = MagicMock(return_value=["valid-server-1", "another-valid-server"]) + + # Mock config_mcp_servers to be empty for this test + manager.config_mcp_servers = {} + + # Create mock servers - one valid, one invalid + valid_server = LiteLLM_MCPServerTable( + server_id="valid-server-1", + server_name="Valid Server", + alias="valid", + description="A valid server", + url="http://valid.com", + transport=MCPTransport.http, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + invalid_server = LiteLLM_MCPServerTable( + server_id="invalid-server-1", + server_name="Invalid Server", + alias="invalid", + description="An invalid server", + url="http://invalid.com", + transport=MCPTransport.http, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + # Mock the database call to return both servers + with patch('litellm.proxy._experimental.mcp_server.db.get_mcp_servers') as mock_get_servers: + mock_get_servers.return_value = [valid_server, invalid_server] + + # Mock prisma_client to not be None + with patch('litellm.proxy.proxy_server.prisma_client', new=MagicMock()): + # Mock _user_has_admin_view to return False (non-admin user) + with patch('litellm.proxy.management_endpoints.common_utils._user_has_admin_view') as mock_admin_view: + mock_admin_view.return_value = False + + # Create mock user auth + user_auth = MagicMock() + user_auth.user_id = "test-user" + + # Call the method + result = await manager.get_all_mcp_servers_with_health_and_teams( + user_api_key_auth=user_auth, + include_health=False + ) + + # Verify that only the valid server is returned + assert len(result) == 1 + assert result[0].server_id == "valid-server-1" + assert result[0].server_name == "Valid Server" + + # Verify invalid server is not in results + server_ids = [server.server_id for server in result] + assert "invalid-server-1" not in server_ids + if __name__ == "__main__": pytest.main([__file__])