From c25ceb5596d7b95b201014d3073fa792cf8256f4 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 26 Dec 2025 16:55:51 +0900 Subject: [PATCH] refactor: MCP health check --- .../mcp_server/mcp_server_manager.py | 210 ++++-------- .../mcp_management_endpoints.py | 122 +------ .../mcp_server/test_mcp_server_manager.py | 270 ++++++++------- .../test_mcp_management_endpoints.py | 317 +----------------- 4 files changed, 226 insertions(+), 693 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a579080367..1171eef75d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2127,7 +2127,7 @@ class MCPServerManager: async def health_check_server( self, server_id: str, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: + ) -> LiteLLM_MCPServerTable: """ Perform a health check on a specific MCP server. @@ -2138,96 +2138,82 @@ class MCPServerManager: Returns: Dict containing health check results """ - import time from datetime import datetime server = self.get_mcp_server_by_id(server_id) if not server: - return { - "server_id": server_id, - "server_name": None, - "status": "unknown", - "error": "Server not found", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": None, - } - - start_time = time.time() - try: - # Try to get tools from the server as a health check - tools = await self._get_tools_from_server(server, mcp_auth_header) - response_time = (time.time() - start_time) * 1000 - - return { - "server_id": server_id, - "server_name": server.name, - "status": "healthy", - "tools_count": len(tools), - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": None, - } - except Exception as e: - response_time = (time.time() - start_time) * 1000 - error_message = str(e) - - return { - "server_id": server_id, - "server_name": server.name, - "status": "unhealthy", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": error_message, - } - - async def health_check_all_servers( - self, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers. - - Args: - mcp_auth_header: Optional authentication header for the MCP servers - - Returns: - Dict containing health check results for all servers - """ - all_servers = self.get_registry() - results = {} - - for server_id, server in all_servers.items(): - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + verbose_logger.warning(f"MCP Server {server_id} not found") + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name=None, + transport=MCPTransport.http, # Default transport for not found servers + status="unknown", + health_check_error="Server not found", + last_health_check=datetime.now(), ) - return results + status = "unknown" + health_check_error = None - async def health_check_allowed_servers( - self, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers that the user has access to. + # Check if we should skip health check based on auth configuration + should_skip_health_check = False - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional authentication header for the MCP servers + # Skip if auth_type is oauth2 + if server.auth_type == MCPAuth.oauth2: + should_skip_health_check = True + # Skip if auth_type is not none and authentication_token is missing + elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: + should_skip_health_check = True - Returns: - Dict containing health check results for accessible servers - """ - # Get allowed servers for the user - allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + if not should_skip_health_check: + extra_headers = {} + if server.static_headers: + extra_headers.update(server.static_headers) - # Perform health checks on allowed servers - results = {} - for server_id in allowed_server_ids: - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + client = self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, ) - return results + try: + async def _noop(session): + return "ok" + + await client.run_with_session(_noop) + status = "healthy" + except Exception as e: + health_check_error = str(e) + status = "unhealthy" + + return LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + alias=server.alias, + description=( + server.mcp_info.get("description") + if server.mcp_info + else None + ), + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + created_at=datetime.now(), + updated_at=datetime.now(), + teams=[], + mcp_access_groups=server.access_groups or [], + allowed_tools=server.allowed_tools or [], + extra_headers=server.extra_headers or [], + mcp_info=server.mcp_info, + static_headers=server.static_headers, + status=status, + last_health_check=datetime.now(), + health_check_error=health_check_error, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + ) async def get_all_mcp_servers_with_health_and_teams( self, @@ -2248,72 +2234,8 @@ class MCPServerManager: # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) - async def _check_server_health(server_id: str) -> Optional[LiteLLM_MCPServerTable]: - """Helper function to check health of a single server""" - server = self.get_mcp_server_by_id(server_id) - if server is None: - verbose_logger.warning(f"MCP Server {server_id} not found") - return None - - status = "unknown" - health_check_error = None - - # Check if we should skip health check based on auth configuration - should_skip_health_check = False - - # Skip if auth_type is oauth2 - if server.auth_type == MCPAuth.oauth2: - should_skip_health_check = True - # Skip if auth_type is not none and authentication_token is missing - elif server.auth_type and server.auth_type != MCPAuth.none and not server.authentication_token: - should_skip_health_check = True - - if not should_skip_health_check: - extra_headers = {} - if server.static_headers: - extra_headers.update(server.static_headers) - - client = self._create_mcp_client( - server=server, - mcp_auth_header=None, - extra_headers=extra_headers, - stdio_env=None, - ) - - try: - async def _noop(session): - return "ok" - - await client.run_with_session(_noop) - status = "healthy" - except Exception as e: - health_check_error = str(e) - status = "unhealthy" - - return LiteLLM_MCPServerTable( - **{ - **server.model_dump(), - "created_at": datetime.datetime.now(), - "updated_at": datetime.datetime.now(), - "description": ( - server.mcp_info.get("description") - if server.mcp_info - else None - ), - "allowed_tools": server.allowed_tools or [], - "mcp_info": server.mcp_info, - "mcp_access_groups": server.access_groups or [], - "extra_headers": server.extra_headers or [], - "command": getattr(server, "command", None), - "args": getattr(server, "args", None) or [], - "env": getattr(server, "env", None) or {}, - "status": status, - "health_check_error": health_check_error, - } - ) - # Run health checks concurrently - tasks = [_check_server_health(server_id) for server_id in allowed_server_ids] + tasks = [self.health_check_server(server_id) for server_id in allowed_server_ids] results = await asyncio.gather(*tasks) # Filter out None results (servers that were not found) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f0eddcc868..a77c7c8754 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -296,116 +296,6 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} - @router.get( - "/server/{server_id}/health", - description="Perform health check on a specific MCP server", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_mcp_server( - server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform a health check on the MCP server specified by the `server_id` - Parameters: - - server_id: str - Required. The unique identifier of the mcp server to health check. - ``` - curl --location 'http://localhost:4000/v1/mcp/server/{server_id}/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Check if server exists and user has access - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) - - # check to see if server exists for all users - mcp_server = await get_mcp_server(prisma_client, server_id) - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server with id {server_id} not found"}, - ) - - # Implement authz restriction from requested user - if not _user_has_admin_view(user_api_key_dict): - # Perform authz check to filter the mcp servers user has access to - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - exists = does_mcp_server_exist(mcp_server_records, server_id) - - if not exists: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": f"User does not have permission to access mcp server with id {server_id}. You can only access mcp servers that you have access to." - }, - ) - - # Perform health check using server manager - try: - health_result = await global_mcp_server_manager.health_check_server( - server_id - ) - return health_result - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health check on MCP server {server_id}: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health check: {str(e)}"}, - ) - - @router.get( - "/server/health", - description="Perform health check on all accessible MCP servers", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_all_mcp_servers( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform health checks on all MCP servers accessible to the user - ``` - curl --location 'http://localhost:4000/v1/mcp/server/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Use server manager to get health checks for allowed servers - try: - all_health_results = ( - await global_mcp_server_manager.health_check_allowed_servers( - user_api_key_auth=user_api_key_dict - ) - ) - - return { - "total_servers": len(all_health_results), - "healthy_count": len( - [r for r in all_health_results.values() if r["status"] == "healthy"] - ), - "unhealthy_count": len( - [ - r - for r in all_health_results.values() - if r["status"] == "unhealthy" - ] - ), - "unknown_count": len( - [r for r in all_health_results.values() if r["status"] == "unknown"] - ), - "servers": all_health_results, - } - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health checks on MCP servers: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health checks: {str(e)}"}, - ) ## FastAPI Routes @router.get( @@ -484,15 +374,9 @@ if MCP_AVAILABLE: server_id ) # Update the server object with health check results - mcp_server.status = health_result.get("status", "unknown") - mcp_server.last_health_check = ( - datetime.fromisoformat( - health_result.get("last_health_check", datetime.now().isoformat()) - ) - if health_result.get("last_health_check") - else None - ) - mcp_server.health_check_error = health_result.get("error") + mcp_server.status = health_result.status if health_result.status else "unknown" + mcp_server.last_health_check = health_result.last_health_check + mcp_server.health_check_error = health_result.health_check_error except Exception as e: verbose_proxy_logger.debug( f"Error performing health check on server {server_id}: {e}" 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 ff016a1a13..37a9344151 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 @@ -641,37 +641,31 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock successful _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - tool1 = MagicMock() - tool1.name = "tool1" - tool2 = MagicMock() - tool2.name = "tool2" - return [tool1, tool2] - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock successful client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") - # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 2 - assert result["error"] is None - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + # Verify results - result is now LiteLLM_MCPServerTable + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_unhealthy(self): @@ -679,32 +673,33 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock failed _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - raise Exception("Connection timeout") - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock failed client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock( + side_effect=Exception("Connection timeout") + ) + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "unhealthy" - assert result["error"] == "Connection timeout" - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "unhealthy" + assert result.health_check_error == "Connection timeout" + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_not_found(self): @@ -718,104 +713,121 @@ class TestMCPServerManager: result = await manager.health_check_server("non-existent-server") # Verify results - assert result["server_id"] == "non-existent-server" - assert result["status"] == "unknown" - assert result["error"] == "Server not found" - assert result["response_time_ms"] is None - assert "last_health_check" in result + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "non-existent-server" + assert result.server_name is None + assert result.status == "unknown" + assert result.health_check_error == "Server not found" + assert result.last_health_check is not None @pytest.mark.asyncio - async def test_health_check_all_servers(self): - """Test health check for all servers""" + async def test_health_check_server_oauth2_skips_check(self): + """Test that health check is skipped for OAuth2 servers and returns unknown status""" manager = MCPServerManager() - # Mock servers - server1 = MagicMock() - server1.server_id = "server1" - server1.name = "server1" - - server2 = MagicMock() - server2.server_id = "server2" - server2.name = "server2" - - # Mock registry - manager.registry = {"server1": server1, "server2": server2} - - # Mock get_mcp_server_by_id - def mock_get_server_by_id(server_id): - if server_id == "server1": - return server1 - elif server_id == "server2": - return server2 - return None - - manager.get_mcp_server_by_id = mock_get_server_by_id - - # Mock _get_tools_from_server with different results - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - if server.server_id == "server1": - tool = MagicMock() - tool.name = "tool1" - return [tool] - elif server.server_id == "server2": - raise Exception("Connection failed") - return [] - - manager._get_tools_from_server = mock_get_tools_from_server - - # Perform health check for all servers - result = await manager.health_check_all_servers() - - # Verify results - assert len(result) == 2 - assert "server1" in result - assert "server2" in result - - # Check server1 (healthy) - assert result["server1"]["status"] == "healthy" - assert result["server1"]["tools_count"] == 1 - assert result["server1"]["error"] is None - - # Check server2 (unhealthy) - assert result["server2"]["status"] == "unhealthy" - assert result["server2"]["error"] == "Connection failed" - - @pytest.mark.asyncio - async def test_health_check_server_with_auth_header(self): - """Test health check with authentication header""" - manager = MCPServerManager() - - # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + # Mock OAuth2 server + server = MCPServer( + server_id="oauth2-server", + name="oauth2-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth2-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock _get_tools_from_server to verify auth header is passed - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - assert mcp_auth_header == "test-token" - tool = MagicMock() - tool.name = "tool1" - return [tool] + # _create_mcp_client should not be called for OAuth2 servers + manager._create_mcp_client = MagicMock() - manager._get_tools_from_server = mock_get_tools_from_server + # Perform health check + result = await manager.health_check_server("oauth2-server") - # Perform health check with auth header - result = await manager.health_check_server("test-server", "test-token") + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 1 + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "oauth2-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_no_token_skips_check(self): + """Test that health check is skipped when auth_type is set but authentication_token is missing""" + manager = MCPServerManager() + + # Mock server with auth_type but no authentication_token + server = MCPServer( + server_id="no-token-server", + name="no-token-server", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=None, # No token + url="http://no-token-server.com", + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called + manager._create_mcp_client = MagicMock() + + # Perform health check + result = await manager.health_check_server("no-token-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "no-token-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_with_static_headers(self): + """Test health check with static headers configured""" + manager = MCPServerManager() + + # Mock server with static_headers + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + static_headers={"X-Custom-Header": "custom-value"}, + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + + # Capture the extra_headers passed to _create_mcp_client + captured_extra_headers = None + + def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env): + nonlocal captured_extra_headers + captured_extra_headers = extra_headers + return mock_client + + manager._create_mcp_client = MagicMock(side_effect=capture_create_mcp_client) + + # Perform health check + result = await manager.health_check_server("test-server") + + # Verify static headers were passed + assert captured_extra_headers == {"X-Custom-Header": "custom-value"} + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 61342e8025..8d5b3dcedc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -486,11 +486,14 @@ class TestListMCPServers: mock_server.credentials = {"auth_value": "top-secret"} mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", alias="Server 1" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -531,11 +534,14 @@ class TestListMCPServers: delattr(mock_server, "credentials") mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-2", alias="Server 2" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -568,296 +574,6 @@ class TestListMCPServers: assert result.status == "healthy" -class TestMCPHealthCheckEndpoints: - """Test MCP health check endpoints""" - - @pytest.mark.asyncio - async def test_health_check_mcp_server_success(self): - """Test successful health check for a specific MCP server""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_server = AsyncMock( - return_value={ - "server_id": "test-server", - "server_name": "Test Server", - "status": "healthy", - "tools_count": 3, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 150.5, - "error": None, - } - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - result = await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["server_id"] == "test-server" - assert result["server_name"] == "Test Server" - assert result["status"] == "healthy" - assert result["tools_count"] == 3 - assert result["response_time_ms"] == 150.5 - assert result["error"] is None - - @pytest.mark.asyncio - async def test_health_check_mcp_server_not_found(self): - """Test health check for a server that doesn't exist""" - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="non-existent-server", user_api_key_dict=mock_user_auth - ) - - assert "not found" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_mcp_server_unauthorized(self): - """Test health check for a server user doesn't have access to""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER # Non-admin user - ) - - # Mock user doesn't have access to this server - mock_user_servers = [] - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", - return_value=mock_user_servers, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - assert "permission" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_all_mcp_servers(self): - """Test health check for all accessible MCP servers""" - # Mock team records - team_records = [ - generate_mock_team_record( - team_id="team1", - team_alias="Team 1", - organization_id="org1", - mcp_servers=["server1", "server2"], - ) - ] - - # Mock DB servers - db_servers = [ - generate_mock_mcp_server_db_record(server_id="server1"), - generate_mock_mcp_server_db_record(server_id="server2"), - ] - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=team_records, - mcp_servers=db_servers, - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_allowed_servers = AsyncMock( - return_value={ - "server1": { - "server_id": "server1", - "server_name": "Test DB Server", - "status": "healthy", - "tools_count": 2, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 100.0, - "error": None, - }, - "server2": { - "server_id": "server2", - "server_name": "Test DB Server", - "status": "unhealthy", - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 5000.0, - "error": "Connection timeout", - }, - } - ) - mock_manager.get_allowed_mcp_servers = AsyncMock( - return_value=["server1", "server2"] - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_all_mcp_servers, - ) - - result = await health_check_all_mcp_servers( - user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["total_servers"] == 2 - assert result["healthy_count"] == 1 - assert result["unhealthy_count"] == 1 - assert result["unknown_count"] == 0 - assert "server1" in result["servers"] - assert "server2" in result["servers"] - - # Check individual server results - assert result["servers"]["server1"]["status"] == "healthy" - assert result["servers"]["server1"]["tools_count"] == 2 - assert result["servers"]["server1"]["server_name"] == "Test DB Server" - assert result["servers"]["server2"]["status"] == "unhealthy" - assert result["servers"]["server2"]["error"] == "Connection timeout" - assert result["servers"]["server2"]["server_name"] == "Test DB Server" - - @pytest.mark.asyncio - async def test_fetch_all_mcp_servers_with_health_status(self): - """Test that fetch_all_mcp_servers includes health check status""" - # Mock server with health status - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - # Add health status to the mock server - mock_server.status = "healthy" - mock_server.last_health_check = datetime.now() - mock_server.health_check_error = None - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=[], - mcp_servers=[], # Don't add servers here since we're mocking get_all_mcp_servers - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.config_mcp_servers = {} - mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=[mock_server] - ) - - mock_server.credentials = {"auth_value": "secret"} - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - fetch_all_mcp_servers, - ) - - result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) - - # Verify health check status is included - assert len(result) == 1 - server = result[0] - assert server.server_id == "test-server" - assert server.status == "healthy" - assert server.last_health_check is not None - assert server.health_check_error is None - assert server.credentials is None - - class TestTemporaryMCPSessionEndpoints: def test_inherit_credentials_from_existing_server(self): payload = NewMCPServerRequest( @@ -1170,7 +886,6 @@ class TestTemporaryMCPSessionEndpoints: fallback_client_id="server-1", ) - class TestUpdateMCPServer: """Test suite for update MCP server functionality"""