From a1dd80689686fd39910eeafbcda46837f6eaff03 Mon Sep 17 00:00:00 2001 From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com> Date: Sun, 11 Oct 2026 15:43:52 -0500 Subject: [PATCH] feat(budgets): reset tokens_used on the team budget window tick On the budget_reset_at tick that zeroes spend, hard-reset the team's tokens_used and invalidate the tokens:team:{id} counter so the next read re-seeds from the freshly-zeroed column. The reset is team-scoped (token budgets only apply to teams), so key/user resets are untouched. --- litellm/proxy/common_utils/reset_budget_job.py | 3 +++ litellm/repositories/unit_of_work.py | 8 +++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 61330d12ac3..4c5db3c8089 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1336,6 +1336,9 @@ class ResetBudgetJob: team_id = getattr(t.row, "team_id", None) if team_id: await self._invalidate_spend_counter(f"spend:team:{team_id}") + # Drop the token-budget counter too, so the next read + # re-seeds from the freshly-zeroed tokens_used column. + await self._invalidate_spend_counter(f"tokens:team:{team_id}") end_time = time.time() outcome: Final = _ChunkOutcome( diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 10f092c1cb2..37f2cfbeaef 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -56,9 +56,15 @@ class TeamSpendResetWrites: table: BatchTable def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: + data = dict(_spend_reset_data(budget_reset_at, spend_decrement)) + # Token budgets are team-only, so zero the cumulative tokens_used on the + # same tick that resets spend. A hard reset (vs spend's decrement) is + # acceptable at a budget boundary: a token allowance need not carry the + # handful of tokens that can race the reset the way billing-exact spend does. + data["tokens_used"] = 0 self.table.update( where={"team_id": team_id}, - data=_spend_reset_data(budget_reset_at, spend_decrement), + data=data, )