mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(proxy): hydrate MCP server registry from DB on startup when store_model_in_db is false (#31775)
MCP servers created through the UI are persisted to the database independent of store_model_in_db, but the in-memory registry that GET /v1/mcp/server reads was hydrated from the database only through add_deployment, which runs solely when store_model_in_db is True. On a DB-backed single-instance proxy with store_model_in_db unset the registry started empty after a restart, so the MCP Servers page showed nothing until an add or edit triggered a reload. Hydrate the registry from the database on startup regardless of store_model_in_db via a new ProxyConfig.init_mcp_servers_from_db, honoring supported_db_objects.
This commit is contained in:
parent
cca71a07c2
commit
13b590c8ec
2 changed files with 67 additions and 0 deletions
|
|
@ -6314,6 +6314,10 @@ class ProxyConfig:
|
|||
"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format(str(e))
|
||||
)
|
||||
|
||||
async def init_mcp_servers_from_db(self) -> None:
|
||||
if self._should_load_db_object(object_type="mcp"):
|
||||
await self._init_mcp_servers_in_db()
|
||||
|
||||
async def _init_agents_in_db(self, prisma_client: PrismaClient):
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
global_agent_registry as AGENT_REGISTRY,
|
||||
|
|
@ -7561,6 +7565,9 @@ class ProxyStartupEvent:
|
|||
)
|
||||
await proxy_config.get_credentials(prisma_client=prisma_client)
|
||||
|
||||
if store_model_in_db is not True:
|
||||
await proxy_config.init_mcp_servers_from_db()
|
||||
|
||||
await cls._initialize_slack_alerting_jobs(
|
||||
scheduler=scheduler,
|
||||
general_settings=general_settings,
|
||||
|
|
|
|||
|
|
@ -754,6 +754,66 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
|
|||
assert len(mock_scheduler_calls) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_false(monkeypatch):
|
||||
"""
|
||||
Regression (LIT-4128): MCP servers created via the UI are persisted to the DB
|
||||
regardless of store_model_in_db, but the in-memory registry that GET
|
||||
/v1/mcp/server reads is hydrated from the DB only by the store_model_in_db
|
||||
model-sync loop (add_deployment). On a DB-backed proxy with store_model_in_db
|
||||
unset the registry must still be hydrated on startup so previously-added
|
||||
servers survive a restart instead of showing an empty list until a write.
|
||||
"""
|
||||
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()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
):
|
||||
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_not_called()
|
||||
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):
|
||||
"""
|
||||
init_mcp_servers_from_db hydrates MCP from the DB by default but skips it when
|
||||
an explicit supported_db_objects allowlist omits "mcp".
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
config = ProxyConfig()
|
||||
with patch.object(config, "_init_mcp_servers_in_db", new=AsyncMock()) as mock_init:
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
await config.init_mcp_servers_from_db()
|
||||
mock_init.assert_awaited_once()
|
||||
|
||||
mock_init.reset_mock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"supported_db_objects": ["models"]},
|
||||
)
|
||||
await config.init_mcp_servers_from_db()
|
||||
mock_init.assert_not_awaited()
|
||||
|
||||
|
||||
def test_update_config_fields_deep_merge_db_wins():
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue