fix(proxy): keep team member budget enforced at the cap and across Redis counter expiry (#40304)

* fix(proxy): keep team member budget enforced at the cap and across Redis counter expiry

The cached-key auth path admitted a request when the team member counter sat exactly at max_budget, and a Redis counter that expired during a long stream was reconciled against this pod's stale in-memory copy, driving the shared counter negative and reopening the budget. Reject at >= like every other budget check, read Redis before the per-pod copy when judging the reconcile delta, and add the settled request cost after a DB reseed since reserved keys skip the normal increment

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): reseeded reservation counter also carries the settled request cost

The reseed test asserted counter == DB floor. The floor is read before the async spend flush writes this request, so the counter now lands at floor plus settled cost, matching the after leg in the PR proof

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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,

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,

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(

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],

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,

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