From 45b55bd60ba388ff0c60f3f9a95cc4cfbae68632 Mon Sep 17 00:00:00 2001 From: Ashish Soni Date: Tue, 4 Aug 2026 01:56:10 +0530 Subject: [PATCH] fix(proxy): clear stale proxy budget from db when max_budget removed from config Proxy startup syncs litellm_settings.max_budget onto the proxy budget rows in LiteLLM_UserTable, but only did so forward: when max_budget is removed from config, litellm.max_budget resolves to 0 and no sync runs at all. The last-synced value keeps enforcing against master-key traffic, which resolves to default_user_id, until spend hits the stale cap and every request 429s with "Max budget limit reached" even though the global proxy budget check reports disabled. ProxyStartupEvent._sync_proxy_budget_state now branches on litellm.max_budget: configured syncs as before, unset clears max_budget/budget_duration/budget_reset_at from both the current sync row (litellm-proxy-budget) and the legacy sync row (default_user_id, written by versions before #33732) and evicts both from the user cache so a stale Redis-cached value can't keep serving the old cap. Fixes #35680 --- litellm/proxy/proxy_server.py | 56 +++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 98 +++++++++++++++++++ 2 files changed, 147 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 901ca39326b..5d07dcb719e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1062,13 +1062,10 @@ async def proxy_startup_event(app: FastAPI): 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, ) ### START BATCH WRITING DB + CHECKING NEW MODELS### @@ -7917,6 +7914,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""" @@ -7925,6 +7944,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 ab64224fdc3..8d8ee3b771c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2764,6 +2764,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(): """