Merge pull request #26829 from BerriAI/litellm_budgetEnforcementMultiPod

[Fix] Refresh Redis TTL on counter writes, skip stale in-memory in Redis

litellm/caching/dual_cache.py

@@ -392,6 +392,7 @@ class DualCache(BaseCache):
         value: float,
         parent_otel_span: Optional[Span] = None,
         local_only: bool = False,
+        refresh_ttl: bool = False,
         **kwargs,
     ) -> Optional[float]:
         """
@@ -399,6 +400,9 @@ class DualCache(BaseCache):
 
         Value - float - the value you want to increment by
 
+        Refresh_ttl - bool - if True, resets the Redis TTL on every write.
+        Default False preserves window-style semantics.
+
         Returns - the incremented value, or None if no cache backend is
         available (in_memory_cache is None and Redis failed/is absent).
         """
@@ -415,6 +419,7 @@ class DualCache(BaseCache):
                     value,
                     parent_otel_span=parent_otel_span,
                     ttl=kwargs.get("ttl", None),
+                    refresh_ttl=refresh_ttl,
                 )
 
             return result

litellm/caching/redis_cache.py

@@ -824,6 +824,7 @@ class RedisCache(BaseCache):
         value: float,
         ttl: Optional[int] = None,
         parent_otel_span: Optional[Span] = None,
+        refresh_ttl: bool = False,
     ) -> float:
         from redis.asyncio import Redis
 
@@ -834,11 +835,12 @@ class RedisCache(BaseCache):
         try:
             result = await _redis_client.incrbyfloat(name=key, amount=value)
             if _used_ttl is not None:
-                # check if key already has ttl, if not -> set ttl
-                current_ttl = await _redis_client.ttl(key)
-                if current_ttl == -1:
-                    # Key has no expiration
+                if refresh_ttl:
                     await _redis_client.expire(key, _used_ttl)
+                else:
+                    current_ttl = await _redis_client.ttl(key)
+                    if current_ttl == -1:
+                        await _redis_client.expire(key, _used_ttl)
 
             ## LOGGING ##
             end_time = time.time()

litellm/proxy/common_utils/reset_budget_job.py

@@ -52,6 +52,37 @@ class ResetBudgetJob:
             ### RESET MULTI-WINDOW BUDGETS ###
             await self.reset_budget_windows()
 
+    @staticmethod
+    async def _invalidate_spend_counter(counter_key: str) -> None:
+        """Zero a spend counter so a DB-row reset takes effect immediately.
+
+        Call AFTER the DB write commits. Clearing Redis before the DB
+        commit opens a window where get_current_spend reads 0 from Redis
+        while the DB still holds the pre-reset value, allowing bypass.
+        """
+        try:
+            from litellm.proxy.proxy_server import spend_counter_cache
+
+            spend_counter_cache.in_memory_cache.set_cache(
+                key=counter_key, value=0.0, ttl=60
+            )
+            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, ttl=60
+                    )
+                except Exception as redis_err:
+                    verbose_proxy_logger.warning(
+                        "Failed to reset spend counter %s in Redis: %s. "
+                        "Budget may be over-enforced until counter expires.",
+                        counter_key,
+                        redis_err,
+                    )
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                "Failed to reset spend counter %s: %s", counter_key, e
+            )
+
     async def reset_budget_for_litellm_team_members(
         self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
     ):
@@ -64,46 +95,30 @@ class ResetBudgetJob:
             if budget.budget_id is not None
         ]
 
-        # Reset spend counters for affected team members.
-        # 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
-
             memberships = await self.prisma_client.db.litellm_teammembership.find_many(
                 where={"budget_id": {"in": budget_ids}}
             )
-            for m in memberships:
-                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
-                )
-                # 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
-                        )
-                    except Exception as redis_err:
-                        verbose_proxy_logger.warning(
-                            "Failed to reset team member spend counter in Redis %s: %s. "
-                            "Budget may be over-enforced until counter expires.",
-                            counter_key,
-                            redis_err,
-                        )
         except Exception as e:
+            memberships = []
             verbose_proxy_logger.warning(
-                "Failed to reset team member spend counters: %s", e
+                "Failed to fetch team memberships for counter invalidation: %s", e
             )
 
