scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

fix(proxy): hydrate DB-backed MCP servers on startup when store_model_in_db is true

Browse source 
The startup gate only called init_mcp_servers_from_db when store_model_in_db was
not True, so a DB-backed proxy with store_model_in_db=True never hydrated the
in-memory MCP registry on boot. MCP servers created via the Admin UI then
disappeared from GET /v1/mcp/server after a restart until an unrelated write
triggered a reload. Call init_mcp_servers_from_db unconditionally; whether MCP
objects load is already governed by _should_load_db_object("mcp") inside it, so
the supported_db_objects allowlist behavior is unchanged.

Fixes #32575
This commit is contained in:
Yash Raj Pandey 2026-07-09 11:11:21 -04:00
parent cd63c7e5a7
commit c052429f2f
2 changed files with 42 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/proxy_server.py
View file

@ -9176,8 +9176,12 @@ class ProxyStartupEvent:
redis_cache=redis_usage_cache,
)
# init_mcp_servers_from_db already gates on _should_load_db_object("mcp"), so calling
# it unconditionally keeps the allowlist behavior while letting store_model_in_db=True
# deployments hydrate the in-memory MCP registry on boot.
await proxy_config.init_mcp_servers_from_db()
if store_model_in_db is not True:
await proxy_config.init_mcp_servers_from_db()
# Without this branch's own refresh, a UI-created search tool never reaches the router:
# the add_deployment job that carries it in store_model_in_db=True mode is not scheduled.
await proxy_config.reload_search_tools_from_db()

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

@ -926,6 +926,43 @@ async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_fal
mock_proxy_config.init_mcp_servers_from_db.assert_awaited_once()
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_true(monkeypatch): # test-quality-ok: the awaited call is the behaviour; the bug was that it never ran
"""
Regression (issue #32575): with store_model_in_db=True the startup gate skipped
init_mcp_servers_from_db entirely, so DB-backed MCP servers created via the Admin
UI vanished from GET /v1/mcp/server after a restart until the next write. MCP
hydration must run regardless of store_model_in_db.
"""
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()
# set on the instance, so a class spec does not expose it
mock_proxy_logging.db_spend_update_writer = MagicMock()
mock_proxy_config = AsyncMock()
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), # test-quality-ok: proxy_config is module state the startup path reads, no injection point
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: store_model_in_db is the module flag whose gate this test exercises
):
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,
)
mock_proxy_config.add_deployment.assert_awaited()
mock_proxy_config.init_mcp_servers_from_db.assert_awaited_once()
@pytest.mark.asyncio
async def test_init_mcp_servers_from_db_respects_supported_db_objects(monkeypatch):
"""