Skip to content

fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter - #29358

Merged
yuneng-berri merged 3 commits into
litellm_internal_stagingfrom
litellm_/angry-goldberg-d30aa0
May 31, 2026
Merged

fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter#29358
yuneng-berri merged 3 commits into
litellm_internal_stagingfrom
litellm_/angry-goldberg-d30aa0

Conversation

@yuneng-berri

Copy link
Copy Markdown
Collaborator

Relevant issues

Fixes #27730. Customers on v1.84.0+ saw ResetBudgetJob silently 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

  • I have added meaningful tests
  • My PR passes the targeted unit tests on the touched files (uv run pytest tests/test_litellm/proxy/common_utils/test_reset_budget_job.py -v -> 47 passed)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review and received a Confidence Score of at least 4/5

CI (LiteLLM team)

  • Branch creation CI run -- link:
  • CI run for the last commit -- link:
  • Merge / cherry-pick CI run -- links:

Screenshots / Proof of Fix

Deterministic harness in ~/.fix-verify-scratch/budget-reset-184/deterministic-repro.sh. It creates a key with object_permission + budget_limits + budget_duration=60s, forces DB spend = 5.0 over max_budget = 1.0, waits for the scheduler tick, then probes the budget gate.

Pre-fix run:

===== Step 6: confirm NEW reset DataError in log =====
NEW DataError/Failed-to-reset lines since run start: 6
prisma.errors.DataError: Unable to match input value to any allowed input type for the field.
  Parse errors: [`updateOneLiteLLM_VerificationToken.data.object_permission_id`: Field does not exist in enclosing type.,
                 Invalid argument type. budget_limits should be of any of the following types:
                   NullableJsonNullValueInput, Json]
===== Step 7: post-reset state =====
DB spend = 5
Redis counter = 0.0
===== Step 8: post-reset budget-test request =====
Post-reset HTTP=200 (mock_response completion -- bypass)
===== Step 9: burst of 20 =====
Burst: 4-15 past gate / 20 (varies by burst race)
>>> CONFIRMED: bypass observed

Post-fix run, same harness, same proxy, same key shape:

===== Step 6: confirm NEW reset DataError in log =====
NEW DataError/Failed-to-reset lines since run start: 0
===== Step 7: post-reset state =====
DB spend = 0          <- reset committed
Redis counter = 0.0
===== Step 8: post-reset budget-test request =====
Post-reset HTTP=200 (legitimate: spend is genuinely 0 after the successful reset)
===== Step 9: burst of 20 =====
Burst: 19 past gate / 1 blocked-by-budget   <- normal enforcement, cap is fresh
>>> FIX VERIFIED: reset committed (DB spend zeroed, no DataError, gate behavior normal)

Adjacent test on a plain key (no object_permission, no budget_limits) still passes -- the path that worked pre-fix continues to work post-fix.

Type

Bug Fix

Changes

ResetBudgetJob had two compounding problems against the customer-typical row shape (UI-created keys that auto-populate object_permission_id, plus any team admin who turned on multi-window budget_limits).

Trigger. reset_budget_for_litellm_keys / _users / _teams shipped the full Pydantic model through update_data(query_type="update_many", ...). The batched prisma.update(data=...) call rejects object_permission_id (relation FK not in the update input type) and budget_limits (Json column needs NullableJsonNullValueInput wrapping). Prisma raises DataError, the batched commit fails atomically, no DB row gets updated. The exception is caught and logged at ERROR but swallowed -- no metric, no alert -- so spend stays over the cap and budget_reset_at stays in the past forever.

Mechanism. _reset_budget_common zeroed the cross-pod spend counter (spend_counter_cache.in_memory_cache and redis_cache) before the batched DB write was even attempted. When the write failed, the counter was already at 0 with no rollback. get_current_spend reads Redis first; with a non-None 0.0 in the counter, it never falls through to the cached valid_token.spend fallback. _virtual_key_max_budget_check saw 0 < max_budget and admitted requests until the counter naturally re-saturated from new reservations. Each reset tick (~10 min default) re-opened that window.

The fix:

  1. Three new narrow-write helpers (_write_{key,user,team}_reset_updates) on ResetBudgetJob that open db.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.
  2. _reset_budget_common no longer touches the spend counter at all. The post-DB-commit _invalidate_spend_counter call (already added by 4e26835098 and already ordered after update_data in 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.py gets a MockDB.batch_() so the new write path can be exercised under mocks.
  • Three existing tests that asserted against the old update_data shape were rewritten to assert the new per-row batch write contract (only {spend, budget_reset_at}, scoped by token / user_id / team_id).
  • Two new regression tests pin the bug:
    • 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_counter call fires for any affected key. Pins the bypass-mechanism close.

47/47 reset-budget-job pytest pass post-fix.

Notes for review

…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

codecov Bot commented May 30, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented May 30, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes a silent ResetBudgetJob failure on UI-created keys (v1.84.0+) that had two compounding effects: spend was never reset at the budget-cycle boundary, and a bypass window opened on every scheduler tick. Both root causes are addressed here.

  • Trigger fix: Replaces the full-Pydantic-model update_data(query_type="update_many") path with three narrow helpers (_write_{key,user,team}_reset_updates) that open a prisma.db.batch_() and write only {spend, budget_reset_at} per row, sidestepping Prisma's rejection of object_permission_id and budget_limits on the update input type.
  • Bypass fix: Removes the pre-emptive spend-counter zeroing from _reset_budget_common so that _invalidate_spend_counter only runs after the DB write commits successfully; if the write raises, enforcement continues uninterrupted.
  • Tests: MockDB gains a batch_() implementation; three existing tests are updated to assert the new contract; two new regression tests pin the trigger-half (payload must be exactly {spend, budget_reset_at}) and bypass-half (counter must not be zeroed when the DB write fails), with the bypass test correctly using assert_not_called() rather than iterating call_args_list.

Confidence Score: 5/5

Safe 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.

Important Files Changed

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

Comment thread tests/test_litellm/proxy/common_utils/test_reset_budget_job.py Outdated
Comment thread litellm/proxy/common_utils/reset_budget_job.py
- 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.
@yuneng-berri

Copy link
Copy Markdown
Collaborator Author

@greptile

@ryan-crabbe-berri ryan-crabbe-berri left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@yuneng-berri
yuneng-berri merged commit 7d1bd9d into litellm_internal_staging May 31, 2026
145 of 146 checks passed
yuneng-berri added a commit that referenced this pull request May 31, 2026
…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.
yuneng-berri added a commit that referenced this pull request May 31, 2026
…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.
yuneng-berri added a commit that referenced this pull request May 31, 2026
* 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.
mateo-berri added a commit that referenced this pull request Jun 2, 2026
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] ResetBudgetJob fails on key/team batch updates when rows include relation and JSON fields

2 participants