From fed5f36a3dfbca6b189ab741ade45399dd6aa078 Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Wed, 29 Apr 2026 15:41:05 -0700 Subject: [PATCH] Invalidate spend counters on budget reset --- .../proxy/common_utils/reset_budget_job.py | 89 ++++++---- .../common_utils/test_reset_budget_job.py | 156 ++++++++++++++++++ 2 files changed, 217 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index e486336cec0..0bd2f18bae9 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/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, diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 379ccf4d9af..ba9e0e80228 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/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 + )