diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f632317d9a8..30556f8a353 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11316,14 +11316,26 @@ class ProxyStartupEvent: # but YAML config has False. if store_model_in_db is not True and prisma_client is not None: try: - _db_gs_record: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( - where={"param_name": "general_settings"} + _db_gs_record: Final[_ConfigParamRow | None] = await call_with_db_reconnect_retry( + prisma_client, + lambda: _config_param_table(prisma_client).find_first(where={"param_name": "general_settings"}), + reason="startup_store_model_in_db_general_settings_lookup_failure", ) - if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict): - _db_val: Final = _db_gs_record.param_value.get("store_model_in_db") - if _db_val is True or (isinstance(_db_val, str) and _db_val.lower() == "true"): - store_model_in_db = True - verbose_proxy_logger.info("store_model_in_db=True loaded from DB, overriding config/env") + if _db_gs_record is not None: + _db_param_value = _db_gs_record.param_value + # Prisma can return the JSON column as a raw string; normalize + # it the way ConfigRepository.get_param does before reading, so + # the override isn't silently skipped for string-valued rows. + if isinstance(_db_param_value, str): + try: + _db_param_value = json.loads(_db_param_value) + except json.JSONDecodeError: + _db_param_value = {} + if isinstance(_db_param_value, dict): + _db_val: Final = _db_param_value.get("store_model_in_db") + if _db_val is True or (isinstance(_db_val, str) and _db_val.lower() == "true"): + store_model_in_db = True + verbose_proxy_logger.info("store_model_in_db=True loaded from DB, overriding config/env") except Exception as e: verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e)) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 14196d4a1ac..0fef62a4e9d 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1160,6 +1160,171 @@ 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_recovers_from_transient_read_error(monkeypatch): + """ + Regression (#32412): the startup `store_model_in_db` lookup reads + general_settings straight off the Prisma HTTP engine. A transient transport + blip (httpx.ReadError) on that first read must be healed by one + reconnect-and-retry, like every other DB read path, instead of leaving the + connection poisoned and the DB-persisted store_model_in_db=True ignored. + """ + import types + + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + 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 + + record = types.SimpleNamespace(param_value={"store_model_in_db": True}) + find_first = AsyncMock(side_effect=[httpx.ReadError("Server disconnected"), record]) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = find_first + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_config = _mock_scheduled_proxy_config() + + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + 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, + ) + + assert scheduler.get_job("add_deployment_job") is not None + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + assert find_first.await_count == 2 + mock_proxy_config.add_deployment.assert_awaited() + mock_proxy_config.init_mcp_servers_from_db.assert_not_called() + finally: + scheduler.shutdown(wait=False) + + +@pytest.mark.asyncio +async def test_initialize_scheduled_jobs_normalizes_string_general_settings(monkeypatch): + """ + Regression (#32412): Prisma can return the general_settings JSON column as a + raw string rather than a parsed dict. The startup store_model_in_db lookup + must json.loads that string before reading the flag, otherwise a + DB-persisted store_model_in_db=True is silently skipped for string-valued + rows and the proxy comes up as if model storage were off. + """ + import types + + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + 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 + + record = types.SimpleNamespace(param_value='{"store_model_in_db": true}') + find_first = AsyncMock(return_value=record) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = find_first + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_config = _mock_scheduled_proxy_config() + + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + 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, + ) + + assert scheduler.get_job("add_deployment_job") is not None + assert find_first.await_count == 1 + mock_prisma_client.attempt_db_reconnect.assert_not_awaited() + mock_proxy_config.add_deployment.assert_awaited() + mock_proxy_config.init_mcp_servers_from_db.assert_not_called() + finally: + scheduler.shutdown(wait=False) + + +@pytest.mark.asyncio +async def test_initialize_scheduled_jobs_ignores_malformed_string_general_settings(monkeypatch): + """ + A string-valued general_settings row that is not valid JSON must be treated + as empty, leaving store_model_in_db off, instead of failing the lookup. + """ + import types + + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + 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 + + record = types.SimpleNamespace(param_value="{not json") + find_first = AsyncMock(return_value=record) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = find_first + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_config = _mock_scheduled_proxy_config() + + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + patch("litellm.proxy.proxy_server.verbose_proxy_logger") as mock_logger, + ): + 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, + ) + + assert scheduler.get_job("add_deployment_job") is None + assert find_first.await_count == 1 + assert not any( + "Failed to check DB for store_model_in_db" in str(call.args[0]) for call in mock_logger.debug.call_args_list + ) + mock_proxy_config.add_deployment.assert_not_called() + mock_proxy_config.init_mcp_servers_from_db.assert_awaited_once() + finally: + scheduler.shutdown(wait=False) + + @pytest.mark.asyncio async def test_init_mcp_servers_from_db_respects_supported_db_objects(monkeypatch): """