fix: periodically refresh MCP registry across instances via scheduler job

This commit is contained in:
Yuta Saito 2026-01-09 17:12:17 +09:00
parent d9b275e62a
commit 8c11ddfc5b
3 changed files with 55 additions and 4 deletions

View file

@ -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]:

View file

@ -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,

View file

@ -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)