-        return await self.prisma_client.db.litellm_teammembership.update_many(
+        update_result = await self.prisma_client.db.litellm_teammembership.update_many(
             where={"budget_id": {"in": budget_ids}},
             data={
                 "spend": 0,
             },
         )
 
+        for m in memberships:
+            await self._invalidate_spend_counter(
+                f"spend:team_member:{m.user_id}:{m.team_id}"
+            )
+
+        return update_result
+
     async def reset_budget_for_keys_linked_to_budgets(
         self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
     ):
@@ -126,17 +141,36 @@ class ResetBudgetJob:
         if not budget_ids:
             return
 
-        return await self.prisma_client.db.litellm_verificationtoken.update_many(
-            where={
-                "budget_id": {"in": budget_ids},
-                "budget_duration": None,  # only keys without their own reset schedule
-                "spend": {"gt": 0},  # only reset keys that have accumulated spend
-            },
-            data={
-                "spend": 0,
-            },
+        where_clause: dict = {
+            "budget_id": {"in": budget_ids},
+            "budget_duration": None,  # only keys without their own reset schedule
+            "spend": {"gt": 0},  # only reset keys that have accumulated spend
+        }
+
+        try:
+            keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
+                where=where_clause
+            )
+        except Exception as e:
+            keys = []
+            verbose_proxy_logger.warning(
+                "Failed to fetch keys for counter invalidation: %s", e
+            )
+
+        update_result = (
+            await self.prisma_client.db.litellm_verificationtoken.update_many(
+                where=where_clause,
+                data={
+                    "spend": 0,
+                },
+            )
         )
 
+        for k in keys:
+            await self._invalidate_spend_counter(f"spend:key:{k.token}")
+
+        return update_result
+
     async def reset_budget_for_litellm_budget_table(self):
         """
         Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
@@ -365,6 +399,10 @@ class ResetBudgetJob:
                         data_list=updated_keys,
                         table_name="key",
                     )
+                    for k in updated_keys:
+                        token = getattr(k, "token", None)
+                        if token:
+                            await self._invalidate_spend_counter(f"spend:key:{token}")
 
             end_time = time.time()
             if len(failed_keys) > 0:  # If any keys failed to reset
@@ -450,6 +488,12 @@ class ResetBudgetJob:
                         data_list=updated_users,
                         table_name="user",
                     )
+                    for u in updated_users:
+                        user_id = getattr(u, "user_id", None)
+                        if user_id:
+                            await self._invalidate_spend_counter(
+                                f"spend:user:{user_id}"
+                            )
 
             end_time = time.time()
             if len(failed_users) > 0:  # If any users failed to reset
@@ -541,6 +585,12 @@ class ResetBudgetJob:
                         data_list=updated_teams,
                         table_name="team",
                     )
+                    for t in updated_teams:
+                        team_id = getattr(t, "team_id", None)
+                        if team_id:
+                            await self._invalidate_spend_counter(
+                                f"spend:team:{team_id}"
+                            )
 
             end_time = time.time()
             if len(failed_teams) > 0:  # If any teams failed to reset

litellm/proxy/db/spend_counter_reseed.py

@@ -129,7 +129,9 @@ class SpendCounterReseed:
         """
         lock = await SpendCounterReseed._get_lock(counter_key)
         async with lock:
-            # Re-check after acquiring the lock - another waiter may have warmed it.
+            # Re-check after acquiring the lock. Skip in-memory on a clean
+            # Redis miss - in-memory is per-pod-stale.
+            redis_clean_miss = False
             if spend_counter_cache.redis_cache is not None:
                 try:
                     val = await spend_counter_cache.redis_cache.async_get_cache(
@@ -137,11 +139,13 @@ class SpendCounterReseed:
                     )
                     if val is not None:
                         return float(val)
+                    redis_clean_miss = True
                 except Exception:
                     pass
-            val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
-            if val is not None:
-                return float(val)
+            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)
 
             db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
             if db_spend is None:
