diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b39b1f330b3..d8679627bc3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2038,7 +2038,7 @@ async def _user_api_key_auth_builder( fallback_spend=team_member_spend, max_budget=team_member_budget, ) - if team_member_spend > team_member_budget: + if team_member_spend >= team_member_budget: _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 38c4c2e96c9..238fd841ce1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2582,10 +2582,11 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None ) -async def reseed_spend_counter_from_db(counter_key: str) -> None: +async def reseed_spend_counter_from_db(counter_key: str) -> bool: """Recover a counter that the reservation reconcile found in an inconsistent state (missing, or where applying the reconcile delta would drive it - negative) by reseeding it from the DB instead of deleting it. + negative) by reseeding it from the DB instead of deleting it. Returns + whether a DB row was found and the counter was reseeded. The DB row is a LAGGING authoritative floor, not post-request truth: the entity .spend column is flushed in batches (every PROXY_BATCH_WRITE_AT), so @@ -2600,8 +2601,9 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None: """ db_spend: Final = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) if db_spend is None: - return + return False await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend) + return True async def _floor_spend_from_db( @@ -2663,21 +2665,14 @@ async def _authoritative_floor_spend( return db_spend -async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]: - """Return (spend, authoritative). ``authoritative`` is True when the value - came from Redis or a fresh DB read (cross-pod truth), False when it came - from the per-pod in-memory copy or the caller's fallback. Only the - fail-closed path reads the flag; normal callers ignore it.""" - # 1. Redis first (cross-pod authoritative). On clean miss, skip - # in-memory: per-pod in-memory only has this pod's writes, so it - # would mask cross-pod increments. - redis_clean_miss = False +async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None, bool]: + """Return (value, authoritative) for the live counter, None when absent. A clean + Redis miss is final: the per-pod in-memory copy outlives the Redis TTL and only + holds this pod's writes, so it is consulted only when Redis is unreachable.""" if spend_counter_cache.redis_cache is not None: try: - val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) - if val is not None: - return float(val), True - redis_clean_miss = True + redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + return (float(redis_val) if redis_val is not None else None), True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -2685,13 +2680,20 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) e, ) - # 2. In-memory only when Redis is unreachable. - if not redis_clean_miss: - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val), False + in_memory_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + return (float(in_memory_val) if in_memory_val is not None else None), False - # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. + +async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]: + """Return (spend, authoritative). ``authoritative`` is True when the value + came from Redis or a fresh DB read (cross-pod truth), False when it came + from the per-pod in-memory copy or the caller's fallback. Only the + fail-closed path reads the flag; normal callers ignore it.""" + cached_val, cached_authoritative = await read_spend_counter_cache_value(counter_key=counter_key) + if cached_val is not None: + return cached_val, cached_authoritative + + # Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend: Final = await SpendCounterReseed.coalesced( prisma_client=prisma_client, spend_counter_cache=spend_counter_cache, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 152f3befa15..d9f6e33c43c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -905,13 +905,13 @@ async def _set_reserved_entry_actual_cost( increment=adjustment, ) elif reseed_on_inconsistent: - # Post-call reconcile / release: the counter was flushed or reseeded - # between reservation and reconcile (Redis restart / cross-pod reset), - # so the optimistic delta no longer applies. Recover by reseeding from - # the DB's lagging authoritative floor rather than deleting the counter - # and failing open — deleting it is what left budgets unenforced after a - # Redis reload. - await reseed_spend_counter_from_db(counter_key=counter_key) + # Post-call reconcile / release: the counter was flushed, expired or reseeded + # between reservation and reconcile, so the optimistic delta no longer applies. + # Reseed from the DB floor (which cannot include this request's cost yet) and + # add the settled cost, since increment_spend_counters skips reserved keys. + reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key) + if reseeded and actual_cost > 0: + await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost) else: # Pre-call admission resize: the in-flight reservation cost is not yet # persisted, so the DB floor would discard it. Keep the original @@ -925,18 +925,16 @@ async def _counter_can_apply_adjustment( counter_key: str, adjustment: float, ) -> bool: - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import read_spend_counter_cache_value - current_value: Final = await spend_counter_cache.async_get_cache(key=counter_key) + try: + current_value, _ = await read_spend_counter_cache_value(counter_key=counter_key) + except (TypeError, ValueError): + return False if current_value is None: return False - try: - current_float: Final = float(current_value) - except (TypeError, ValueError): - return False - - return not (adjustment < 0 and current_float + adjustment < -1e-12) + return not (adjustment < 0 and current_value + adjustment < -1e-12) async def _release_applied_entries_best_effort( 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 541aeabcbcd..f1a269cc00a 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 @@ -6767,6 +6767,109 @@ async def test_temp_budget_increase_applied_for_cached_key(): assert cached_after.max_budget == 2.0 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_member_spend, expect_blocked", + [ + (2.4, True), + (2.4000000000000004, True), + (2.39, False), + ], +) +async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked): + """A team member counter sitting exactly at the cap (where a resized reservation + lands it) must be rejected by the cached-key auth path like every other budget check.""" + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.utils import hash_token + + api_key = "sk-team-member-exact-cap" + hashed_token = hash_token(api_key) + team_id = "team-exact-cap" + user_id = "user-exact-cap" + max_budget = 2.4 + + user_api_key_cache = DualCache() + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=UserAPIKeyAuth( + token=hashed_token, + team_id=team_id, + user_id=user_id, + team_member_spend=team_member_spend, + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + await user_api_key_cache.async_set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id), + ) + await user_api_key_cache.async_set_cache( + key=user_id, + value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER), + ) + await user_api_key_cache.async_set_cache( + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=team_member_spend, + budget_id="budget-exact-cap", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget), + ), + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {api_key}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + async def _auth(): + return await _user_api_key_auth_builder( + request=mock_request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, + ) + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True} + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: seed the cached key, team and membership without a DB + "litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), + patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=team_member_spend), + ), + ): + if not expect_blocked: + result = await _auth() + assert result.team_member_spend == team_member_spend + return + with pytest.raises(ProxyException) as exc_info: + await _auth() + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message + + async def _proxy_exception_for_key( api_key: str, general_settings: dict[str, bool], diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index b8fb6170d34..6dab054d8ea 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2198,6 +2198,74 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): assert reservation["finalized"] is True +class _ExpiringRedisCache: + def __init__(self) -> None: + self.store: dict[str, float] = {} + + async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + return self.store.get(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self.store[key] = self.store.get(key, 0.0) + float(value) + return self.store[key] + + async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + return self.store[key] + + async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: + self.store[key] = float(value) + return True + + async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: + self.store.pop(key, None) + + +@pytest.mark.asyncio +async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( + spend_counter_state, +): + """Redis key expired mid-stream while the pod's in-memory copy still holds the + reserved value: reconcile must reseed from the DB floor plus the settled cost + instead of applying ``actual - reserved`` to the empty key.""" + import litellm.proxy.proxy_server as ps + + counter_cache, _ = spend_counter_state + counter_key = "spend:team_member:user-expiry:team-expiry" + redis_cache = _ExpiringRedisCache() + counter_cache.redis_cache = redis_cache + counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + + reservation = { + "reserved_cost": 0.6, + "entries": [ + { + "counter_key": counter_key, + "entity_type": "TeamMember", + "entity_id": "user-expiry:team-expiry", + "reserved_cost": 0.6, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for + ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.3) + ): + await ps.increment_spend_counters( + token="key-expiry", + team_id="team-expiry", + user_id="user-expiry", + response_cost=0.05, + budget_reservation=reservation, + ) + + assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.35) + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_should_invalidate_reserved_counters_after_persisted_spend_failure( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6a901c5e7de..54579e6cb7c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8356,8 +8356,8 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( """When the reservation reconcile finds the counter in an inconsistent state (here: missing), it must NOT delete the counter and fail open (the old behavior, which left the counter unenforced after a Redis reload). It reseeds - from the authoritative DB so the counter reflects the recorded total and - budget gating continues.""" + from the authoritative DB and adds this request's settled cost, which the + async spend flush has not written yet, so budget gating continues.""" from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import increment_spend_counters from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -8394,9 +8394,7 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( ) assert budget_reservation["finalized"] is True - # counter reseeded to the authoritative DB value, not deleted/left None - # and not double-counted via a direct increment - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.85) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma