fix(proxy): reset carried spend when a budget window is (re)armed on key/user/team update

[yryzhan-vitech]

TLDR

Problem this solves:

• Applying budget_duration on an existing key/user/team leaves carried spend → instant 429 on a fresh window
• The cross-pod spend counter also isn't invalidated on that path, so it stays warm at the stale value

How it solves it:

• Zero spend + invalidate the counter, but ONLY when the window is genuinely (re)armed — never on a no-op re-send

Relevant issues

Fixes #34492

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Reproduction: create a key/user/team without a budget, accumulate spend (e.g. 100), then /key/update (or /user/update, /team/update) with max_budget=50 + budget_duration="30d". Pre-fix the next request 429s because the carried spend exceeds the new cap on a brand-new window.

$ pytest tests/test_litellm/proxy/common_utils/test_timezone_utils.py -k "newly_armed or rearms or resend" \
         tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py -k "budget_window or window_resend" \
         tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py -k "budget_window or unchanged_window" \
         tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py -k "zeroes_spend_on_new or preserves_spend_on_unchanged"
11 passed

Live proof on a running proxy to be added.

Type

🐛 Bug Fix

Changes

Applying a budget window on an existing entity via update set budget_reset_at but never reset the carried lifetime spend, so an entity with prior spend was instantly over-cap and 429'd on a brand-new window (a key/user/team created with a budget starts at spend=0; the update path now matches).

• New shared predicate is_budget_window_newly_armed(new_duration, existing_duration, existing_reset_at) in litellm/proxy/common_utils/timezone_utils.py: True on first application, an expired window (existing_reset_at < now), or a changed duration; False on a no-op re-send of the same active duration.
• key_management_endpoints.py, internal_user_endpoints.py, team_endpoints.py update paths: zero spend only when the predicate says the window is newly armed.
• Post-commit, invalidate/re-seed the cross-pod spend counter (spend:key:* / spend:user:* / spend:team:*) so budget checks don't read a stale value on the fresh window. The user path reuses the existing _invalidate_user_spend_counter_if_changed; key and team invalidate explicitly (the existing key sync only fires when spend is passed explicitly, and teams have no spend-triggered invalidation).

Notes for review:

• The key correctness requirement is that a no-op re-send of the same active budget_duration (e.g. a UI form resubmitting the whole object alongside an unrelated edit) must NOT wipe accumulated mid-window spend. The predicate gates exactly that; regression tests cover the preserve case for all three paths.
• The self-escalation guard (GHSA-wvg4) ordering is preserved: injecting spend=0.0 before the guard means a non-admin self-update that would arm a window is correctly blocked (zeroing one's own spend is itself an escalation), failing closed.
• Counter invalidation is done AFTER the DB commit, matching the existing anti-bypass ordering.

Extends the counter-invalidation mechanism from #30785 and the {spend, budget_reset_at} write pattern from #29358 to the budget_duration-change trigger, which neither covered.

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 2 · PR risk: 0/10

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

Updates key, user, and team budget-window handling to reset carried spend only when a window is newly armed, synchronize spend counters after persistence, tighten re-arm authorization, and include server-inferred budget fields in key audit records

Confidence Score: 5/5

The PR appears safe to merge

No blocking failures remain

Important Files Changed

Filename	Overview
litellm/proxy/common_utils/timezone_utils.py	Adds the shared predicate distinguishing newly armed or expired budget windows from unchanged active-window resubmissions
litellm/proxy/management_endpoints/internal_user_endpoints.py	Applies conditional spend resets to single and bulk user updates and invalidates the corresponding spend counters
litellm/proxy/management_endpoints/key_management_endpoints.py	Applies conditional key spend resets, synchronizes key spend counters, enforces budget-update authorization, and forwards inferred values for auditing
litellm/proxy/management_endpoints/team_endpoints.py	Resets team spend on authorized window re-arms and invalidates the plain team spend counter after persistence
litellm/proxy/hooks/key_management_event_hooks.py	Merges filtered server-applied budget-window values into key update and regeneration audit payloads

_{Reviews (12): Last reviewed commit: "fix(proxy): record which users had spend..." | Re-trigger Greptile}

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing yryzhan-vitech:bugfix/reset-spend-on-budget-window-apply (075df61) with litellm_internal_staging (f4a68a7)}

[yryzhan-vitech]

@greptileai review

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

@greptileai review

[yryzhan-vitech]

Hand review of this PR (self-review, not a bot pass)

I read the whole diff myself rather than only reacting to bot findings. Two real defects came out of it — one from @greptileai's P1 on the rotation path, one I found on my own. Both are fixed in this push. Recording the rest so reviewers do not have to re-derive it.

Fixed as a result of this review

1. Rotation audit log dropped budget_reset_at (from the P1 above, a0ee6805ab)
/key/regenerate shares prepare_key_update_data, so a rotation carrying budget_duration persists spend=0.0 and a fresh budget_reset_at. The rotation audit record is built from GenerateKeyResponse, which declares spend and budget_duration but not budget_reset_at — so the armed window silently never reached the audit trail. Both hooks now share one serializer.

2. Bulk-update audit log did not record the spend reset (mine, 075df61058)
Same defect class, different endpoint, and the bot did not flag it. The all-users bulk audit entry logs non_default_values, which by design carries no spend key — the reset is a second, subset-scoped write. So an operator saw "budget window armed" with no record that spend was zeroed or for whom. Now emits spend_reset_user_ids.

Checked and deliberately left alone

• Cross-pod counter on the rotation path. I checked whether regenerate also needs the _sync_key_spend_counter call update_key_fn makes. It does not: the counter is keyed spend:key:{token_hash} and regenerate writes a new hash, so the post-rotation read is a clean miss and get_current_spend reseeds from the authoritative DB row. Syncing there would write a counter for an already-superseded token.
• Naive (tz-less) budget_reset_at. is_budget_window_newly_armed derives now from existing_reset_at.tzinfo, so a naive DB value compares against a naive now instead of raising TypeError. Intentional.
• Boundary reset_at == now arms the window (treated as expired). Deliberate: erring toward a reset is the safe direction, since the alternative locks an entity out of a window that has just rolled over.
• Read-then-write race in the bulk path. find_many → update_many(all) → update_many(armed subset) is not atomic, so a window armed by a concurrent request in that gap is not accounted for. This is the pre-existing shape of the endpoint and I did not widen it; a transaction here would be a separate change.
• Two-query split vs. one. I first tried splitting the main write into armed/not-armed halves. Reverted: it changes the failure semantics of the existing except fallback. The base bulk write is now untouched and the reset is additive after it.
• Empty-string budget_duration returns False (no arm), matching the len(budget_duration) > 0 guard the key path already had.

Verification

• 736 tests green across the five touched test files; the 5 unrelated failures in test_cache_settings_endpoints.py / test_coordination_redis_endpoints.py reproduce identically on the base with my changes stashed.
• Patch coverage: no uncovered lines in the changed ranges.
• ruff check / ruff format, ruff_strict_gate.py and type_discipline_gate.py all clean against this PR's base SHA. Note ANN401 and LIT001 both tightened under fbfb63c948 mid-review; the read-only payload arguments are now Mapping, and the three helpers that genuinely write back into the caller's payload carry a mutable-ok reason rather than a blanket suppression.
• Negative controls for every behavioural claim, mutated narrowly one mechanism at a time rather than by deleting the fix: removing only the subset reset query, only the counter invalidation, only the re-arm predicate, only the rotation merge, and only the bulk audit merge each fail exactly the test that asserts them.

[greptile-apps]

The analysis is accurate and the fix is correctly scoped. Looking at the current state of async_key_rotated_hook in key_management_event_hooks.py, commit a0ee6805ab added:

1. applied_values: Mapping[str, Any] | None = None parameter to async_key_rotated_hook
2. Shared _merge_applied_values serializer: KeyManagementEventHooks._merge_applied_values(base=response.model_dump(exclude_none=True), applied_values=applied_values)

Your measurement is correct — spend and budget_duration were already on GenerateKeyResponse so they were never lost; budget_reset_at was the only actual gap. The fix properly threads through the server-inferred values via applied_values without touching GenerateKeyResponse's schema, which is the right approach since that type is part of the public SDK surface.

Your note on the counter also holds: spend:key:{token_hash} is keyed on the token hash and regenerate_key_fn writes a new hash, so the post-rotation read is a guaranteed cache miss that reseeds from the DB row at spend=0.0 — no explicit invalidation needed.

Fix is correct. Resolving.

[yryzhan-vitech]

@greptileai