diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 02f9c9bef2..493519f232 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -211,6 +211,7 @@ from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -6750,27 +6751,64 @@ class ProxyStartupEvent: "budget_duration not set on Proxy. budget_duration is required to use max_budget." ) - # add proxy budget to db in the user table asyncio.create_task( - generate_key_helper_fn( # type: ignore - request_type="user", - table_name="user", - user_id=litellm_proxy_budget_name, - duration=None, - models=[], - aliases={}, - config={}, - spend=0, - max_budget=litellm.max_budget, - budget_duration=litellm.budget_duration, - query_type="update_data", - update_key_values={ - "max_budget": litellm.max_budget, - "budget_duration": litellm.budget_duration, - }, - ) + cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name) ) + @classmethod + async def _upsert_proxy_budget_with_reset_at_backfill( + cls, litellm_proxy_budget_name: str + ) -> None: + """ + Upsert the proxy admin user row with the configured max_budget / + budget_duration, then backfill budget_reset_at if currently NULL. + + The backfill uses `WHERE budget_reset_at IS NULL` so it only fires + when the row pre-existed without a reset schedule (e.g. row created + via a different path before the proxy budget was configured). On + subsequent restarts it no-ops, so an active reset window is never + slid forward. + """ + await generate_key_helper_fn( # type: ignore + request_type="user", + table_name="user", + user_id=litellm_proxy_budget_name, + duration=None, + models=[], + aliases={}, + config={}, + spend=0, + max_budget=litellm.max_budget, + budget_duration=litellm.budget_duration, + query_type="update_data", + update_key_values={ + "max_budget": litellm.max_budget, + "budget_duration": litellm.budget_duration, + }, + ) + + # Without this, the upsert leaves budget_reset_at=NULL on rows that + # took the UPDATE path, and reset_budget_for_litellm_users never + # matches them (NULL < now() is unknown in SQL) — so the proxy-wide + # spend cap blocks forever once it's hit. + if prisma_client is not None and litellm.budget_duration is not None: + try: + await prisma_client.db.litellm_usertable.update_many( + where={ + "user_id": litellm_proxy_budget_name, + "budget_reset_at": None, + }, + data={ + "budget_reset_at": get_budget_reset_time( + budget_duration=litellm.budget_duration + ) + }, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to backfill budget_reset_at on proxy admin row: %s", e + ) + @classmethod async def _warm_global_spend_cache( cls, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6718f52cbf..859594f7a0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1728,6 +1728,67 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): assert call_args.kwargs["query_type"] == "update_data" +@pytest.mark.asyncio +async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): + """ + Test that _upsert_proxy_budget_with_reset_at_backfill issues a conditional + update_many with `WHERE budget_reset_at IS NULL` to backfill the column on + rows that pre-existed without a reset schedule. Without this, the proxy + admin row stays at NULL and reset_budget_for_litellm_users never matches + it (NULL < now() is unknown in SQL), so the global proxy budget never + resets. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + import litellm + from litellm.proxy.proxy_server import ProxyStartupEvent + + litellm.budget_duration = "30d" + litellm.max_budget = 100.0 + litellm_proxy_budget_name = "litellm-proxy-budget" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.update_many = AsyncMock(return_value={"count": 1}) + + mock_generate_key_helper = AsyncMock( + return_value={ + "user_id": litellm_proxy_budget_name, + "max_budget": 100.0, + "budget_duration": "30d", + "spend": 0, + "models": [], + } + ) + + with ( + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + mock_generate_key_helper, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + await ProxyStartupEvent._upsert_proxy_budget_with_reset_at_backfill( + litellm_proxy_budget_name + ) + + # Upsert ran with the configured budget + mock_generate_key_helper.assert_called_once() + + # Backfill update_many ran with the conditional WHERE + mock_prisma.db.litellm_usertable.update_many.assert_called_once() + backfill_call = mock_prisma.db.litellm_usertable.update_many.call_args + assert backfill_call.kwargs["where"]["user_id"] == litellm_proxy_budget_name + assert backfill_call.kwargs["where"]["budget_reset_at"] is None + + # The backfilled value must be a real future datetime — anything else and + # reset_budget_for_litellm_users would still skip the row. + from datetime import datetime, timezone + + backfilled_reset_at = backfill_call.kwargs["data"]["budget_reset_at"] + assert isinstance(backfilled_reset_at, datetime) + assert backfilled_reset_at > datetime.now(timezone.utc) + + @pytest.mark.asyncio async def test_custom_ui_sso_sign_in_handler_config_loading(): """