diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 61c618eeb18..241ef00e7eb 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1293,7 +1293,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: + # Note: intentionally do NOT gate on `valid_token.team_member_spend is not None`. + # That field is LEFT-JOINed from LiteLLM_TeamMembership and is None for a + # team member that has never recorded any spend — gating on it would skip the + # authoritative Redis counter read and let the first over-budget request through. + if ( + not skip_budget_checks + and valid_token.team_id is not None + and valid_token.user_id is not None + ): if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" @@ -1330,16 +1338,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend - team_member_spend = valid_token.team_member_spend - if ( - valid_token.user_id is not None - and valid_token.team_id is not None - ): - team_member_spend = await get_current_spend( - counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", - fallback_spend=team_member_spend, - ) - if team_member_spend > team_member_budget: + # valid_token.team_member_spend is None for team members + # with no prior spend row — treat as 0.0 and let the + # Redis counter read be authoritative. + team_member_spend = valid_token.team_member_spend or 0.0 + team_member_spend = await get_current_spend( + counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", + fallback_spend=team_member_spend, + ) + if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..0505092b816 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -64,7 +64,10 @@ class ResetBudgetJob: # Reset Redis directly so a transient failure doesn't leave stale # counters that get_current_spend would read as authoritative. try: - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import ( + SPEND_COUNTER_REDIS_TTL_SECONDS, + spend_counter_cache, + ) memberships = await self.prisma_client.db.litellm_teammembership.find_many( where={"budget_id": {"in": budget_ids}} @@ -73,13 +76,17 @@ class ResetBudgetJob: counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" # Always reset in-memory spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 + key=counter_key, + value=0.0, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, ) # Explicitly reset Redis with warning on failure if spend_counter_cache.redis_cache is not None: try: await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 + key=counter_key, + value=0.0, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, ) except Exception as redis_err: verbose_proxy_logger.warning( @@ -567,7 +574,10 @@ class ResetBudgetJob: # Reset Redis directly (not via DualCache) so a Redis failure # doesn't silently leave a stale counter that get_current_spend # would read as authoritative, permanently blocking the user. - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import ( + SPEND_COUNTER_REDIS_TTL_SECONDS, + spend_counter_cache, + ) counter_key = None if item_type == "key" and hasattr(item, "token") and item.token is not None: @@ -582,13 +592,17 @@ class ResetBudgetJob: if counter_key is not None: # Always reset in-memory (local fallback) spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 + key=counter_key, + value=0.0, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, ) # Explicitly reset Redis with warning on failure if spend_counter_cache.redis_cache is not None: try: await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 + key=counter_key, + value=0.0, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, ) except Exception as redis_err: verbose_proxy_logger.warning( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 85a12f70f58..758004961c5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1540,6 +1540,11 @@ shared_aiohttp_session: Optional["ClientSession"] = ( user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) +# Long TTL for cross-pod spend counter keys. Counters are reset explicitly by +# reset_budget_job, so the TTL is only a safety net for orphaned counters +# (deleted users/teams/keys). Must be larger than any realistic budget period +# (weekly/monthly) to avoid silent mid-cycle expiry. See PR #24682 regression. +SPEND_COUNTER_REDIS_TTL_SECONDS: int = 60 * 60 * 24 * 30 # 30 days spend_counter_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1815,7 +1820,28 @@ async def _init_and_increment_spend_counter( rather than under-counting (would allow overspend). 4. Increment atomically (both in-memory + Redis) """ - current = await spend_counter_cache.async_get_cache(key=counter_key) + # Redis-first: when Redis is configured it is the authoritative cross-pod + # source. A stale per-pod in-memory hit must not short-circuit the seed + # block, or a fresh Redis key (post-TTL-expiry) would start from just + # `increment` rather than base_spend + increment. In-memory is only + # consulted when Redis is absent or errors on read. + current: Optional[float] = None + if spend_counter_cache.redis_cache is not None: + try: + current = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + except Exception as e: + verbose_proxy_logger.debug( + "_init_and_increment_spend_counter: Redis read failed for %s, " + "falling back to in-memory: %s", + counter_key, + e, + ) + current = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + else: + current = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if current is None: source = await user_api_key_cache.async_get_cache(key=source_cache_key) base_spend = 0.0 @@ -1826,10 +1852,16 @@ async def _init_and_increment_spend_counter( base_spend = getattr(source, "spend", 0.0) or 0.0 if base_spend > 0: await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend + key=counter_key, + value=base_spend, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + await spend_counter_cache.async_increment_cache( + key=counter_key, + value=increment, + ttl=SPEND_COUNTER_REDIS_TTL_SECONDS, + ) async def update_cache( # noqa: PLR0915 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 14912b67ab8..0541acb9c43 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 @@ -1431,3 +1431,176 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): finally: for k, v in _originals.items(): setattr(_proxy_server_mod, k, v) + + +async def _run_team_member_budget_auth( + valid_token, + team_member_info, + current_spend: float, +): + """Shared helper: run _user_api_key_auth_builder through the team_member_budget + branch, patching everything up to that branch. Returns after the auth completes + (or raises BudgetExceededError from within the target branch).""" + from starlette.datastructures import URL + from starlette.requests import Request + + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + # Cache stub: return team_member_info for the team-member lookup key only. + # Everything else returns None so the auth flow falls through to patched helpers. + cache_key = f"{valid_token.team_id}_{valid_token.user_id}" + + async def fake_async_get_cache(key=None, **kwargs): + if key == cache_key: + return team_member_info + return None + + mock_cache = AsyncMock() + mock_cache.async_get_cache = AsyncMock(side_effect=fake_async_get_cache) + mock_cache.async_set_cache = AsyncMock(return_value=None) + + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.internal_usage_cache = MagicMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( + AsyncMock() + ) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(return_value=None) + + import litellm.proxy.proxy_server as _proxy_server_mod + + _attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": mock_cache, + "proxy_logging_obj": mock_proxy_logging_obj, + "master_key": "sk-master-key", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + _originals = {k: getattr(_proxy_server_mod, k, None) for k in _attrs} + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + try: + for k, v in _attrs.items(): + setattr(_proxy_server_mod, k, v) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + new_callable=AsyncMock, + return_value=valid_token, + ), patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=None, + ), patch( + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + return_value=None, + ), patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + return_value=True, + ), patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=current_spend, + ): + return await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {valid_token.token}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in _originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_team_member_budget_enforced_when_team_member_spend_is_none(): + """Regression: user_api_key_auth must read the Redis spend counter even when + valid_token.team_member_spend is None (fresh team member, no DB spend row). + Pre-fix, the `is not None` gate skipped the entire check, letting the first + over-budget request through.""" + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_TeamMembership, + LitellmUserRoles, + ) + + valid_token = UserAPIKeyAuth( + token="sk-test-none-spend-regression", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-xyz", + user_id="user-xyz", + team_member_spend=None, # fresh team member, no DB spend row + ) + + team_member_info = LiteLLM_TeamMembership( + user_id="user-xyz", + team_id="team-xyz", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + # The auth exception handler wraps BudgetExceededError as ProxyException + # with type=budget_exceeded, so that's what the caller actually sees. + with pytest.raises(ProxyException) as exc_info: + await _run_team_member_budget_auth( + valid_token=valid_token, + team_member_info=team_member_info, + current_spend=5.0, # Redis counter is over the $1 cap + ) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert "Budget has been exceeded" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_team_member_budget_blocks_at_exact_boundary(): + """Regression: user_api_key_auth must block at spend == budget (>=), matching + _check_team_member_budget in auth_checks.py. Pre-fix, the `>` comparison let + the exact boundary through, creating an off-by-one at rollover.""" + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_TeamMembership, + LitellmUserRoles, + ) + + valid_token = UserAPIKeyAuth( + token="sk-test-boundary-regression", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-boundary", + user_id="user-boundary", + team_member_spend=0.5, + ) + + team_member_info = LiteLLM_TeamMembership( + user_id="user-boundary", + team_id="team-boundary", + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_team_member_budget_auth( + valid_token=valid_token, + team_member_info=team_member_info, + current_spend=1.0, # exactly at the budget + ) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index daabed0def1..0c26c28649e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4925,3 +4925,168 @@ async def test_increment_spend_counters_team_and_member(): finally: ps.user_api_key_cache = original_key_cache ps.spend_counter_cache = original_counter_cache + + +@pytest.mark.asyncio +async def test_increment_spend_counters_sets_long_redis_ttl(): + """Regression for PR #24682 TTL bug: counter writes must pass an explicit + long TTL, not rely on BaseCache.default_ttl=60 which silently expires the + counter every minute and reseeds it from stale DB state.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_VerificationTokenView, hash_token + + key_cache = DualCache() + counter_cache = DualCache() + + # Mock Redis that records every async_increment call with its ttl kwarg. + recorded_ttls: list = [] + + async def fake_async_increment(key, value, ttl=None, **kwargs): + recorded_ttls.append(ttl) + return value + + mock_redis = AsyncMock() + mock_redis.async_increment = fake_async_increment + mock_redis.async_get_cache = AsyncMock(return_value=None) + counter_cache.redis_cache = mock_redis + + hashed_token = hash_token("sk-ttl-regression") + key_cache.in_memory_cache.set_cache( + key=hashed_token, + value=LiteLLM_VerificationTokenView( + token=hashed_token, spend=3.0, max_budget=10.0 + ), + ) + + import litellm.proxy.proxy_server as ps + + original_key_cache = ps.user_api_key_cache + original_counter_cache = ps.spend_counter_cache + ps.user_api_key_cache = key_cache + ps.spend_counter_cache = counter_cache + try: + from litellm.proxy.proxy_server import ( + SPEND_COUNTER_REDIS_TTL_SECONDS, + increment_spend_counters, + ) + + await increment_spend_counters( + token=hashed_token, team_id=None, user_id=None, response_cost=0.25 + ) + + # First call seeds from base_spend (3.0), second increments by 0.25. + # Both writes must carry the long TTL. + assert len(recorded_ttls) == 2 + assert all(t == SPEND_COUNTER_REDIS_TTL_SECONDS for t in recorded_ttls) + # And definitely not the 60s default that caused the bug. + assert 60 not in recorded_ttls + finally: + ps.user_api_key_cache = original_key_cache + ps.spend_counter_cache = original_counter_cache + + +@pytest.mark.asyncio +async def test_init_and_increment_redis_first_init_check(): + """Seed block must not be short-circuited by a stale in-memory hit. + After a Redis TTL expiry, a pod that still has an in-memory value must + re-read Redis (empty), fall through to the seed block, and re-initialize + from team_membership.spend — otherwise the new Redis key would start at + just the current increment, dropping all prior accumulation.""" + from litellm.caching.dual_cache import DualCache + + key_cache = DualCache() + counter_cache = DualCache() + + # Stale in-memory value on the counter (simulates pre-expiry state). + counter_cache.in_memory_cache.set_cache(key="spend:team_member:u1:t1", value=7.0) + + # Redis is empty (simulates post-TTL-expiry state on every pod). + mock_redis = AsyncMock() + mock_redis.async_get_cache = AsyncMock(return_value=None) + incremented_values: list = [] + + async def fake_async_increment(key, value, ttl=None, **kwargs): + incremented_values.append(value) + return value + + mock_redis.async_increment = fake_async_increment + counter_cache.redis_cache = mock_redis + + # DB-loaded team_membership with authoritative spend. + key_cache.in_memory_cache.set_cache( + key="team_membership:u1:t1", + value={"user_id": "u1", "team_id": "t1", "spend": 70.0}, + ) + + import litellm.proxy.proxy_server as ps + + original_key_cache = ps.user_api_key_cache + original_counter_cache = ps.spend_counter_cache + ps.user_api_key_cache = key_cache + ps.spend_counter_cache = counter_cache + try: + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, team_id="t1", user_id="u1", response_cost=0.50 + ) + + # Must seed with 70.0 THEN increment with 0.50 — not just 0.50. + assert 70.0 in incremented_values + assert 0.50 in incremented_values + finally: + ps.user_api_key_cache = original_key_cache + ps.spend_counter_cache = original_counter_cache + + +@pytest.mark.asyncio +async def test_reset_budget_for_team_members_sets_long_ttl(): + """Budget reset must pass the long TTL explicitly — otherwise the freshly + reset counter inherits the 60s Redis default and the TTL bug reappears + immediately after every reset cycle.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + + counter_cache = DualCache() + recorded_redis: list = [] + + async def fake_async_set_cache(key, value, ttl=None, **kwargs): + recorded_redis.append((key, value, ttl)) + + mock_redis = AsyncMock() + mock_redis.async_set_cache = fake_async_set_cache + counter_cache.redis_cache = mock_redis + + import litellm.proxy.proxy_server as ps + + original_counter_cache = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + + # Minimal prisma mock — one team member linked to the reset budget. + mock_prisma = MagicMock() + mock_membership = MagicMock(user_id="u1", team_id="t1") + mock_prisma.db.litellm_teammembership.find_many = AsyncMock( + return_value=[mock_membership] + ) + mock_prisma.db.litellm_teammembership.update_many = AsyncMock() + + # Minimal budget mock — reset path only reads .budget_id + mock_budget = MagicMock(budget_id="b1") + + try: + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=mock_prisma) + await job.reset_budget_for_litellm_team_members(budgets_to_reset=[mock_budget]) + + from litellm.proxy.proxy_server import SPEND_COUNTER_REDIS_TTL_SECONDS + + assert recorded_redis == [ + ("spend:team_member:u1:t1", 0.0, SPEND_COUNTER_REDIS_TTL_SECONDS) + ] + # In-memory reset must also land at 0.0 (ttl is an override hint for + # in-memory; the value is what the gate reads). + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:team_member:u1:t1") + == 0.0 + ) + finally: + ps.spend_counter_cache = original_counter_cache