Cherry-pick #29358 onto patch/v1.84.3

[yuneng-berri]

Cherry-picks three commits from #29358 into patch/v1.84.3.

Commits

• fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter (9bfac01)

  • ResetBudgetJob's batched update_data path shipped the full key/user/team Pydantic model on each reset. Prisma rejects object_permission_id (relation FK) and budget_limits (Json needs NullableJsonNullValueInput) on the update input type, so any row carrying those fields detonated the entire batch with DataError. Switches the write to per-row narrow updates via db.batch_().<table>.update(where=..., data={"spend": 0, "budget_reset_at": ...}) for keys, users, and teams.
  • _reset_budget_common also zeroed the cross-pod spend counter (spend_counter_cache) before the DB write was even attempted. On DB-write failure the counter was already at 0 with no rollback, so get_current_spend read 0 from Redis while the DB still held the over-budget value -- _virtual_key_max_budget_check saw 0 < max_budget and admitted requests until the counter naturally re-saturated. Removes the inline pre-write counter zero from _reset_budget_common; the post-DB-commit _invalidate_spend_counter call (already on this branch and already ordered after the DB write) becomes the sole counter writer. On DB-write failure the counter retains its pre-reset value, enforcement keeps blocking, and the next scheduler tick can retry without a bypass window in between.
  • Resolved a small conflict in tests/test_litellm/proxy/common_utils/test_reset_budget_job.py: kept the new batch_() and _record_and_commit machinery on MockDB; dropped two MockLiteLLMOrganizationTable / MockLiteLLMTagTable references that don't exist at v1.84.3 (they were introduced by separate staging commits).

• fix(reset_budget): address Greptile review on fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter #29358 (73b6e3f)

  • Strengthens the bypass-half regression test: replaces a 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.
  • Adds the same #27730 explanatory docstring on _write_user_reset_updates and _write_team_reset_updates that _write_key_reset_updates already had.

• test(reset_budget): update test_proxy_budget_reset for new batch-write path (2dcaf2f)

  • _wire_batcher_for_test helper that returns a list accumulating 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 existing 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 the partial-failure / continues-on-failure / service-logger-success tests for keys/users/teams to assert against the captured batch-call list instead of update_data.assert_awaited_*. End-user and budget-table paths still go through update_data and their assertions are unchanged.

Test plan

• uv run pytest tests/test_litellm/proxy/common_utils/test_reset_budget_job.py tests/litellm_utils_tests/test_proxy_budget_reset.py against the cherry-picked tree -> 61 tests pass (47 in the first file, 14 in the second).

[greptile-apps]

Greptile Summary

This cherry-pick fixes a two-part budget-reset bug (#27730): the full-model update_data batch path was rejected by Prisma's DataError on rows carrying object_permission_id or budget_limits, leaving budgets permanently un-reset; and the Redis spend counter was zeroed before the DB write, opening a bypass window on write failure.

• Three new _write_{key,user,team}_reset_updates helpers replace the full-model batch path with narrow per-row prisma.db.batch_().<table>.update writes that only send {spend, budget_reset_at}, resolving the DataError trigger.
• _reset_budget_common no longer pre-zeros the spend counter; counter invalidation (_invalidate_spend_counter) is now called only after batcher.commit() succeeds, closing the bypass window.
• Tests in both test files are updated to assert against the new batch-write path; two new regression tests cover the trigger-half (payload narrowness) and bypass-half (counter not zeroed on DB failure).

Confidence Score: 4/5

Safe to merge; the core production fix is correct and the partial-failure tests verify it end-to-end.

The production code change is well-reasoned and the key regression tests correctly guard against re-introduction of both failure modes. The only concern is in the service-logger success tests where fake_reset_{key,user,team} returns a plain dict, causing the batch write path to be silently no-op'd despite _wire_batcher_for_test being wired and ID fields being added to the fixtures.

tests/litellm_utils_tests/test_proxy_budget_reset.py — the service-logger success tests silently skip the batch write path due to plain-dict/getattr incompatibility.

Important Files Changed

Filename	Overview
litellm/proxy/common_utils/reset_budget_job.py	Introduces three narrow-write helpers that replace the full-model batched update_data path with per-row prisma batch_() writes for only {spend, budget_reset_at}. Removes the pre-write spend-counter zero from _reset_budget_common, making counter invalidation strictly post-commit. Logic is correct and the bypass-window fix is sound.
tests/litellm_utils_tests/test_proxy_budget_reset.py	Adds _attrify and wire_batcher_for_test helpers and updates partial-failure / cross-category tests to assert against the new batch() write path. Service-logger success tests add ID fields and wire the batcher, but returned plain dicts cause getattr to return None, so the batch write path is silently skipped in those tests.
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py	Adds MockBatcher and wires MockDB.batch_() to accumulate per-row calls; updates key/user/team reset tests to assert narrow {spend, budget_reset_at} payloads; adds regression tests for the bypass-half and trigger-half. Well-structured and comprehensive.

Comments Outside Diff (1)

1. tests/litellm_utils_tests/test_proxy_budget_reset.py, line 605-654 (link)

   "token" / "user_id" / "team_id" fields in service-logger tests are dead code

   In test_service_logger_keys_success (and the analogous user/team tests), fake_reset_key returns the original plain dict. The narrow-write helpers use getattr(k, "token", None), which returns None for a plain dict — dict keys are not exposed as attributes — so every row is silently skipped and batcher.commit() is called with an empty payload. The "token": "key1" fields added to the fixtures and the _wire_batcher_for_test() call prevent a crash, but they do not actually exercise the batch write path. If _write_key_reset_updates regressed, these tests would still pass while the partial-failure tests would catch it — but the asymmetry is easy to miss.

   The fix is to have fake_reset_key return _attrify(key) instead of key, matching the pattern used in the partial-failure tests.

_{Reviews (1): Last reviewed commit: "test(reset_budget): update test_proxy_bu..." | Re-trigger Greptile}