@@ -149,7 +153,7 @@ class SpendCounterReseed:
             # Warm even when 0 so subsequent reads hit cache, not DB.
             try:
                 await spend_counter_cache.async_increment_cache(
-                    key=counter_key, value=db_spend
+                    key=counter_key, value=db_spend, refresh_ttl=True
                 )
             except Exception:
                 verbose_proxy_logger.exception(

litellm/proxy/proxy_server.py

@@ -1798,12 +1798,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
     3. Reseed from authoritative DB spend (counter expired, cross-pod stale)
     4. Caller-supplied fallback (DB unavailable, cold start)
     """
-    # 1. Try Redis first (cross-pod authoritative)
+    # 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
     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)
+            redis_clean_miss = True
         except Exception as e:
             verbose_proxy_logger.debug(
                 "get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
@@ -1811,10 +1815,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
                 e,
             )
 
-    # 2. Fall back to in-memory counter (single-instance or Redis failure)
-    val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
-    if val is not None:
-        return float(val)
+    # 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)
 
     # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
     db_spend = await SpendCounterReseed.coalesced(
@@ -1976,10 +1981,12 @@ 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, refresh_ttl=True
                 )
 
-    await spend_counter_cache.async_increment_cache(key=counter_key, value=increment)
+    await spend_counter_cache.async_increment_cache(
+        key=counter_key, value=increment, refresh_ttl=True
+    )
 
 
 async def update_cache(  # noqa: PLR0915

tests/test_litellm/caching/test_redis_cache.py

@@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping
         )
 
 
+@pytest.mark.asyncio
+async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl(
+    monkeypatch, redis_no_ping
+):
+    """With refresh_ttl=True, every increment should call expire() to bump
+    the TTL, even when the key already has a TTL (counter-style use)."""
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    redis_cache = RedisCache()
+    mock_redis_instance = AsyncMock()
+    mock_redis_instance.__aenter__.return_value = mock_redis_instance
+    mock_redis_instance.__aexit__.return_value = None
+    mock_redis_instance.ttl.return_value = 42  # key already has ~42s left
+
+    with patch.object(
+        redis_cache, "init_async_client", return_value=mock_redis_instance
+    ):
+        await redis_cache.async_increment(
+            key="spend:team_member:u:t", value=0.05, refresh_ttl=True
+        )
+
+    mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60)
+
+
+@pytest.mark.asyncio
+async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl(
+    monkeypatch, redis_no_ping
+):
+    """Default (refresh_ttl=False) preserves window-style semantics: TTL is
+    set only on first creation, never refreshed (used by rate-limit windows)."""
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    redis_cache = RedisCache()
+    mock_redis_instance = AsyncMock()
+    mock_redis_instance.__aenter__.return_value = mock_redis_instance
+    mock_redis_instance.__aexit__.return_value = None
+    mock_redis_instance.ttl.return_value = 42  # key already has ~42s left
+
+    with patch.object(
+        redis_cache, "init_async_client", return_value=mock_redis_instance
+    ):
+        await redis_cache.async_increment(key="rate_limit:window", value=1)
+
+    mock_redis_instance.expire.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping):
     monkeypatch.setenv("REDIS_HOST", "my-fake-host")

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
     asyncio.run(job.reset_budget_windows())  # must not raise
 
     prisma_client.db.litellm_teamtable.update.assert_awaited_once()
+
+
+# ---------------------------------------------------------------------------
+# Counter invalidation on budget reset
+# ---------------------------------------------------------------------------
+
+
+def _make_counter_invalidation_job(monkeypatch):
+    """Stub spend_counter_cache so we can observe invalidation calls."""
+    spend_counter_cache = MagicMock()
+    spend_counter_cache.in_memory_cache.set_cache = MagicMock()
+    spend_counter_cache.redis_cache = MagicMock()
+    spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
+
+    fake_module = types.ModuleType("litellm.proxy.proxy_server")
+    fake_module.spend_counter_cache = spend_counter_cache
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
+
+    return spend_counter_cache
+
+
+def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch):
+    """Team-member budget reset clears the Redis spend counter."""
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    expired_budget = type("B", (), {"budget_id": "budget-1"})
+    membership = type(
+        "Membership",
+        (),
+        {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
+    )
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_teammembership.find_many = AsyncMock(
+        return_value=[membership]
+    )
+    prisma_client.db.litellm_teammembership.update_many = AsyncMock(
+        return_value={"count": 1}
+    )
+
+    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
+    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
+
+    counter_cache.in_memory_cache.set_cache.assert_any_call(
+        key="spend:team_member:alice:team-x", value=0.0, ttl=60
+    )
+    counter_cache.redis_cache.async_set_cache.assert_any_await(
+        key="spend:team_member:alice:team-x", value=0.0, ttl=60
+    )
+
+
+def test_reset_budget_for_keys_invalidates_redis_counter(
+    reset_budget_job, mock_prisma_client, monkeypatch
+):
+    """Key budget reset must clear the Redis spend counter."""
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    now = datetime.now(timezone.utc)
+    mock_prisma_client.data["key"] = [
+        type(
+            "Key",
+            (),
+            {
+                "spend": 100.0,
+                "budget_duration": "30d",
+                "budget_reset_at": now,
+                "id": "key-1",
+                "token": "sk-abc",
+            },
+        )
+    ]
+
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
+
+    counter_cache.in_memory_cache.set_cache.assert_any_call(
+        key="spend:key:sk-abc", value=0.0, ttl=60
+    )
+
+
+def test_reset_budget_for_users_invalidates_redis_counter(
+    reset_budget_job, mock_prisma_client, monkeypatch
+):
+    """User budget reset must clear the Redis spend counter."""
+    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())
+
+    counter_cache.in_memory_cache.set_cache.assert_any_call(
+        key="spend:user:alice", value=0.0, ttl=60
+    )
+
+
+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)
+
+    now = datetime.now(timezone.utc)
+    mock_prisma_client.data["team"] = [
+        type(
+            "Team",
+            (),
+            {
+                "spend": 200.0,
+                "budget_duration": "1mo",
+                "budget_reset_at": now,
+                "id": "team-1",
+                "team_id": "team-x",
+            },
+        )
+    ]
+
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
+
+    counter_cache.in_memory_cache.set_cache.assert_any_call(
+        key="spend:team:team-x", value=0.0, ttl=60
+    )
+
+
+def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
+    """Resetting keys via budget tier must clear each linked key's counter."""
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    expired_budget = type("B", (), {"budget_id": "budget-1"})
+    linked_key = type("Key", (), {"token": "sk-linked"})
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
+        return_value=[linked_key]
+    )
+    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
+        return_value={"count": 1}
+    )
+
+    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
+    asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
+
+    counter_cache.in_memory_cache.set_cache.assert_any_call(
+        key="spend:key:sk-linked", value=0.0, ttl=60
+    )

