mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix: periodically refresh MCP registry across instances via scheduler job
This commit is contained in:
parent
d9b275e62a
commit
8c11ddfc5b
3 changed files with 55 additions and 4 deletions
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue