diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 0bd2f18bae..0928ce914d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -54,15 +54,22 @@ class ResetBudgetJob: @staticmethod async def _invalidate_spend_counter(counter_key: str) -> None: - """Zero a spend counter so a DB-row reset takes effect immediately.""" + """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) + 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 + key=counter_key, value=0.0, ttl=60 ) except Exception as redis_err: verbose_proxy_logger.warning( @@ -92,22 +99,26 @@ class ResetBudgetJob: memberships = await self.prisma_client.db.litellm_teammembership.find_many( where={"budget_id": {"in": budget_ids}} ) - for m in memberships: - await self._invalidate_spend_counter( - f"spend:team_member:{m.user_id}:{m.team_id}" - ) except Exception as e: + memberships = [] verbose_proxy_logger.warning( "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] ): @@ -140,20 +151,26 @@ class ResetBudgetJob: 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: + keys = [] verbose_proxy_logger.warning( "Failed to fetch keys for counter invalidation: %s", e ) - return await self.prisma_client.db.litellm_verificationtoken.update_many( - where=where_clause, - data={ - "spend": 0, - }, + 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 @@ -377,15 +394,15 @@ 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, 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 @@ -466,17 +483,17 @@ class ResetBudgetJob: "Updated users %s", json.dumps(updated_users, indent=4, default=str) ) if updated_users: + await self.prisma_client.update_data( + query_type="update_many", + 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}" ) - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_users, - table_name="user", - ) end_time = time.time() if len(failed_users) > 0: # If any users failed to reset @@ -563,17 +580,17 @@ class ResetBudgetJob: "Updated teams %s", json.dumps(updated_teams, indent=4, default=str) ) if updated_teams: + await self.prisma_client.update_data( + query_type="update_many", + 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}" ) - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_teams, - table_name="team", - ) end_time = time.time() if len(failed_teams) > 0: # If any teams failed to reset 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 ba9e0e8022..5c86f9057a 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 @@ -1093,10 +1093,10 @@ def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): 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 + 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 + key="spend:team_member:alice:team-x", value=0.0, ttl=60 ) @@ -1124,7 +1124,7 @@ def test_reset_budget_for_keys_invalidates_redis_counter( 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 + key="spend:key:sk-abc", value=0.0, ttl=60 ) @@ -1152,7 +1152,7 @@ def test_reset_budget_for_users_invalidates_redis_counter( 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 + key="spend:user:alice", value=0.0, ttl=60 ) @@ -1180,7 +1180,7 @@ def test_reset_budget_for_teams_invalidates_redis_counter( 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 + key="spend:team:team-x", value=0.0, ttl=60 ) @@ -1203,5 +1203,5 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monke 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 + key="spend:key:sk-linked", value=0.0, ttl=60 )