fix(reset_budget): cherry-pick #29358 into patch/v1.87.0rc2 - #29364
Conversation
…eroing counter (#29358) * fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter ResetBudgetJob's batched update_data path shipped the full key/user/team model on each reset. Prisma rejects object_permission_id and budget_limits on the update input type, so any row carrying those fields detonated the entire batch -- spend never reset, budget_reset_at never advanced. After v1.84.0 started populating object_permission_id on UI-created keys, this fires routinely. _reset_budget_common also zeroed the cross-pod spend counter before the DB write, so failed resets left enforcement reading 0 from the counter while the DB still held the over-budget spend, admitting requests past the cap until the counter naturally re-saturated from new reservations. Switch the write to per-row narrow updates ({spend, budget_reset_at}) via db.batch_, and move the counter invalidation out of _reset_budget_common so it only fires after the DB write commits. On DB-write failure the counter is left untouched, enforcement continues to block, and the next scheduler tick can retry without leaving a bypass window. Fixes #27730. * fix(reset_budget): address Greptile review on #29358 - Strengthen the bypass-half regression test: replace the for-loop over call_args_list (vacuously true when empty) with assert_not_called(), so the test would actually flag a re-introduction of counter-zeroing via any code path. - Add the same explanatory docstring on _write_user_reset_updates and _write_team_reset_updates that _write_key_reset_updates already has, so all three helpers point future maintainers at #27730. * test(reset_budget): update test_proxy_budget_reset for new batch-write path Same shape as the previous test_reset_budget_job.py update: keys/users/teams now write through prisma.db.batch_().<table>.update, not update_data, so the tests need a batcher mock and updated assertions. Adds: - _wire_batcher_for_test helper that returns a list which accumulates per-row batch updates captured from prisma_client.db.batch_(). - _attrify helper that wraps dict fixtures so getattr(item, "token") works alongside the dict item-access the fake_reset_* mocks rely on. The new narrow-write helpers use getattr to pull out the row's id, and would silently skip plain dicts otherwise. - Updates 3 partial_failure tests to assert against the batch-call list (rows by id, payload contains only {spend, budget_reset_at}) instead of update_data.assert_awaited_once + data_list inspection. - Updates test_reset_budget_continues_other_categories_on_failure: only budget + enduser still flow through update_data; key/user/team go through the batch path now. - Wires the batcher mock into 3 service_logger_*_success tests so commit() is actually awaitable and the success hook fires. These tests were silently passing locally only because the editable install in .venv pointed at the main repo, not the worktree — running pytest with PYTHONPATH overridden to the worktree (matching CI) reproduces the failures.
Greptile SummaryThis PR cherry-picks #29358 onto
Confidence Score: 4/5The cherry-pick correctly addresses both halves of the reset-budget bug; the write path is narrowed to only the two fields Prisma accepts, and counter invalidation is correctly sequenced after the DB commit. The core logic is sound: narrow-payload batch writes replace the full-model calls that triggered DataError, and counter zeroing now only runs after a successful commit. Two new regression tests cover both failure scenarios. The only observations are a minor empty-batch edge case and an undocumented TTL change on counter invalidation — neither is a functional bug in the changed path. The new
|
| Filename | Overview |
|---|---|
| litellm/proxy/common_utils/reset_budget_job.py | Core fix: three new narrow-write helpers replace full-model update_data calls; counter invalidation moved post-commit; new _invalidate_spend_counter with ttl=60 (new TTL behavior vs original no-TTL). Logic is sound. |
| tests/litellm_utils_tests/test_proxy_budget_reset.py | Existing tests updated to use the new batch-write path via _wire_batcher_for_test; assertions correctly verify narrow payload and assert update_data is not called for keys/users/teams. All mock-based, no real network calls. |
| tests/test_litellm/proxy/common_utils/test_reset_budget_job.py | MockBatcher added to MockDB; existing tests updated to assert via batch_calls instead of updated_data for key/user/team rows; two new regression tests cover the trigger-half and bypass-half of the bug. Tests are thorough and correctly structured. |
Comments Outside Diff (1)
-
litellm/proxy/common_utils/reset_budget_job.py, line 56-84 (link)ttl=60on counter invalidation is a new behaviour vs. the original codeThe replaced code in
_reset_budget_commoncalledset_cache(key=counter_key, value=0.0)with no TTL, keeping the zeroed entry in-memory indefinitely. The new_invalidate_spend_countersetsttl=60on both the in-memory and Redis entries. After 60 seconds the entry expires; the nextget_current_spendcall must rebuild the counter from the DB (which holds 0 post-reset) or from fresh accumulation. This should be safe because the DB row is always the ground truth, but it is a subtle behavioural change that isn't called out in the PR description and could matter if the in-memoryset_cacheimplementation resets the TTL on subsequent incremental writes differently from what callers expect.
Reviews (1): Last reviewed commit: "fix(reset_budget): write only {spend, bu..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async def _write_key_reset_updates( | ||
| self, updated_keys: List[LiteLLM_VerificationToken] | ||
| ) -> None: | ||
| """ | ||
| Write per-row {spend, budget_reset_at} updates for keys. | ||
|
|
||
| Avoids the batched full-model update path, which trips | ||
| prisma.errors.DataError on any row carrying object_permission_id or | ||
| budget_limits (see #27730). Both fields are rejected by Prisma's | ||
| update input type for LiteLLM_VerificationToken, and the failure | ||
| aborts the entire batch — silently leaving spend over the cap and | ||
| budget_reset_at unchanged forever. | ||
| """ | ||
| batcher = self.prisma_client.db.batch_() | ||
| for k in updated_keys: |
There was a problem hiding this comment.
None token/ID
Each _write_*_reset_updates helper calls await batcher.commit() unconditionally after the loop. If every entry in updated_keys has token=None (all rows skipped via continue), the commit is called on an empty batch. Prisma treats this as a no-op, so there's no functional bug, but it makes an unnecessary round-trip. This is an extreme edge case in production (keys without tokens shouldn't exist), but tracking whether any .update() was appended before calling .commit() would make the intent clearer.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
1 participant
Relevant issues
Cherry-pick of #29358 (merge commit
7d1bd9d9f4) onto thepatch/v1.87.0rc2branch (= tagv1.87.0-rc.2). Original PR fixes #27730. Supersedes #29361, which targeted the olderpatch/v1.87.0rc1base.Linear ticket
n/a
Pre-Submission checklist
git diffbetween this PR's tree and7d1bd9d9f4's tree on the three touched files is 0 lines)CI (LiteLLM team)
Bug verification on v1.87.0-rc.2
Confirmed the bug is still present at the v1.87.0-rc.2 tag (head of
patch/v1.87.0rc2) before this cherry-pick lands. rc.2 = rc.1 + 2 patch commits, neither of which touchedreset_budget_job.py, so the rc.1 verification carries over unchanged.Trigger half — full Pydantic model is still shipped through
update_dataon rc.2 at lines 458 / 547 / 644:Each affected call ships the full Pydantic row, triggering Prisma DataError on rows with
object_permission_id/budget_limits.Mechanism half —
_reset_budget_commonstill zeroes the cross-pod counter pre-write on rc.2 at lines 64-71:Counter is zeroed inside
_reset_budget_common, which the caller runs before the (failing)update_databatch. Both halves of the bug are present.Screenshots / Proof of Fix
End-to-end proof-of-fix harness output (pre-fix DataError + bypass; post-fix clean reset + normal enforcement) is captured in the original PR body at #29358 and was run against a real Postgres + Redis proxy. This PR is a verbatim cherry-pick (
git diff HEAD 7d1bd9d9f4 -- <touched files>is 0 lines) so the same harness output applies.Type
Bug Fix
Changes
Verbatim cherry-pick of #29358's merge commit
7d1bd9d9f4c313dc1e2cd4aeec5e70a199ee6529ontopatch/v1.87.0rc2. No conflicts. 3 files, +409/-132 (identical stat to the original merge commit). Full rationale and trigger/mechanism analysis live in the #29358 PR body.