From 7fdea85b5c6549ecc2c3a39bdf32c733745c0ae0 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Wed, 31 Dec 2025 06:55:47 +0900 Subject: [PATCH] feat: optimize MCP server listing by separating health checks --- .../mcp_server/mcp_server_manager.py | 79 ++++++- .../mcp_management_endpoints.py | 55 ++++- .../test_mcp_management_endpoints.py | 199 ++++++++++++++++-- .../mcpServers/useMCPServerHealth.test.ts | 127 +++++++++++ .../hooks/mcpServers/useMCPServerHealth.ts | 22 ++ .../mcp_tools/mcp_server_columns.tsx | 14 ++ .../components/mcp_tools/mcp_servers.test.tsx | 105 +++++++++ .../src/components/mcp_tools/mcp_servers.tsx | 50 +++-- .../src/components/networking.tsx | 38 ++++ 9 files changed, 653 insertions(+), 36 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 58e9a345dc..2d3ea1e827 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2186,8 +2186,12 @@ class MCPServerManager: async def _noop(session): return "ok" - await client.run_with_session(_noop) + # Add timeout wrapper to prevent hanging + await asyncio.wait_for(client.run_with_session(_noop), timeout=10.0) status = "healthy" + except asyncio.TimeoutError: + health_check_error = "Health check timed out after 10 seconds" + status = "unhealthy" except Exception as e: health_check_error = str(e) status = "unhealthy" @@ -2221,14 +2225,15 @@ class MCPServerManager: async def get_all_mcp_servers_with_health_and_teams( self, user_api_key_auth: Optional[UserAPIKeyAuth] = None, - include_health: bool = True, + server_ids: Optional[List[str]] = None, ) -> List[LiteLLM_MCPServerTable]: """ Get all MCP servers that the user has access to, with health status and team information. Args: user_api_key_auth: User authentication info for access control - include_health: Whether to include health check information + server_ids: Optional list of server IDs to filter. If provided, only these servers + will be checked (subject to access control). If None, all accessible servers are checked. Returns: List of MCP server objects with health and team data @@ -2237,10 +2242,16 @@ class MCPServerManager: # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + # Filter by requested server_ids if provided + if server_ids: + # Only check servers that are both requested AND accessible + target_server_ids = [sid for sid in server_ids if sid in allowed_server_ids] + else: + # Check all accessible servers + target_server_ids = allowed_server_ids + # Run health checks concurrently - tasks = [ - self.health_check_server(server_id) for server_id in allowed_server_ids - ] + tasks = [self.health_check_server(server_id) for server_id in target_server_ids] results = await asyncio.gather(*tasks) # Filter out None results (servers that were not found) @@ -2248,6 +2259,62 @@ class MCPServerManager: return list_mcp_servers + async def get_all_allowed_mcp_servers( + self, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[LiteLLM_MCPServerTable]: + """ + Get all MCP servers that the user has access to. + + Args: + user_api_key_auth: User authentication info for access control + + Returns: + List of MCP server objects without health status + """ + from datetime import datetime + + # Get allowed server IDs + allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + + list_mcp_servers: List[LiteLLM_MCPServerTable] = [] + + for server_id in allowed_server_ids: + server = self.get_mcp_server_by_id(server_id) + if not server: + verbose_logger.warning(f"MCP Server {server_id} not found in registry") + continue + + # Build LiteLLM_MCPServerTable without health check + mcp_server_table = 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=None, # No health check performed + last_health_check=None, # No health check performed + health_check_error=None, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + ) + list_mcp_servers.append(mcp_server_table) + + return list_mcp_servers + async def reload_servers_from_database(self): """ Public method to reload all MCP servers from database into registry. diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 59b9659dd8..500323d3be 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -16,7 +16,7 @@ Endpoints here: import importlib from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( APIRouter, @@ -24,6 +24,7 @@ from fastapi import ( Form, Header, HTTPException, + Query, Request, Response, status, @@ -318,7 +319,7 @@ if MCP_AVAILABLE: aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( user_api_key_auth=auth_context ) for server in servers: @@ -336,6 +337,56 @@ if MCP_AVAILABLE: server.mcp_info["is_public"] = True return redacted_mcp_servers + @router.get( + "/server/health", + description="Health check for MCP servers", + dependencies=[Depends(user_api_key_auth)], + ) + async def health_check_servers( + server_ids: Optional[List[str]] = Query( + None, + description="Server IDs to check. If not provided, checks all accessible servers.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Perform health checks on one or more MCP servers. + + Parameters: + - server_ids: Optional list of server IDs. If not provided, checks all accessible servers. + + Returns: + - Health check results for requested servers + + ``` + # Check all accessible servers + curl --location 'http://localhost:4000/v1/mcp/server/health' \ + --header 'Authorization: Bearer your_api_key_here' + + # Check specific servers + curl --location 'http://localhost:4000/v1/mcp/server/health?server_ids=server-1&server_ids=server-2' \ + --header 'Authorization: Bearer your_api_key_here' + ``` + """ + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + + server_status_map: Dict[ + str, Optional[Literal["healthy", "unhealthy", "unknown"]] + ] = {} + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + user_api_key_auth=auth_context, + server_ids=server_ids, + ) + for server in servers: + if server.server_id not in server_status_map: + server_status_map[server.server_id] = server.status + + return [ + {"server_id": server_id, "status": status} + for server_id, status in server_status_map.items() + ] + @router.get( "/server/{server_id}", description="Returns the mcp server info", 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 8d5b3dcedc..cd1a1f5e10 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 @@ -169,8 +169,8 @@ class TestListMCPServers: return_value=["config_server_1", "config_server_2"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ generate_mock_mcp_server_db_record( server_id="config_server_1", alias="Zapier MCP", @@ -184,11 +184,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -200,6 +200,9 @@ class TestListMCPServers: ), 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.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -300,8 +303,8 @@ class TestListMCPServers: ] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_1, db_server_2, generate_mock_mcp_server_db_record( @@ -317,11 +320,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -333,6 +336,9 @@ class TestListMCPServers: ), 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.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -425,8 +431,8 @@ class TestListMCPServers: return_value=["db_server_allowed", "config_server_allowed"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_allowed, generate_mock_mcp_server_db_record( server_id="config_server_allowed", @@ -434,11 +440,11 @@ class TestListMCPServers: url="https://actions.zapier.com/mcp/sse", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -450,6 +456,9 @@ class TestListMCPServers: ), 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.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -975,3 +984,163 @@ class TestUpdateMCPServer: # Verify the result includes extra_headers assert result.extra_headers == ["X-Custom-Header", "X-Another-Header"] assert result.alias == "Updated Test Server" + + +class TestHealthCheckServers: + """Test suite for health check servers endpoint""" + + @pytest.mark.asyncio + async def test_health_check_all_servers(self): + """ + Test health check for all accessible servers + + Scenario: User has access to 2 servers, checks all + Expected: Returns health status for both servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check results + mock_health_result_1 = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result_1.status = "healthy" + mock_health_result_1.last_health_check = datetime.now() + mock_health_result_1.health_check_error = None + + mock_health_result_2 = generate_mock_mcp_server_db_record( + server_id="server-2", + alias="Server 2", + url="https://server2.example.com", + ) + mock_health_result_2.status = "unhealthy" + mock_health_result_2.last_health_check = datetime.now() + mock_health_result_2.health_check_error = "Connection timeout" + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result_1, mock_health_result_2] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=None, + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 2 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + assert result[1]["server_id"] == "server-2" + assert result[1]["status"] == "unhealthy" + + @pytest.mark.asyncio + async def test_health_check_specific_servers(self): + """ + Test health check for specific servers + + Scenario: User requests health check for specific server IDs + Expected: Returns health status only for requested servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + + @pytest.mark.asyncio + async def test_health_check_unauthorized_servers(self): + """ + Test health check with unauthorized servers + + Scenario: User requests health check for servers they don't have access to + Expected: Only checks accessible servers, unauthorized servers are filtered out + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result for authorized server + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager - server_ids filter is applied inside get_all_mcp_servers_with_health_and_teams + # So it only returns servers the user has access to + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] # Only server-1 is returned (accessible) + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1", "server-unauthorized"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results - only accessible server is returned + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts new file mode 100644 index 0000000000..be910acf7e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts @@ -0,0 +1,127 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { renderHook, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { useMCPServerHealth } from "./useMCPServerHealth"; +import * as networking from "@/components/networking"; + +// Mock the networking module +vi.mock("@/components/networking", () => ({ + fetchMCPServerHealth: vi.fn(), +})); + +// Mock useAuthorized hook +vi.mock("../useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "test-token-123", + })), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = createQueryClient(); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +describe("useMCPServerHealth", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should fetch health status for given server IDs", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(["server-1", "server-2"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", ["server-1", "server-2"]); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should fetch health status for all servers when no server IDs provided", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "healthy" }, + { server_id: "server-3", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", undefined); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should handle empty server list", async () => { + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + + const { result } = renderHook(() => useMCPServerHealth([]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", []); + expect(result.current.data).toEqual([]); + }); + + it("should handle errors when fetching health status", async () => { + const mockError = new Error("Failed to fetch health status"); + vi.mocked(networking.fetchMCPServerHealth).mockRejectedValue(mockError); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(mockError); + }); + + it("should not fetch when accessToken is not available", async () => { + // Mock useAuthorized to return no token + const useAuthorizedModule = await import("../useAuthorized"); + vi.mocked(useAuthorizedModule.default).mockReturnValue({ + accessToken: null, + } as any); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + // Should remain in idle state since query is not enabled + expect(result.current.status).toBe("pending"); + expect(networking.fetchMCPServerHealth).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts new file mode 100644 index 0000000000..ad3c633eb9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -0,0 +1,22 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPServerHealth } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +const mcpServerHealthKeys = createQueryKeys("mcpServerHealth"); + +interface MCPServerHealth { + server_id: string; + status: string; +} + +export const useMCPServerHealth = (serverIds?: string[]) => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: mcpServerHealthKeys.list({ serverIds }), + queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds), + enabled: !!accessToken, + // Refetch health status every 30 seconds to keep it up to date + refetchInterval: 30000, + }); +}; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index 1bf719ef90..f6a5d6622d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -10,6 +10,7 @@ export const mcpServerColumns = ( onView: (serverId: string) => void, onEdit: (serverId: string) => void, onDelete: (serverId: string) => void, + isLoadingHealth?: boolean, ): ColumnDef[] => [ { accessorKey: "server_id", @@ -58,6 +59,19 @@ export const mcpServerColumns = ( const lastCheck = server.last_health_check; const error = server.health_check_error; + // Show loading spinner if health check is in progress + if (isLoadingHealth) { + return ( +
+ + + + + Loading... +
+ ); + } + const getStatusColor = (status: string) => { switch (status) { case "healthy": diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx index 776d579fc1..4b8698b976 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx @@ -8,6 +8,7 @@ import * as networking from "../networking"; // Mock the networking module vi.mock("../networking", () => ({ fetchMCPServers: vi.fn(), + fetchMCPServerHealth: vi.fn(), deleteMCPServer: vi.fn(), getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), })); @@ -123,4 +124,108 @@ describe("MCPServers", () => { // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock expect(networking.fetchMCPServers).toHaveBeenCalledWith("123"); }); + + it("should fetch and merge health status for servers", async () => { + // Mock MCP servers data without health status + const mockServers = [ + { + server_id: "server-1", + server_name: "Test Server 1", + alias: "test-server-1", + url: "https://example.com/mcp", + transport: "http", + auth_type: "none", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + teams: [], + mcp_access_groups: [], + status: undefined, + }, + { + server_id: "server-2", + server_name: "Test Server 2", + alias: "test-server-2", + url: "https://example2.com/mcp", + transport: "sse", + auth_type: "api_key", + created_at: "2024-01-02T00:00:00Z", + created_by: "user-2", + updated_at: "2024-01-02T00:00:00Z", + updated_by: "user-2", + teams: [], + mcp_access_groups: ["group-1"], + status: undefined, + }, + ]; + + // Mock health status data + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const queryClient = createQueryClient(); + const { getByText } = render( + + + , + ); + + // Wait for the component to load + await waitFor(() => { + expect(getByText("MCP Servers")).toBeInTheDocument(); + }); + + // Verify the health check API was called with server IDs + await waitFor(() => { + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("123", ["server-1", "server-2"]); + }); + }); + + it("should display loading state while health check is in progress", async () => { + const mockServers = [ + { + server_id: "server-1", + server_name: "Test Server 1", + alias: "test-server-1", + url: "https://example.com/mcp", + transport: "http", + auth_type: "none", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + teams: [], + mcp_access_groups: [], + }, + ]; + + vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers); + // Mock health check to never resolve (to test loading state) + vi.mocked(networking.fetchMCPServerHealth).mockImplementation( + () => new Promise(() => {}), // Never resolves + ); + + const queryClient = createQueryClient(); + const { getByText } = render( + + + , + ); + + // Wait for the component to load + await waitFor(() => { + expect(getByText("MCP Servers")).toBeInTheDocument(); + }); + + // Verify that health check was initiated + await waitFor(() => { + expect(networking.fetchMCPServerHealth).toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index d5b147f85c..f6669fb282 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -2,8 +2,9 @@ import { isAdminRole } from "@/utils/roles"; import { QuestionCircleOutlined } from "@ant-design/icons"; import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; import { Descriptions, Modal, Select, Tooltip, Typography } from "antd"; -import React, { useEffect, useState } from "react"; +import React, { useEffect, useState, useMemo } from "react"; import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMCPServerHealth"; import NotificationsManager from "../molecules/notifications_manager"; import { deleteMCPServer } from "../networking"; import { DataTable } from "../view_logs/table"; @@ -19,7 +20,29 @@ const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { - const { data: mcpServers, isLoading: isLoadingServers, refetch, dataUpdatedAt } = useMCPServers(); + const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); + + // Fetch health status for all servers + const serverIds = useMemo(() => mcpServers?.map((server) => server.server_id), [mcpServers]); + const { data: healthStatuses, isLoading: isLoadingHealth } = useMCPServerHealth(serverIds); + + // Merge health status data into servers + const serversWithHealth = useMemo(() => { + if (!mcpServers) return []; + if (!healthStatuses) return mcpServers; + + const healthMap = new Map(healthStatuses.map((h) => [h.server_id, h.status])); + + return mcpServers.map((server) => { + const healthStatus = healthMap.get(server.server_id); + return { + ...server, + status: healthStatus + ? (healthStatus as "healthy" | "unhealthy" | "unknown") + : server.status, + }; + }); + }, [mcpServers, healthStatuses]); // Log allowed_tools from fetched servers React.useEffect(() => { @@ -65,10 +88,10 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { - if (!mcpServers) return []; + if (!serversWithHealth) return []; const teamsSet = new Set(); const uniqueTeamsArray: Team[] = []; - mcpServers.forEach((server: MCPServer) => { + serversWithHealth.forEach((server: MCPServer) => { if (server.teams) { server.teams.forEach((team: Team) => { const teamKey = team.team_id; @@ -80,17 +103,17 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } }); return uniqueTeamsArray; - }, [mcpServers]); + }, [serversWithHealth]); // Get unique MCP access groups from all servers const uniqueMcpAccessGroups = React.useMemo(() => { - if (!mcpServers) return []; + if (!serversWithHealth) return []; return Array.from( new Set( - mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), + serversWithHealth.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), ), ); - }, [mcpServers]); + }, [serversWithHealth]); // Handle team filter change const handleTeamChange = (teamId: string) => { @@ -106,8 +129,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Filtering logic for both team and access group const filterServers = (teamId: string, group: string) => { - if (!mcpServers) return setFilteredServers([]); - let filtered = mcpServers; + if (!serversWithHealth) return setFilteredServers([]); + let filtered = serversWithHealth; if (teamId === "personal") { setFilteredServers([]); return; @@ -123,10 +146,10 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) setFilteredServers(filtered); }; - // Initial and effect-based filtering (trigger on query data updates) + // Initial and effect-based filtering (trigger on query data updates and health data updates) useEffect(() => { filterServers(selectedTeam, selectedMcpAccessGroup); - }, [dataUpdatedAt]); + }, [serversWithHealth, selectedTeam, selectedMcpAccessGroup]); const columns = React.useMemo( () => @@ -141,8 +164,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) setEditServer(true); }, handleDelete, + isLoadingHealth, ), - [userRole], + [userRole, isLoadingHealth], ); function handleDelete(server_id: string) { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index aeba207db2..e244a09613 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5687,6 +5687,44 @@ export const fetchMCPServers = async (accessToken: string) => { } }; +export const fetchMCPServerHealth = async (accessToken: string, serverIds?: string[]) => { + try { + // Construct base URL + let url = proxyBaseUrl ? `${proxyBaseUrl}/v1/mcp/server/health` : `/v1/mcp/server/health`; + + // Add server_ids query parameters if provided + if (serverIds && serverIds.length > 0) { + const params = new URLSearchParams(); + serverIds.forEach((id) => params.append("server_ids", id)); + url = `${url}?${params.toString()}`; + } + + console.log("Fetching MCP server health from:", url); + + const response = await fetch(url, { + method: HTTP_REQUEST.GET, + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Fetched MCP server health:", data); + return data; + } catch (error) { + console.error("Failed to fetch MCP server health:", error); + throw error; + } +}; + export const fetchMCPAccessGroups = async (accessToken: string) => { try { // Construct base URL