diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dad7abd210..85066837ffe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1225,13 +1225,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: prompt_injection_detection_obj.update_environment(router=llm_router) verbose_proxy_logger.debug("prisma_client: %s", prisma_client) - if prisma_client is not None and litellm.max_budget > 0: - ProxyStartupEvent._add_proxy_budget_to_db() - asyncio.create_task( - ProxyStartupEvent._warm_global_spend_cache( - user_api_key_cache=user_api_key_cache, - prisma_client=prisma_client, - ) + if prisma_client is not None: + ProxyStartupEvent._sync_proxy_budget_state( + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, ) ProxyStartupEvent._warn_budget_without_db( max_budget=litellm.max_budget, @@ -8872,6 +8869,28 @@ class ProxyStartupEvent: litellm_jwtauth=litellm_jwtauth, ) + @classmethod + def _sync_proxy_budget_state( + cls, + user_api_key_cache: UserApiKeyCache, + prisma_client: PrismaClient, + ) -> None: + if litellm.max_budget > 0: + cls._add_proxy_budget_to_db() + asyncio.create_task( + cls._warm_global_spend_cache( + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + ) + else: + asyncio.create_task( + cls._clear_stale_proxy_budget_from_db( + user_api_key_cache=user_api_key_cache, + prisma_client=prisma_client, + ) + ) + @classmethod def _add_proxy_budget_to_db(cls): """Adds a global proxy budget to db""" @@ -8880,6 +8899,29 @@ class ProxyStartupEvent: asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill()) + @classmethod + async def _clear_stale_proxy_budget_from_db( + cls, + user_api_key_cache: UserApiKeyCache, + prisma_client: PrismaClient, + ) -> None: + try: + await UserRepository(prisma_client).table.update_many( + where={ + "user_id": {"in": [LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_ADMIN_NAME]}, + "max_budget": {"not": None}, + }, + data={ + "max_budget": None, + "budget_duration": None, + "budget_reset_at": None, + }, + ) + await user_api_key_cache.async_delete_cache(key=LITELLM_PROXY_BUDGET_NAME) + await user_api_key_cache.async_delete_cache(key=LITELLM_PROXY_ADMIN_NAME) + except Exception as e: + verbose_proxy_logger.warning("Failed to clear stale proxy budget rows: %s", e) + @classmethod async def _upsert_proxy_budget_with_reset_at_backfill(cls) -> None: """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2c8ae9f0370..6bfabbe6d91 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2934,6 +2934,104 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): assert backfill_call.kwargs["data"]["spend"] == 0 +@pytest.mark.asyncio +async def test_sync_proxy_budget_state_clears_stale_budget_when_max_budget_unset(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/35680. + + When litellm_settings.max_budget is removed from config (litellm.max_budget + resolves to 0), startup must clear any max_budget/budget_duration previously + synced onto the proxy budget rows, on both the current sync target + ("litellm-proxy-budget") and the legacy sync target ("default_user_id"). + Without this, master-key traffic (which resolves to default_user_id) keeps + getting 429'd against a budget that no longer exists in config. + """ + from litellm.proxy.proxy_server import ProxyStartupEvent + + original_max_budget = litellm.max_budget + litellm.max_budget = 0.0 + + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock(return_value={"count": 2}) + mock_cache = MagicMock() + mock_cache.async_delete_cache = AsyncMock() + + try: + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + ProxyStartupEvent._sync_proxy_budget_state( + user_api_key_cache=mock_cache, + prisma_client=mock_prisma, + ) + await asyncio.sleep(0.1) + + mock_prisma.db.litellm_usertable.update_many.assert_called_once() + clear_call = mock_prisma.db.litellm_usertable.update_many.call_args + assert set(clear_call.kwargs["where"]["user_id"]["in"]) == { + "litellm-proxy-budget", + "default_user_id", + } + assert clear_call.kwargs["where"]["max_budget"] == {"not": None} + assert clear_call.kwargs["data"] == { + "max_budget": None, + "budget_duration": None, + "budget_reset_at": None, + } + + deleted_keys = {call.kwargs["key"] for call in mock_cache.async_delete_cache.call_args_list} + assert deleted_keys == {"litellm-proxy-budget", "default_user_id"} + finally: + litellm.max_budget = original_max_budget + + +@pytest.mark.asyncio +async def test_sync_proxy_budget_state_syncs_when_max_budget_configured(): + """ + When litellm.max_budget is set, startup must sync the budget onto the db + (existing behavior) and must NOT run the stale-budget clear path. + """ + from litellm.proxy.proxy_server import ProxyStartupEvent + + original_max_budget = litellm.max_budget + original_budget_duration = litellm.budget_duration + litellm.max_budget = 100.0 + litellm.budget_duration = "30d" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock(return_value={"count": 0}) + mock_cache = MagicMock() + mock_cache.async_delete_cache = AsyncMock() + + mock_generate_key_helper = AsyncMock( + return_value={ + "user_id": "litellm-proxy-budget", + "max_budget": 100.0, + "budget_duration": "30d", + "spend": 0, + "models": [], + } + ) + + try: + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + mock_generate_key_helper, + ), + ): + ProxyStartupEvent._sync_proxy_budget_state( + user_api_key_cache=mock_cache, + prisma_client=mock_prisma, + ) + await asyncio.sleep(0.1) + + mock_generate_key_helper.assert_called_once() + mock_cache.async_delete_cache.assert_not_called() + finally: + litellm.max_budget = original_max_budget + litellm.budget_duration = original_budget_duration + + @pytest.mark.asyncio async def test_custom_ui_sso_sign_in_handler_config_loading(): """