fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter - #29358
Conversation
…eroing 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.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR fixes a silent
Confidence Score: 5/5Safe to merge — the change is narrowly scoped to the reset-budget write path, makes DB writes strictly less likely to fail, and does not touch the enforcement or auth layers. The narrow-write helpers correctly limit the Prisma payload to two scalar columns that have always been accepted by the update input type, eliminating the DataError root cause. The ordering of DB commit before counter invalidation is sound and directly prevents the bypass window. Both regression tests use deterministic assertions (assert_not_called(), exact payload key-set equality) rather than weaker loop-based checks. No pre-existing behaviour for the enduser or budget-table paths is changed. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/common_utils/reset_budget_job.py | Core fix: three new narrow-write helpers use prisma.db.batch_() to write only {spend, budget_reset_at} per row; pre-emptive counter zeroing removed from _reset_budget_common so invalidation only runs after a successful DB commit |
| tests/test_litellm/proxy/common_utils/test_reset_budget_job.py | MockDB gains a batch_() implementation; three existing test assertions updated to verify the new narrow-write contract; two regression tests added (trigger-half: payload must be exactly {spend, budget_reset_at}; bypass-half: counter must not be zeroed when DB write raises), using assert_not_called() correctly |
| tests/litellm_utils_tests/test_proxy_budget_reset.py | Existing integration-style tests updated to wire _wire_batcher_for_test(), add token/user_id/team_id fields to test fixtures, and assert against the new batch-write path rather than the removed update_data path |
Reviews (2): Last reviewed commit: "test(reset_budget): update test_proxy_bu..." | Re-trigger Greptile
- 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.
…e 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.
7d1bd9d
into
litellm_internal_staging
…eroing counter (#29358) (#29361) * 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.
…eroing counter (#29358) (#29364) * 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.
* 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.
…eroing counter (BerriAI#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 BerriAI#27730. * fix(reset_budget): address Greptile review on BerriAI#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 BerriAI#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.
Relevant issues
Fixes #27730. Customers on v1.84.0+ saw
ResetBudgetJobsilently fail for every UI-created key, with two compounding effects: spend never reset at the budget cycle boundary, and a periodic budget-enforcement bypass window opened on every scheduler tick.Linear ticket
n/a
Pre-Submission checklist
uv run pytest tests/test_litellm/proxy/common_utils/test_reset_budget_job.py -v-> 47 passed)CI (LiteLLM team)
Screenshots / Proof of Fix
Deterministic harness in
~/.fix-verify-scratch/budget-reset-184/deterministic-repro.sh. It creates a key withobject_permission+budget_limits+budget_duration=60s, forces DB spend = 5.0 overmax_budget = 1.0, waits for the scheduler tick, then probes the budget gate.Pre-fix run:
Post-fix run, same harness, same proxy, same key shape:
Adjacent test on a plain key (no
object_permission, nobudget_limits) still passes -- the path that worked pre-fix continues to work post-fix.Type
Bug Fix
Changes
ResetBudgetJobhad two compounding problems against the customer-typical row shape (UI-created keys that auto-populateobject_permission_id, plus any team admin who turned on multi-windowbudget_limits).Trigger.
reset_budget_for_litellm_keys/_users/_teamsshipped the full Pydantic model throughupdate_data(query_type="update_many", ...). The batchedprisma.update(data=...)call rejectsobject_permission_id(relation FK not in the update input type) andbudget_limits(Json column needsNullableJsonNullValueInputwrapping). Prisma raisesDataError, the batched commit fails atomically, no DB row gets updated. The exception is caught and logged at ERROR but swallowed -- no metric, no alert -- sospendstays over the cap andbudget_reset_atstays in the past forever.Mechanism.
_reset_budget_commonzeroed the cross-pod spend counter (spend_counter_cache.in_memory_cacheandredis_cache) before the batched DB write was even attempted. When the write failed, the counter was already at 0 with no rollback.get_current_spendreads Redis first; with a non-None 0.0 in the counter, it never falls through to the cachedvalid_token.spendfallback._virtual_key_max_budget_checksaw0 < max_budgetand admitted requests until the counter naturally re-saturated from new reservations. Each reset tick (~10 min default) re-opened that window.The fix:
_write_{key,user,team}_reset_updates) onResetBudgetJobthat opendb.batch_()and call<table>.update(where=..., data={"spend": 0, "budget_reset_at": ...})per row. Two scalar columns -- no Json, no relation FKs -- so Prisma's update input type accepts them on every row regardless of what other fields the row carries._reset_budget_commonno longer touches the spend counter at all. The post-DB-commit_invalidate_spend_countercall (already added by4e26835098and already ordered afterupdate_datain the caller) becomes the sole counter writer. If the DB write raises, the post-write invalidation never runs and the counter retains its pre-reset value -- enforcement continues to block, and the next scheduler tick can retry without a bypass window in between.Test coverage:
tests/test_litellm/proxy/common_utils/test_reset_budget_job.pygets aMockDB.batch_()so the new write path can be exercised under mocks.update_datashape were rewritten to assert the new per-row batch write contract (only{spend, budget_reset_at}, scoped bytoken/user_id/team_id).test_reset_budget_for_keys_writes_only_spend_and_reset_at-- the reset payload must contain exactly those two fields, no extras. Pins the trigger close.test_reset_does_not_zero_counter_when_db_write_fails-- when the batched commit raises, no_invalidate_spend_countercall fires for any affected key. Pins the bypass-mechanism close.47/47 reset-budget-job pytest pass post-fix.
Notes for review
litellm_oss_stagingwith a different shape (extendsupdate_datato accept dict payloads). This PR targetslitellm_internal_stagingand goes further by also covering the counter-bypass mechanism, plus the user/team analogues that fix(proxy): use >= in _team_max_budget_check to match other budget checks #28051 doesn't touch.~/.fix-verify-scratch/budget-reset-184/is the integration coverage and can be replayed against any proxy with real Redis + Postgres.