fix(reset_budget): cherry-pick #29358 into patch/v1.87.0rc1 - #29361
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 cherry-pick of #29358 onto
Confidence Score: 4/5Safe to merge; the fix correctly eliminates both the Prisma DataError and the Redis counter-bypass window, backed by targeted regression tests for each half. The core logic — narrow-field batch writes + post-commit counter invalidation — is sound and all three entity paths (key/user/team) are handled consistently. Two minor style-level observations: the spend literal is 0 (int) rather than 0.0 (float) for consistency with the rest of the codebase, and all three write helpers call batcher.commit() even when every row was skipped due to a None identifier (harmless empty round-trip). Neither affects correctness. No files require special attention; reset_budget_job.py carries the substantive change and the test files provide solid regression coverage for both bug halves.
|
| Filename | Overview |
|---|---|
| litellm/proxy/common_utils/reset_budget_job.py | Core fix: three narrow-write helpers replace full-model update_data calls, and counter invalidation moved post-commit; logic is correct but empty-batch guard is missing. |
| tests/litellm_utils_tests/test_proxy_budget_reset.py | Existing partial-failure tests updated to wire the new batch_ path and assert narrow {spend, budget_reset_at} payload; _attrify helper and _wire_batcher_for_test correctly simulate the Prisma batcher interface. |
| tests/test_litellm/proxy/common_utils/test_reset_budget_job.py | New MockBatcher class mirrors Prisma batch_() ergonomics; two new regression tests verify counter is not zeroed on DB failure and payload never includes extra fields. |
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}
| batcher = self.prisma_client.db.batch_() | ||
| for k in updated_keys: | ||
| token = getattr(k, "token", None) | ||
| if token is None: | ||
| continue | ||
| batcher.litellm_verificationtoken.update( | ||
| where={"token": token}, | ||
| data={"spend": 0, "budget_reset_at": k.budget_reset_at}, | ||
| ) | ||
| await batcher.commit() |
There was a problem hiding this comment.
updated_keys has token=None, no update calls are staged but batcher.commit() is still awaited. Prisma handles an empty batch gracefully, but the cost is an extra round-trip to the database engine. The same applies to _write_user_reset_updates and _write_team_reset_updates. Guarding with a staged flag makes the intent explicit and avoids the needless round-trip.
| batcher = self.prisma_client.db.batch_() | |
| for k in updated_keys: | |
| token = getattr(k, "token", None) | |
| if token is None: | |
| continue | |
| batcher.litellm_verificationtoken.update( | |
| where={"token": token}, | |
| data={"spend": 0, "budget_reset_at": k.budget_reset_at}, | |
| ) | |
| await batcher.commit() | |
| batcher = self.prisma_client.db.batch_() | |
| staged = False | |
| for k in updated_keys: | |
| token = getattr(k, "token", None) | |
| if token is None: | |
| continue | |
| batcher.litellm_verificationtoken.update( | |
| where={"token": token}, | |
| data={"spend": 0, "budget_reset_at": k.budget_reset_at}, | |
| ) | |
| staged = True | |
| if staged: | |
| await batcher.commit() |
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!
| batcher.litellm_verificationtoken.update( | ||
| where={"token": token}, | ||
| data={"spend": 0, "budget_reset_at": k.budget_reset_at}, | ||
| ) |
There was a problem hiding this comment.
spend field is written as the integer literal 0 here, but _reset_budget_common sets item.spend = 0.0 (a float). Prisma coerces both to the Float column type, so there is no functional issue, but keeping the value as 0.0 is more consistent and makes it clear the column is a float field.
| batcher.litellm_verificationtoken.update( | |
| where={"token": token}, | |
| data={"spend": 0, "budget_reset_at": k.budget_reset_at}, | |
| ) | |
| batcher.litellm_verificationtoken.update( | |
| where={"token": token}, | |
| data={"spend": 0.0, "budget_reset_at": k.budget_reset_at}, | |
| ) |
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.0rc1branch. Original PR fixes #27730.Linear ticket
n/a
Pre-Submission checklist
CI (LiteLLM team)
Bug verification on v1.87.0-rc.1
Confirmed the bug is present at the v1.87.0-rc.1 tag (the head of
patch/v1.87.0rc1) before this cherry-pick lands.Trigger half — full Pydantic model is shipped through
update_data:Same shape at lines 545-555 (
updated_users) and 642-652 (updated_teams). Each ships the full Pydantic row; Prisma rejectsobject_permission_idandbudget_limitson the update input type and detonates the batch.Mechanism half —
_reset_budget_commonzeroes the cross-pod counter before the DB write:Counter is zeroed inside
_reset_budget_common, which the caller runs before the (failing)update_databatch. After this cherry-pick lands, the counter writes only fire via_invalidate_spend_counterafter a successful DB commit.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. The cherry-pick is a verbatim replay of that commit (
git diffbetween this PR and the merge commit's tree on the touched files is empty) so the same harness output applies.Type
Bug Fix
Changes
Verbatim cherry-pick of #29358's merge commit
7d1bd9d9f4c313dc1e2cd4aeec5e70a199ee6529ontopatch/v1.87.0rc1. 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.