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.

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)

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:

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:

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)