tests/test_litellm/proxy/test_proxy_server.py

@@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware:
         assert attempts == [
             "called"
         ], f"failing register_fn should be invoked once, not on every request; got {attempts}"
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory():
+    """When Redis is reachable and cleanly returns None (TTL expired,
+    counter genuinely absent), the read must reseed from DB - NOT fall
+    through to per-pod in-memory which only contains this pod's writes.
+
+    Pre-fix in multi-pod deployments, in-memory contained a stale local
+    subset (e.g. $30) while DB had the true cross-pod total ($500). The
+    fall-through returned $30, enforcement passed, bypass.
+    """
+    from litellm.caching.dual_cache import DualCache
+    from litellm.proxy.proxy_server import get_current_spend
+
+    counter_cache = DualCache()
+    counter_key = "spend:team_member:user-1:team-1"
+
+    # Per-pod stale in-memory: only this pod's writes, not cross-pod truth.
+    counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0)
+
+    # Redis cleanly returns None (key expired or never written on this pod).
+    fake_redis = AsyncMock()
+    fake_redis.async_get_cache = AsyncMock(return_value=None)
+    fake_redis.async_increment = AsyncMock(return_value=500.0)
+    counter_cache.redis_cache = fake_redis
+
+    # DB has the authoritative cross-pod spend.
+    db_row = MagicMock()
+    db_row.spend = 500.0
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row)
+
+    import litellm.proxy.proxy_server as ps
+
+    orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
+    ps.spend_counter_cache = counter_cache
+    ps.prisma_client = fake_prisma
+    try:
+        spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
+        assert spend == 500.0, (
+            f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} "
+            f"(stale per-pod in-memory $30 would have caused multi-pod bypass)"
+        )
+    finally:
+        ps.spend_counter_cache = orig_counter
+        ps.prisma_client = orig_prisma
+
+
+@pytest.mark.asyncio
+async def test_get_current_spend_redis_error_falls_back_to_in_memory():
+    """When Redis raises, the read should still degrade to in-memory rather
+    than going straight to DB - in-memory is at least same-pod-fresh and
+    cheaper than a DB query during a Redis outage."""
+    from litellm.caching.dual_cache import DualCache
+    from litellm.proxy.proxy_server import get_current_spend
+
+    counter_cache = DualCache()
+    counter_key = "spend:team_member:user-1:team-1"
+
+    counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0)
+
+    fake_redis = AsyncMock()
+    fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down"))
+    counter_cache.redis_cache = fake_redis
+
+    fake_prisma = MagicMock()
+    fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
+        return_value=MagicMock(spend=999.0)
+    )
+
+    import litellm.proxy.proxy_server as ps
+
+    orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
+    ps.spend_counter_cache = counter_cache
+    ps.prisma_client = fake_prisma
+    try:
+        spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
+        assert spend == 42.0, (
+            f"expected in-memory fallback 42.0 on Redis error, got {spend} "
+            f"(should not have hit DB when Redis errored)"
+        )
+        # DB query should NOT have fired - in-memory short-circuits.
+        fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited()
+    finally:
+        ps.spend_counter_cache = orig_counter
+        ps.prisma_client = orig_prisma