From 8c11ddfc5bf9770bbc46ae2168cf132a50340e46 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 17:12:17 +0900 Subject: [PATCH 1/4] fix: periodically refresh MCP registry across instances via scheduler job --- .../mcp_server/mcp_server_manager.py | 28 +++++++++++++++--- litellm/proxy/proxy_server.py | 29 +++++++++++++++++++ .../types/mcp_server/mcp_server_manager.py | 2 ++ 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3a548e203c5..6e7db26a95c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -551,6 +551,7 @@ class MCPServerManager: allowed_tools=getattr(mcp_server, "allowed_tools", None), disallowed_tools=getattr(mcp_server, "disallowed_tools", None), allow_all_keys=mcp_server.allow_all_keys, + updated_at=getattr(mcp_server, "updated_at", None), ) return new_server @@ -2074,15 +2075,34 @@ class MCPServerManager: db_mcp_servers = await get_all_mcp_servers(prisma_client) verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database") - # ensure the global_mcp_server_manager is up to date with the db + previous_registry = self.registry + new_registry: Dict[str, MCPServer] = {} + for server in db_mcp_servers: + existing_server = previous_registry.get(server.server_id) + + if ( + existing_server is not None + and existing_server.updated_at is not None + and server.updated_at is not None + and existing_server.updated_at == server.updated_at + ): + # Re-use existing server instance to avoid re-running build_mcp_server_from_table() + # which can perform network discovery for OAuth2 servers. + new_registry[server.server_id] = existing_server + continue + verbose_logger.debug( - f"Adding server to registry: {server.server_id} ({server.server_name})" + f"Building server from DB: {server.server_id} ({server.server_name})" ) - await self.add_server(server) + new_registry[server.server_id] = await self.build_mcp_server_from_table( + server + ) + + self.registry = new_registry verbose_logger.debug( - f"Registry now contains {len(self.get_registry())} servers" + "MCP registry refreshed (%s servers in registry)", len(new_registry) ) def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e58ae7c643..77a0f5a1c09 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4109,6 +4109,23 @@ class ProxyConfig: return [] +async def _reload_mcp_servers_job(): + """Background job entrypoint for MCP registry refreshes.""" + if proxy_config._should_load_db_object(object_type="mcp") is False: + return + + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + await global_mcp_server_manager._add_mcp_servers_from_db_to_in_memory_registry() # noqa: SLF001 + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reload MCP servers from database: %s", str(e) + ) + + proxy_config = ProxyConfig() @@ -4646,6 +4663,18 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) await proxy_config.get_credentials(prisma_client=prisma_client) + + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if is_mcp_available(): + scheduler.add_job( + _reload_mcp_servers_job, + "interval", + seconds=30, + id="reload_mcp_servers_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) await cls._initialize_slack_alerting_jobs( scheduler=scheduler, general_settings=general_settings, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 96fd79f466b..94f33ffb297 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,3 +1,4 @@ +from datetime import datetime from typing import Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict @@ -50,4 +51,5 @@ class MCPServer(BaseModel): env: Optional[Dict[str, str]] = None access_groups: Optional[List[str]] = None allow_all_keys: bool = False + updated_at: Optional[datetime] = None model_config = ConfigDict(arbitrary_types_allowed=True) From 8b90e5f4dd53766694c9fe833b1601c695b8df56 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 17:16:38 +0900 Subject: [PATCH 2/4] refactor: expose MCP registry reload helper --- .../_experimental/mcp_server/mcp_server_manager.py | 10 ++-------- litellm/proxy/proxy_server.py | 4 ++-- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6e7db26a95c..fb4943100ad 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2060,7 +2060,8 @@ class MCPServerManager: return None - async def _add_mcp_servers_from_db_to_in_memory_registry(self): + async def reload_servers_from_database(self): + """Re-synchronize the in-memory MCP server registry with the database.""" from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_prisma_client_or_throw, @@ -2389,13 +2390,6 @@ class MCPServerManager: servers.append(self._build_mcp_server_table(server)) return servers - async def reload_servers_from_database(self): - """ - Public method to reload all MCP servers from database into registry. - This can be called from management endpoints to ensure registry is up to date. - """ - await self._add_mcp_servers_from_db_to_in_memory_registry() - async def get_all_mcp_servers_with_health_unfiltered( self, server_ids: Optional[List[str]] = None ) -> List[LiteLLM_MCPServerTable]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77a0f5a1c09..4a44586f0d4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3972,7 +3972,7 @@ class ProxyConfig: ) try: - await global_mcp_server_manager._add_mcp_servers_from_db_to_in_memory_registry() + await global_mcp_server_manager.reload_servers_from_database() except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format( @@ -4119,7 +4119,7 @@ async def _reload_mcp_servers_job(): global_mcp_server_manager, ) - await global_mcp_server_manager._add_mcp_servers_from_db_to_in_memory_registry() # noqa: SLF001 + await global_mcp_server_manager.reload_servers_from_database() except Exception as e: verbose_proxy_logger.exception( "Failed to reload MCP servers from database: %s", str(e) From 5927a557fbd78e03549bb0f9d25d840b42678698 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 17:25:21 +0900 Subject: [PATCH 3/4] tests: add test --- .../mcp_server/test_mcp_server.py | 104 +++++++++++++++++- 1 file changed, 103 insertions(+), 1 deletion(-) 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 8062243dfdd..f1558ac5791 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 @@ -1,4 +1,5 @@ import asyncio +from datetime import datetime, timedelta from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -7,7 +8,12 @@ from fastapi import HTTPException from mcp import ReadResourceResult, Resource from mcp.types import Prompt, ResourceTemplate, TextResourceContents -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + MCPTransport, + UserAPIKeyAuth, +) +from litellm.types.mcp_server.mcp_server_manager import MCPServer @pytest.mark.asyncio @@ -1688,3 +1694,99 @@ def test_filter_tools_by_allowed_tools(): assert len(filtered_tools) == 2 assert filtered_tools[0].name == "my_api_mcp-getpetbyid" assert filtered_tools[1].name == "my_api_mcp-findpetsbystatus" + + +def _make_db_mcp_server(server_id: str, updated_at: datetime) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name="server", + alias="server", + url="https://example.com", + transport=MCPTransport.http, + created_at=updated_at, + updated_at=updated_at, + mcp_info={}, + ) + + +class TestMCPServerManagerReload: + @pytest.mark.asyncio + async def test_reuses_existing_server_when_updated_at_matches(self): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + timestamp = datetime.utcnow() + existing_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=timestamp, + ) + manager.registry = {existing_server.server_id: existing_server} + + db_row = _make_db_mcp_server("server-1", timestamp) + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", + new=AsyncMock(return_value=[db_row]), + ) as mock_get_all, patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=object(), + ), patch.object( + manager, "build_mcp_server_from_table", AsyncMock() + ) as mock_build: + await manager.reload_servers_from_database() + + mock_get_all.assert_awaited_once() + mock_build.assert_not_awaited() + assert manager.registry["server-1"] is existing_server + + @pytest.mark.asyncio + async def test_rebuilds_server_when_updated_at_changes(self): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + timestamp = datetime.utcnow() + existing_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=timestamp, + ) + manager.registry = {existing_server.server_id: existing_server} + + new_timestamp = timestamp + timedelta(minutes=5) + db_row = _make_db_mcp_server("server-1", new_timestamp) + rebuilt_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=new_timestamp, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", + new=AsyncMock(return_value=[db_row]), + ) as mock_get_all, patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=object(), + ), patch.object( + manager, + "build_mcp_server_from_table", + AsyncMock(return_value=rebuilt_server), + ) as mock_build: + await manager.reload_servers_from_database() + + mock_get_all.assert_awaited_once() + mock_build.assert_awaited_once_with(db_row) + assert manager.registry["server-1"] is rebuilt_server From d4483d8422e5b92de414f8d248113a50e54eef51 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 17:31:30 +0900 Subject: [PATCH 4/4] fix: formatter --- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fb4943100ad..1029f2241a1 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -698,9 +698,7 @@ class MCPServerManager: results = await asyncio.gather(*tasks) # Flatten results into single list - list_tools_result: List[MCPTool] = [ - tool for tools in results for tool in tools - ] + list_tools_result: List[MCPTool] = [tool for tools in results for tool in tools] verbose_logger.info( f"Successfully fetched {len(list_tools_result)} tools total from all servers"