From 3a6b332b3b9f40b6a26e338dd035e2a5930e4e20 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 1 Jul 2026 04:23:54 +0000
Subject: [PATCH] fix: use _init_non_llm_objects_in_db to load all DB objects
 with proper guards

Address Greptile review: call _init_non_llm_objects_in_db instead of
_init_mcp_servers_in_db so all non-LLM DB objects (guardrails, agents,
vector stores, etc.) are loaded on startup when store_model_in_db is
False, and the _should_load_db_object guard is respected
---
 litellm/proxy/proxy_server.py                 |  2 +-
 tests/test_litellm/proxy/test_proxy_server.py | 23 +++++++------------
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index be77a74816a..ccbd26cded3 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -7535,7 +7535,7 @@ class ProxyStartupEvent:
             )
             await proxy_config.get_credentials(prisma_client=prisma_client)
         else:
-            await proxy_config._init_mcp_servers_in_db()
+            await proxy_config._init_non_llm_objects_in_db(prisma_client=prisma_client)
 
         await cls._initialize_slack_alerting_jobs(
             scheduler=scheduler,
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index f424985c5d0..a952564182f 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -6288,22 +6288,18 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch):
         # add_deployment should NOT have been called since store_model_in_db is False
         mock_proxy_config.add_deployment.assert_not_called()
 
-        # _init_mcp_servers_in_db should still be called even when
-        # store_model_in_db is False so DB-added MCP servers are loaded
-        # on restart (LIT-4128)
-        mock_proxy_config._init_mcp_servers_in_db.assert_awaited_once()
+        mock_proxy_config._init_non_llm_objects_in_db.assert_awaited_once()
 
 
 @pytest.mark.asyncio
-async def test_mcp_servers_loaded_on_startup_without_store_model_in_db(monkeypatch):
+async def test_non_llm_objects_loaded_on_startup_without_store_model_in_db(monkeypatch):
     """
-    Regression test for LIT-4128: MCP servers added via the UI are stored
-    in the database but only loaded into the in-memory registry during
-    add_deployment, which is gated by store_model_in_db=True. When
-    store_model_in_db is False (common in single-instance deployments),
-    the registry stays empty after restart until a new server is added.
+    Regression test for LIT-4128: non-LLM DB objects (MCP servers, guardrails,
+    agents, etc.) are only loaded inside add_deployment, which is gated by
+    store_model_in_db=True. When that flag is False (common in single-instance
+    deployments), the in-memory registries stay empty after restart.
 
-    Verify that _init_mcp_servers_in_db is called during startup even when
+    Verify that _init_non_llm_objects_in_db is called during startup even when
     store_model_in_db is False.
     """
     monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
@@ -6331,11 +6327,8 @@ async def test_mcp_servers_loaded_on_startup_without_store_model_in_db(monkeypat
             proxy_logging_obj=mock_proxy_logging,
         )
 
-        # add_deployment should NOT be called (store_model_in_db is False)
         mock_proxy_config.add_deployment.assert_not_called()
-
-        # MCP servers should still be loaded from DB on startup
-        mock_proxy_config._init_mcp_servers_in_db.assert_awaited_once()
+        mock_proxy_config._init_non_llm_objects_in_db.assert_awaited_once()
 
 
 # =====================================================================