From cc4e31e3180bf653758fdd74ff8a6eaa0397647d Mon Sep 17 00:00:00 2001 From: pylrn Date: Thu, 3 Sep 2026 02:12:50 +0530 Subject: [PATCH] fix(proxy): self-heal null budget_duration rows without zeroing spend (#39370) Closes #39370 A team, user, or key row with budget_duration=null but a non-null, past budget_reset_at (e.g. from legacy data, migration, or removal of budget duration) would get selected by the reset-budget scheduler on every run. Because compute_budget_reset_at only advances budget_reset_at when budget_duration is not None, _reset_budget_common would zero spend every tick while leaving budget_reset_at permanently stuck in the past. Changes: - In `_reset_budget_common`: when budget_duration is None, do not wipe spend and self-heal the row by clearing budget_reset_at to None. - Added `queue_clear_budget_reset_at` to unit of work writers (KeySpendResetWrites, UserSpendResetWrites, TeamSpendResetWrites) to safely update budget_reset_at=None in the DB batch without modifying spend. - In `_write_*_reset_updates_once`: dispatch to `queue_clear_budget_reset_at` when budget_duration is None, and skip spend counter invalidation. - In `_count_advanced`: count rows cleared to None as advanced. - In `VerificationTokenRepository` find_many query: support NULL budget_reset_at with non-NULL budget_duration so newly created keys reset consistently like users and teams. - Added regression unit tests for teams, users, and keys in test_reset_budget_job.py. --- .../proxy/common_utils/reset_budget_job.py | 59 ++++++++----- litellm/proxy/utils.py | 19 +++-- litellm/repositories/unit_of_work.py | 18 ++++ .../common_utils/test_reset_budget_job.py | 82 +++++++++++++++++++ 4 files changed, 152 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..fb16699bd40 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -288,7 +288,11 @@ def _count_advanced(reset_ats: Iterable[object], cutoff: datetime) -> int: would re-read the same chunk until the per-run cap on every tick. """ utc_cutoff: Final = _as_utc(cutoff) - return sum(1 for reset_at in reset_ats if isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff) + return sum( + 1 + for reset_at in reset_ats + if reset_at is None or (isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff) + ) def _phase_is_drained(outcome: _ChunkOutcome) -> bool: @@ -844,11 +848,14 @@ class ResetBudgetJob: for k in updated_keys: if k.token is None: continue - uow.keys.queue_spend_reset( - token=k.token, - budget_reset_at=k.budget_reset_at, - spend_decrement=k.max_budget if (k.spend or 0.0) > 0.0 else None, - ) + if getattr(k, "budget_duration", None) is None: + uow.keys.queue_clear_budget_reset_at(token=k.token) + else: + uow.keys.queue_spend_reset( + token=k.token, + budget_reset_at=k.budget_reset_at, + spend_decrement=k.max_budget if (k.spend or 0.0) > 0.0 else None, + ) async def _write_user_reset_updates(self, updated_users: list[LiteLLM_UserTable]) -> None: """ @@ -866,11 +873,14 @@ class ResetBudgetJob: async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for u in updated_users: - uow.users.queue_spend_reset( - user_id=u.user_id, - budget_reset_at=u.budget_reset_at, - spend_decrement=u.max_budget if (u.spend or 0.0) > 0.0 else None, - ) + if getattr(u, "budget_duration", None) is None: + uow.users.queue_clear_budget_reset_at(user_id=u.user_id) + else: + uow.users.queue_spend_reset( + user_id=u.user_id, + budget_reset_at=u.budget_reset_at, + spend_decrement=u.max_budget if (u.spend or 0.0) > 0.0 else None, + ) async def _write_team_reset_updates(self, updated_teams: list[LiteLLM_TeamTable]) -> None: """ @@ -888,11 +898,14 @@ class ResetBudgetJob: async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for t in updated_teams: - uow.teams.queue_spend_reset( - team_id=t.team_id, - budget_reset_at=t.budget_reset_at, - spend_decrement=t.max_budget if (t.spend or 0.0) > 0.0 else None, - ) + if getattr(t, "budget_duration", None) is None: + uow.teams.queue_clear_budget_reset_at(team_id=t.team_id) + else: + uow.teams.queue_spend_reset( + team_id=t.team_id, + budget_reset_at=t.budget_reset_at, + spend_decrement=t.max_budget if (t.spend or 0.0) > 0.0 else None, + ) def _emit_phase_failure( self, @@ -967,7 +980,7 @@ class ResetBudgetJob: await self._write_key_reset_updates(updated_keys=updated_keys) for k in updated_keys: token = getattr(k, "token", None) - if token: + if token and getattr(k, "budget_duration", None) is not None: await self._invalidate_spend_counter(f"spend:key:{token}", new_spend=k.spend or 0.0) end_time = time.time() @@ -1072,9 +1085,9 @@ class ResetBudgetJob: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: user_id = getattr(u, "user_id", None) - if user_id: + if user_id and getattr(u, "budget_duration", None) is not None: await self._invalidate_spend_counter(f"spend:user:{user_id}", new_spend=u.spend or 0.0) - if user_id == LITELLM_PROXY_BUDGET_NAME: + if user_id == LITELLM_PROXY_BUDGET_NAME and getattr(u, "budget_duration", None) is not None: await self._invalidate_global_proxy_spend_cache() end_time = time.time() @@ -1181,7 +1194,7 @@ class ResetBudgetJob: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: team_id = getattr(t, "team_id", None) - if team_id: + if team_id and getattr(t, "budget_duration", None) is not None: await self._invalidate_spend_counter(f"spend:team:{team_id}", new_spend=t.spend or 0.0) end_time = time.time() @@ -1443,11 +1456,15 @@ class ResetBudgetJob: still holds the pre-reset value, admitting requests past the cap. """ try: - item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: + item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0 item.budget_reset_at = compute_budget_reset_at( budget_duration=item.budget_duration, settings=reset_settings ) + else: + # Self-heal: row has no budget_duration configured (or it was removed), + # so clear any stale budget_reset_at without wiping spend. + item.budget_reset_at = None return item except Exception as e: verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 051d36c4d0f..cc6f2c42eac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4045,12 +4045,21 @@ class PrismaClient: response = await VerificationTokenRepository(self).table.find_many( take=limit, where={ - "OR": [ - {"expires": None}, - {"expires": {"gt": expires}}, - ], - "budget_reset_at": {"lt": reset_at}, "NOT": {"budget_duration": None}, + "AND": [ + { + "OR": [ + {"expires": None}, + {"expires": {"gt": expires}}, + ] + }, + { + "OR": [ + {"budget_reset_at": None}, + {"budget_reset_at": {"lt": reset_at}}, + ] + }, + ], }, ) if response is not None and len(response) > 0: diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index a497d0580db..645105f0ef9 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -45,6 +45,12 @@ class KeySpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clear_budget_reset_at(self, token: str) -> None: + self.table.update( + where={"token": token}, # mutable-ok: prisma where filter must be a dict + data={"budget_reset_at": None}, + ) + @dataclass(frozen=True, slots=True) class UserSpendResetWrites: @@ -58,6 +64,12 @@ class UserSpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clear_budget_reset_at(self, user_id: str) -> None: + self.table.update( + where={"user_id": user_id}, # mutable-ok: prisma where filter must be a dict + data={"budget_reset_at": None}, + ) + @dataclass(frozen=True, slots=True) class TeamSpendResetWrites: @@ -71,6 +83,12 @@ class TeamSpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clear_budget_reset_at(self, team_id: str) -> None: + self.table.update( + where={"team_id": team_id}, # mutable-ok: prisma where filter must be a dict + data={"budget_reset_at": None}, + ) + @dataclass(frozen=True, slots=True) class LinkedSpendResetWrites: 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 03b05bd9d87..006c311db61 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 @@ -379,6 +379,88 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): assert set(write["data"].keys()) == {"spend", "budget_reset_at"} +def test_reset_budget_for_team_null_duration_self_heals(reset_budget_job, mock_prisma_client): + """If a team row has budget_duration=None but a stale budget_reset_at in the past, + the reset job must self-heal by setting budget_reset_at=None WITHOUT wiping spend.""" + now = datetime.now(timezone.utc) + test_team = type( + "LiteLLM_TeamTable", + (), + { + "spend": 500.0, + "budget_duration": None, + "budget_reset_at": now - timedelta(days=1), + "id": "test-team-null-dur", + "team_id": "tid-null", + }, + ) + + mock_prisma_client.data["team"] = [test_team] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) + + team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"] + assert len(team_writes) == 1 + write = team_writes[0] + assert write["where"] == {"team_id": "tid-null"} + assert "spend" not in write["data"] + assert write["data"]["budget_reset_at"] is None + + +def test_reset_budget_for_user_null_duration_self_heals(reset_budget_job, mock_prisma_client): + """If a user row has budget_duration=None but a stale budget_reset_at in the past, + the reset job must self-heal by setting budget_reset_at=None WITHOUT wiping spend.""" + now = datetime.now(timezone.utc) + test_user = type( + "LiteLLM_UserTable", + (), + { + "spend": 300.0, + "budget_duration": None, + "budget_reset_at": now - timedelta(days=1), + "user_id": "uid-null", + }, + ) + + mock_prisma_client.data["user"] = [test_user] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"] + assert len(user_writes) == 1 + write = user_writes[0] + assert write["where"] == {"user_id": "uid-null"} + assert "spend" not in write["data"] + assert write["data"]["budget_reset_at"] is None + + +def test_reset_budget_for_key_null_duration_self_heals(reset_budget_job, mock_prisma_client): + """If a key row has budget_duration=None but a stale budget_reset_at in the past, + the reset job must self-heal by setting budget_reset_at=None WITHOUT wiping spend.""" + now = datetime.now(timezone.utc) + test_key = type( + "LiteLLM_VerificationToken", + (), + { + "spend": 100.0, + "budget_duration": None, + "budget_reset_at": now - timedelta(days=1), + "token": "tok-null", + }, + ) + + mock_prisma_client.data["key"] = [test_key] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] + assert len(key_writes) == 1 + write = key_writes[0] + assert write["where"] == {"token": "tok-null"} + assert "spend" not in write["data"] + assert write["data"]["budget_reset_at"] is None + + def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): """End-user spend is zeroed and the tier's window advances, in one batch.""" now = datetime.now(timezone.utc)