Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions litellm/proxy/proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -6311,6 +6311,10 @@ async def _init_mcp_servers_in_db(self):
"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,
Expand Down Expand Up @@ -7558,6 +7562,9 @@ async def initialize_scheduled_background_jobs(
)
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,
Expand Down
60 changes: 60 additions & 0 deletions tests/test_litellm/proxy/test_proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading