From fe5cc1eb0cd9cb4cdbcff4c3d78fb5bf262eca26 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 25 Jul 2026 11:29:39 -0700 Subject: [PATCH] fix(proxy): global max_budget ignores budget_duration; enforce against the resettable proxy budget row (#33732) * fix(proxy): enforce global max_budget against the resettable proxy budget row The global proxy budget check compared litellm.max_budget against SUM(spend) from the MonthlyGlobalSpend view, whose window is hardcoded to a trailing 30 days. litellm.budget_duration was stored and reset on a user row that enforcement never read, and startup budgeted the admin user's own row (default_user_id) instead of the litellm-proxy-budget aggregate row the spend writer increments per request. Net effect: 1d, 7d and 30d all behaved as a trailing 30 day cap that never reset on the configured duration. Startup now upserts the budget onto the litellm-proxy-budget row (and zeroes lifetime accrual when first putting a row on a reset schedule), enforcement loads global spend from that row, and ResetBudgetJob drops the cached global spend accumulator when it resets that row so the cap unblocks immediately after each window. Fixes https://github.com/BerriAI/litellm/issues/31292 * refactor(proxy): address review nits on global proxy budget fix Drop the redundant litellm_proxy_budget_name parameter from _upsert_proxy_budget_with_reset_at_backfill; its only caller always passed LITELLM_PROXY_BUDGET_NAME, and any other value would write the budget to a row enforcement never reads. Introduce GLOBAL_PROXY_SPEND_CACHE_KEY in constants.py and use it at every site that previously built the key from litellm_proxy_admin_name (auth loads, spend-writer increments, startup warm, reset-job invalidation), so the reader and invalidator can no longer drift apart. The literal key value is unchanged. Also drop the now-pointless litellm_proxy_admin_name parameter from _warm_global_spend_cache and the proxy_server import from the reset-job helper. --- litellm/constants.py | 2 + litellm/proxy/auth/user_api_key_auth.py | 21 ++++--- .../proxy/common_utils/reset_budget_job.py | 11 ++++ litellm/proxy/proxy_server.py | 42 +++++++------ .../proxy/auth/test_user_api_key_auth.py | 54 +++++++++++++++++ .../common_utils/test_reset_budget_job.py | 60 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 19 ++++-- 7 files changed, 178 insertions(+), 31 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 62d351cffe9..a9edf135731 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1418,6 +1418,8 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int( os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000) ) LITELLM_PROXY_ADMIN_NAME = "default_user_id" +LITELLM_PROXY_BUDGET_NAME = "litellm-proxy-budget" +GLOBAL_PROXY_SPEND_CACHE_KEY = f"{LITELLM_PROXY_ADMIN_NAME}:spend" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli" diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e368a1a2275..e7c12448435 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -23,7 +23,11 @@ from fastapi.security.api_key import APIKeyHeader import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS +from litellm.constants import ( + GLOBAL_PROXY_SPEND_CACHE_KEY, + LITELLM_PROXY_BUDGET_NAME, + LITELLM_PROXY_MASTER_KEY_ALIAS, +) from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer @@ -500,13 +504,16 @@ async def _fetch_global_spend_with_event_coordination( """ Fetch global spend with event-driven coordination to prevent cache stampede. Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done. + + Reads the proxy budget aggregate user row, which accrues proxy-wide spend + per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``. """ async def _load_global_spend() -> Optional[float]: - sql_query = """SELECT SUM(spend) AS total_spend FROM "MonthlyGlobalSpend";""" - response = await prisma_client.db.query_raw(query=sql_query) - val = response[0]["total_spend"] - return float(val) if val is not None else None + proxy_budget_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": LITELLM_PROXY_BUDGET_NAME} + ) + return float(proxy_budget_row.spend) if proxy_budget_row is not None else None return await _global_spend_coordinator.get_or_load( cache_key=cache_key, @@ -525,7 +532,7 @@ async def get_global_proxy_spend( global_proxy_spend = None if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget # Use event-driven coordination to prevent cache stampede - cache_key = "{}:spend".format(litellm_proxy_admin_name) + cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY global_proxy_spend = await _fetch_global_spend_with_event_coordination( cache_key=cache_key, user_api_key_cache=user_api_key_cache, @@ -1979,7 +1986,7 @@ async def _user_api_key_auth_builder( global_proxy_spend = None if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget - cache_key = "{}:spend".format(litellm_proxy_admin_name) + cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"): global_proxy_spend = await _fetch_global_spend_with_event_coordination( cache_key=cache_key, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 23a5b8f9c53..1210192416f 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -6,6 +6,7 @@ from typing import Any, Callable, List, Literal, Optional, Union import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, @@ -98,6 +99,14 @@ class ResetBudgetJob: except Exception as e: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) + @staticmethod + async def _invalidate_global_proxy_spend_cache() -> None: + """Drop the cached global-proxy spend accumulator after the proxy + budget aggregate row is reset, so the next auth-time load reads the + zeroed row instead of a stale (potentially never-expiring) counter. + """ + await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY) + @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: """Drop a stale management-cache entry so the next read fetches from DB. @@ -553,6 +562,8 @@ class ResetBudgetJob: user_id = getattr(u, "user_id", None) if user_id: await self._invalidate_spend_counter(f"spend:user:{user_id}") + if user_id == LITELLM_PROXY_BUDGET_NAME: + await self._invalidate_global_proxy_spend_cache() end_time = time.time() if len(failed_users) > 0: # If any users failed to reset diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a20b557e38b..c44f5602c95 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -230,7 +230,9 @@ from litellm.constants import ( DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, + GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, + LITELLM_PROXY_BUDGET_NAME, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, PROXY_BATCH_POLLING_ENABLED, PROXY_BATCH_POLLING_INTERVAL, @@ -1054,10 +1056,9 @@ async def proxy_startup_event(app: FastAPI): 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(litellm_proxy_budget_name=litellm_proxy_admin_name) + ProxyStartupEvent._add_proxy_budget_to_db() asyncio.create_task( ProxyStartupEvent._warm_global_spend_cache( - litellm_proxy_admin_name=litellm_proxy_admin_name, user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, ) @@ -2002,7 +2003,7 @@ health_check_results: Dict[str, Union[int, List[Dict[str, Any]]]] = {} background_health_check_loop_active = False background_health_check_cycle_seq = 0 queue: List = [] -litellm_proxy_budget_name = "litellm-proxy-budget" +litellm_proxy_budget_name = LITELLM_PROXY_BUDGET_NAME litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME ui_access_mode: Union[Literal["admin", "all"], Dict] = "all" proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME @@ -2871,15 +2872,13 @@ async def update_cache( ) ) ## UPDATE GLOBAL PROXY ## - global_proxy_spend = await user_api_key_cache.async_get_cache( - key="{}:spend".format(litellm_proxy_admin_name) - ) + global_proxy_spend = await user_api_key_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) if global_proxy_spend is None: # do nothing if not in cache return elif response_cost is not None and global_proxy_spend is not None: increment = global_proxy_spend + response_cost - values_to_update_in_cache.append(("{}:spend".format(litellm_proxy_admin_name), increment)) + values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment)) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update user spend in cache. " @@ -3040,7 +3039,7 @@ async def update_cache( if tags is not None: await _update_tag_cache() - global_proxy_spend_key = "{}:spend".format(litellm_proxy_admin_name) + global_proxy_spend_key = GLOBAL_PROXY_SPEND_CACHE_KEY local_object_updates = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key) shared_scalar_updates = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key) @@ -7772,29 +7771,32 @@ class ProxyStartupEvent: ) @classmethod - def _add_proxy_budget_to_db(cls, litellm_proxy_budget_name: str): + def _add_proxy_budget_to_db(cls): """Adds a global proxy budget to db""" if litellm.budget_duration is None: raise Exception("budget_duration not set on Proxy. budget_duration is required to use max_budget.") - asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name)) + asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill()) @classmethod - async def _upsert_proxy_budget_with_reset_at_backfill(cls, litellm_proxy_budget_name: str) -> None: + async def _upsert_proxy_budget_with_reset_at_backfill(cls) -> None: """ - Upsert the proxy admin user row with the configured max_budget / - budget_duration, then backfill budget_reset_at if currently NULL. + Upsert the proxy budget aggregate 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. + slid forward. It also zeroes spend at that moment: a row that was + never on a reset schedule holds lifetime accrual, which must not + gate the first duration window. """ await generate_key_helper_fn( # type: ignore request_type="user", table_name="user", - user_id=litellm_proxy_budget_name, + user_id=LITELLM_PROXY_BUDGET_NAME, duration=None, models=[], aliases={}, @@ -7817,10 +7819,13 @@ class ProxyStartupEvent: try: await UserRepository(prisma_client).table.update_many( where={ - "user_id": litellm_proxy_budget_name, + "user_id": LITELLM_PROXY_BUDGET_NAME, "budget_reset_at": None, }, - data={"budget_reset_at": get_budget_reset_time(budget_duration=litellm.budget_duration)}, + data={ + "budget_reset_at": get_budget_reset_time(budget_duration=litellm.budget_duration), + "spend": 0, + }, ) except Exception as e: verbose_proxy_logger.warning("Failed to backfill budget_reset_at on proxy admin row: %s", e) @@ -7828,13 +7833,12 @@ class ProxyStartupEvent: @classmethod async def _warm_global_spend_cache( cls, - litellm_proxy_admin_name: str, user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> None: """Warm global spend cache once at startup to reduce impact of first wave of requests.""" try: - cache_key = "{}:spend".format(litellm_proxy_admin_name) + cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY await _fetch_global_spend_with_event_coordination( cache_key=cache_key, user_api_key_cache=user_api_key_cache, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index c0c91ae103f..ca2e282a119 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5059,6 +5059,60 @@ class TestCheckKeyModelBudgetWithFallback: assert "model" not in request_data +@pytest.mark.asyncio +async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): + """Regression for the global proxy budget ignoring budget_duration + (LIT-4309 / gh#31292): the enforced global spend must be loaded from the + "litellm-proxy-budget" user row, which the spend writer increments per + request and ResetBudgetJob zeroes every budget_duration. It must NOT be + loaded from the MonthlyGlobalSpend view, whose window is hardcoded to a + trailing 30 days and never resets on the configured duration.""" + from litellm.proxy.auth.user_api_key_auth import ( + _fetch_global_spend_with_event_coordination, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_budget_row = MagicMock() + proxy_budget_row.spend = 42.5 + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row) + prisma_client.db.query_raw = AsyncMock( + side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view") + ) + + result = await _fetch_global_spend_with_event_coordination( + cache_key="default_user_id:spend", + user_api_key_cache=UserApiKeyCache(), + prisma_client=prisma_client, + ) + + assert result == 42.5 + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with( + where={"user_id": "litellm-proxy-budget"} + ) + + +@pytest.mark.asyncio +async def test_global_proxy_spend_none_when_proxy_budget_row_missing(): + """Before the startup upsert creates the aggregate row, enforcement must + see None (no cap applied) rather than raising.""" + from litellm.proxy.auth.user_api_key_auth import ( + _fetch_global_spend_with_event_coordination, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + result = await _fetch_global_spend_with_event_coordination( + cache_key="default_user_id:spend", + user_api_key_cache=UserApiKeyCache(), + prisma_client=prisma_client, + ) + + assert result is None + + @pytest.mark.asyncio async def test_temp_budget_increase_applied_for_cached_key(): """ diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index be5bc74c385..f04d6f3cf5a 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1367,6 +1367,66 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60) +def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Regression for LIT-4309: resetting the proxy-wide budget aggregate row + ("litellm-proxy-budget") must also drop the cached global-spend + accumulator ("{admin}:spend") that _global_proxy_budget_check enforces + against. Without the invalidation, the cached value survives the DB reset + and the global cap keeps blocking requests for the whole next window.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 150.0, + "budget_duration": "30d", + "budget_reset_at": now, + "id": "row-1", + "user_id": "litellm-proxy-budget", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + counter_cache.user_api_key_cache.async_delete_cache.assert_any_call(key="default_user_id:spend") + + +def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache( + reset_budget_job, mock_prisma_client, monkeypatch +): + """The global-spend accumulator must only be dropped when the proxy + budget aggregate row itself resets, not on every user reset.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 50.0, + "budget_duration": "7d", + "budget_reset_at": now, + "id": "user-1", + "user_id": "alice", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + assert not any( + call.kwargs.get("key") == "default_user_id:spend" + for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list + ) + + def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Team budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 087aaec9215..536d24d4b4e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2642,6 +2642,11 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): This validates that generate_key_helper_fn is called with table_name="user" which should prevent key creation in LiteLLM_VerificationToken table. + + Also guards the row identity: the budget must land on the proxy-wide + aggregate row "litellm-proxy-budget" (the one the spend writer increments + per request), not the admin user's own row ("default_user_id"). Budgeting + the admin row leaves the global budget without a resettable counter. """ from unittest.mock import AsyncMock, patch @@ -2670,7 +2675,7 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): "litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper ): # Call the function under test - ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name) + ProxyStartupEvent._add_proxy_budget_to_db() # Allow async task to complete import asyncio @@ -2696,9 +2701,13 @@ 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 + budget 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. + + The same conditional update must zero spend: a row that was never on a + reset schedule holds lifetime accrual, which must not gate the first + duration window. """ from unittest.mock import AsyncMock, MagicMock, patch @@ -2729,9 +2738,7 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), ): - await ProxyStartupEvent._upsert_proxy_budget_with_reset_at_backfill( - litellm_proxy_budget_name - ) + await ProxyStartupEvent._upsert_proxy_budget_with_reset_at_backfill() # Upsert ran with the configured budget mock_generate_key_helper.assert_called_once() @@ -2750,6 +2757,8 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): assert isinstance(backfilled_reset_at, datetime) assert backfilled_reset_at > datetime.now(timezone.utc) + assert backfill_call.kwargs["data"]["spend"] == 0 + @pytest.mark.asyncio async def test_custom_ui_sso_sign_in_handler_config_loading():