Invalidate spend counters on budget reset

litellm/proxy/common_utils/reset_budget_job.py

@@ -52,6 +52,30 @@ 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."""
+        try:
+            from litellm.proxy.proxy_server import spend_counter_cache
+
+            spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0)
+            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 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,37 +88,17 @@ 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
+                await self._invalidate_spend_counter(
+                    f"spend:team_member:{m.user_id}:{m.team_id}"
                 )
-                # 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:
             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(
@@ -126,12 +130,25 @@ class ResetBudgetJob:
         if not budget_ids:
             return
 
+        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
+            )
+            for k in keys:
+                await self._invalidate_spend_counter(f"spend:key:{k.token}")
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                "Failed to fetch keys for counter invalidation: %s", e
+            )
+
         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
-            },
+            where=where_clause,
             data={
                 "spend": 0,
             },
@@ -360,6 +377,10 @@ class ResetBudgetJob:
                 )
 
                 if updated_keys:
+                    for k in updated_keys:
+                        token = getattr(k, "token", None)
+                        if token:
+                            await self._invalidate_spend_counter(f"spend:key:{token}")
                     await self.prisma_client.update_data(
                         query_type="update_many",
                         data_list=updated_keys,
@@ -445,6 +466,12 @@ class ResetBudgetJob:
                     "Updated users %s", json.dumps(updated_users, indent=4, default=str)
                 )
                 if updated_users:
+                    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}"
+                            )
                     await self.prisma_client.update_data(
                         query_type="update_many",
                         data_list=updated_users,
@@ -536,6 +563,12 @@ class ResetBudgetJob:
                     "Updated teams %s", json.dumps(updated_teams, indent=4, default=str)
                 )
                 if updated_teams:
+                    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}"
+                            )
                     await self.prisma_client.update_data(
                         query_type="update_many",
                         data_list=updated_teams,

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
+    )
+    counter_cache.redis_cache.async_set_cache.assert_any_await(
+        key="spend:team_member:alice:team-x", value=0.0
+    )
+
+
+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
+    )
+
+
+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
+    )
+
+
+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
+    )
+
+
+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
+    )