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():