scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(proxy): add periodic MCP server DB sync for cross-pod consistency (#22325)

Browse source 
* Adds sync_mcp_servers_from_db() and registers as dedicated 30s scheduler job
* Removes duplicate sync from _init_non_llm_objects_in_db
* Eagerly calls sync on startup
* Includes unit tests for sync execution when MCP is available/unavailable
This commit is contained in:
gavksingh 2026-03-01 15:55:10 -05:00
parent 8053be60df
commit 0c5c670b3e
2 changed files with 136 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/proxy_server.py
View file

@ -4382,7 +4382,7 @@ class ProxyConfig:
"""
Use this to read non-llm objects from the db and initialize them
ex. Vector Stores, Guardrails, MCP tools, etc.
ex. Vector Stores, Guardrails, Agents, etc.
"""
if self._should_load_db_object(object_type="guardrails"):
await self._init_guardrails_in_db(prisma_client=prisma_client)
@ -4396,8 +4396,8 @@ class ProxyConfig:
if self._should_load_db_object(object_type="vector_store_indexes"):
await self._init_vector_store_indexes_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="mcp"):
await self._init_mcp_servers_in_db()
# MCP sync runs via its own dedicated scheduler job (mcp_server_sync_job).
# Calling it here too would double the DB queries every 30 s cycle.
if self._should_load_db_object(object_type="agents"):
await self._init_agents_in_db(prisma_client=prisma_client)
@ -4904,12 +4904,18 @@ class ProxyConfig:
)
)
async def _init_mcp_servers_in_db(self):
async def sync_mcp_servers_from_db(self):
"""
Periodically re-synchronize MCP servers from the database.
Called exclusively by the dedicated mcp_server_sync_job scheduler job
(registered in initialize_scheduled_background_jobs), which already
checks _should_load_db_object before scheduling and calling this method.
This keeps MCP configuration eventually consistent across pods.
"""
from litellm.proxy._experimental.mcp_server.utils import is_mcp_available
if not is_mcp_available():
verbose_proxy_logger.debug(
"MCP module not available, skipping MCP server initialization"
"MCP module not available, skipping MCP server synchronization"
)
return
@ -4919,9 +4925,12 @@ class ProxyConfig:
try:
await global_mcp_server_manager.reload_servers_from_database()
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.py::ProxyConfig:sync_mcp_servers_from_db - MCP servers synchronized"
)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format(
"litellm.proxy.proxy_server.py::ProxyConfig:sync_mcp_servers_from_db - {}".format(
str(e)
)
)
@ -5759,6 +5768,17 @@ class ProxyStartupEvent:
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
if proxy_config._should_load_db_object(object_type="mcp"):
scheduler.add_job(
proxy_config.sync_mcp_servers_from_db,
"interval",
seconds=30, # keep in sync with add_deployment polling
id="mcp_server_sync_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
await proxy_config.sync_mcp_servers_from_db()
### GET STORED CREDENTIALS ###
scheduler.add_job(
proxy_config.get_credentials,

109
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -444,6 +444,115 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
assert len(mock_scheduler_calls) > 0
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_mcp_sync_enabled(monkeypatch):
"""
Verify periodic MCP sync is scheduled and run once at startup when MCP DB objects are enabled.
"""
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
mock_prisma_client = MagicMock()
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_config = AsyncMock()
mock_proxy_config._should_load_db_object = MagicMock(return_value=True)
mock_proxy_config.sync_mcp_servers_from_db = AsyncMock()
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch(
"litellm.proxy.proxy_server.store_model_in_db", True
), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
import litellm.proxy.proxy_server as ps
assert mock_proxy_config.sync_mcp_servers_from_db.await_count == 1
assert ps.scheduler.get_job("mcp_server_sync_job") is not None
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_mcp_sync_disabled(monkeypatch):
"""
Verify periodic MCP sync is not scheduled when MCP DB objects are disabled.
"""
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
mock_prisma_client = MagicMock()
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_config = AsyncMock()
mock_proxy_config._should_load_db_object = MagicMock(return_value=False)
mock_proxy_config.sync_mcp_servers_from_db = AsyncMock()
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch(
"litellm.proxy.proxy_server.store_model_in_db", True
), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
import litellm.proxy.proxy_server as ps
mock_proxy_config.sync_mcp_servers_from_db.assert_not_awaited()
assert ps.scheduler.get_job("mcp_server_sync_job") is None
@pytest.mark.asyncio
async def test_sync_mcp_servers_from_db_available(monkeypatch):
"""
Verify that when MCP is available, reload_servers_from_database is called.
"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_mcp_manager = AsyncMock()
with patch("litellm.proxy._experimental.mcp_server.utils.is_mcp_available", return_value=True), patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", mock_mcp_manager
):
await proxy_config.sync_mcp_servers_from_db()
mock_mcp_manager.reload_servers_from_database.assert_awaited_once()
@pytest.mark.asyncio
async def test_sync_mcp_servers_from_db_unavailable(monkeypatch):
"""
Verify that when MCP is unavailable, reload_servers_from_database is NOT called.
"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_mcp_manager = AsyncMock()
with patch("litellm.proxy._experimental.mcp_server.utils.is_mcp_available", return_value=False), patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", mock_mcp_manager
):
await proxy_config.sync_mcp_servers_from_db()
mock_mcp_manager.reload_servers_from_database.assert_not_awaited()
def test_update_config_fields_deep_merge_db_wins():
from litellm.proxy.proxy_server import ProxyConfig