fix(proxy): enforce budgets against authoritative DB spend when the cross-pod counter is stale

[yassin-berriai]

Relevant issues

Resolves LIT-3772

Linear ticket

https://linear.app/litellm-ai/issue/LIT-3772/virtual-key-spend-limits-are-not-being-enforced

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem

Type

🐛 Bug Fix

Changes

Budget enforcement reads spend through get_current_spend, which returns the cross-pod Redis counter whenever Redis hands back a value. In a multi-worker deployment, a Redis instance that restarts and reloads from an older RDB snapshot can come back serving a stale-low counter; that read is a hit, not a clean miss, so the existing DB reseed never fires and a key keeps getting admitted even though its recorded spend is already over max_budget. The symptom is recorded spend sitting above the limit while requests keep succeeding.

The fix gives get_current_spend an optional max_budget. When the counter would admit the request (counter under budget) and the counter sits below this caller's last-known recorded spend, it re-reads the authoritative DB spend and enforces against the higher value. The guard keeps the DB read off the hot path for healthy keys (a healthy counter tracks at or above recorded spend) and for freshly reset keys (recorded spend is back to 0 after a weekly reset, so there is nothing to floor against), so normal traffic and weekly resets are not over-blocked. The authoritative value is cached in-process for a few seconds, so a persistently stale counter (or a key whose auth-cache spend lags a reset for the 60s auth-cache TTL) drives at most one DB read per counter per window rather than one per request. The hard-budget checks pass their max_budget through across every counter type so none can be bypassed by a stale counter: keys, users, team members, teams, and organizations floor against the DB row; per-window key/team budgets floor against aggregated spend logs; end-user and tag budgets floor against the freshly-loaded recorded spend (they have no DB row). Soft-budget and alert reads keep the cheap counter-only behavior.

How it works

The budget check reads current spend through get_current_spend, which prefers the cross-pod Redis counter. The bug is that a value coming back from Redis is trusted unconditionally, even when Redis restarted and reloaded an older snapshot, so the counter can read lower than the spend already recorded in the database.

Before (over-budget request slips through):

flowchart TD
    R["Request on one of N workers"] --> GC["get_current_spend()"]
    GC --> RC{"Redis counter returns a value?"}
    RC -->|"hit, but stale-low after a Redis restart (e.g. 50)"| ST["current = 50"]
    RC -->|"clean miss"| DBR["reseed from DB (accurate)"]
    ST --> CHK{"current >= max_budget (100)?"}
    DBR --> CHK
    CHK -->|"No (50 < 100)"| ADMIT["admit -> request billed -> recorded spend climbs past budget"]
    CHK -->|"Yes"| REJ["raise BudgetExceededError -> 429"]

Loading

A stale counter never triggers the clean-miss reseed, so enforcement compares a too-low number against the limit and admits the request. Across workers, the ones that read an accurate value reject (429) while the ones that read the stale counter admit and overspend.

After (authoritative re-check before admitting):

flowchart TD
    R["Request on one of N workers"] --> GC["get_current_spend(max_budget)"]
    GC --> CUR["current = counter value (stale 50)"]
    CUR --> G{"max_budget set AND current < max_budget AND recorded_spend > current?"}
    G -->|"No: healthy counter or no budget"| USE["use counter as-is (NO DB read)"]
    G -->|"Yes: counter looks stale-low"| FL["_authoritative_floor_spend()"]
    FL --> CACHE{"DB value cached < 5s?"}
    CACHE -->|"yes"| CV["reuse cached DB value"]
    CACHE -->|"no"| ONE["one DB read, cache for 5s"]
    CV --> MAX["spend = max(counter, DB) = 150"]
    ONE --> MAX
    MAX --> RPR["repair: set-max counter -> 150 (atomic, monotonic) so other workers read the corrected value"]
    USE --> CHK{"spend >= max_budget (100)?"}
    RPR --> CHK
    CHK -->|"Yes (150 >= 100)"| REJ["raise BudgetExceededError -> 429"]
    CHK -->|"No"| ADMIT["admit"]

Loading

The re-check only fires when the counter would admit the request and reads below the spend this worker already knows about (the recorded_spend > current guard), so healthy counters and freshly reset keys stay off the DB path. The authoritative value is cached for 5s, so a persistently stale counter costs at most one DB read per counter per window per worker, not one per request. When the DB value is higher than the counter, the counter is also repaired so a stale-low counter is corrected for every worker (including one whose own cached spend is stale) instead of being re-derived per request. The repair is monotonic and atomic: RedisCache.async_set_max raises the counter only when the new value exceeds the stored one, via a single Lua GET/compare/SET, so a slightly-stale repair can neither lower the counter nor clobber a concurrent increment that racing requests rely on.

Fail-open hardening (reconcile + opt-in fail-closed)

Two further gaps in the fail-open behavior, where the Redis counter is left unset (nil / TTL -2) after a reload:

The post-call reservation reconcile, when it found the counter missing or an adjustment that would drive it negative, deleted the counter and continued ("invalidating reserved counters and continuing"). That deletion is what left counters nil and unenforced after a Redis reload. It now reseeds from the DB's lagging authoritative floor instead of deleting; the monotonic set-max can only raise a stale-low counter, and the read-time floor converges to the true total as the spend buffer flushes. The pre-call admission resize keeps its original fail-closed behavior.

A new opt-in general_settings.fail_closed_budget_enforcement (default false) rejects a request with 503 when the spend backing an admit decision can be verified against neither Redis nor the database, rather than admitting on an unverifiable budget. With it off (default) behavior is unchanged; end-user/tag callers whose spend is loaded fresh from the DB in auth are treated as verified and are not rejected.

Screenshots / Proof of Fix

Every enforcement path was exercised on a live proxy (Postgres + Redis, real gpt-4o-mini traffic, 5433/6380 local containers). For each budgeted entity a budget was driven to its limit, then a Redis restart was simulated by setting the entity's counter to a stale-low value while the recorded spend stayed over budget.

Behavior	Result
Key max_budget, counter accurate	drove to limit, 429 at the boundary
Key, stale-low Redis counter (restart sim)	429; counter repaired up to the authoritative value
Key, Redis container fully stopped	429 enforced via Postgres
Healthy under-budget key	200 (not over-blocked)
Team max_budget, stale counter (strict)	429 via DB floor; counter repaired
End-user max_budget, stale counter	429 via fresh recorded-spend fallback
Tag max_budget, stale counter	429 via fresh recorded-spend fallback
Per-window budget_limits, stale counter	429 via spend-logs aggregate (after the spend-log flush)
fail_closed_budget_enforcement on, Redis + Postgres both stopped	503 "Budget enforcement unavailable: current spend could not be verified..."

Representative transcripts:

Key over budget while the Redis counter is stale-low (Redis up):

DB spend = 0.0002277,  max_budget = 0.0002
redis-cli SET 'spend:key:<hash>' 0.00001
POST /v1/chat/completions  ->  HTTP 429
{"message":"Budget has been exceeded! Current cost: 0.0002277, Max budget: 0.0002", ...}
# counter repaired: redis-cli GET 'spend:key:<hash>' -> 0.0002277

Budget enforced via Postgres with Redis fully stopped:

docker stop litfix-redis
# over-budget key  -> HTTP 429  ("Current cost: 0.0002277, Max budget: 0.0002")
# fresh under-budget key -> HTTP 200

Team budget, stale counter (strict mode reads the DB floor):

redis-cli SET 'spend:team:<id>' 0.00001
POST /v1/chat/completions  ->  HTTP 429
{"message":"Budget has been exceeded! Team=<id> Current cost: 0.00024, Max budget: 0.0002"}

Fail-closed when nothing can verify spend (Redis + Postgres both down, flag on):

docker stop litfix-redis litfix-pg
POST /v1/chat/completions  ->  HTTP 503
"Budget enforcement unavailable: current spend could not be verified against Redis or the database,
 and fail_closed_budget_enforcement is enabled, so the request was rejected ..."

Unit coverage in tests/test_litellm/proxy/proxy_server/test_spend_counters.py and tests/test_litellm/proxy/test_budget_reservation.py pins each path: the stale-low floor and monotonic repair, the healthy-counter fast path, the no-max_budget bypass, the post-reset admit, the strict-mode floor that closes the both-stale gap, the fail-closed 503, and the reconcile reseed-instead-of-delete. The core tests fail on unfixed code.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[yassin-berriai]

@greptileai

[codecov]

Codecov Report

❌ Patch coverage is 87.20930% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/caching/redis_cache.py	16.66%	10 Missing ⚠️
litellm/proxy/spend_tracking/budget_reservation.py	90.90%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes a budget enforcement gap where a Redis counter that survived a restart and reloaded from an older RDB snapshot returns a stale-low value — a cache hit, not a miss — so the existing DB reseed never fires and over-budget keys keep getting admitted. The fix adds an optional max_budget parameter to get_current_spend that re-checks the authoritative DB floor before admitting an under-budget request, with a 5s in-process cache to bound the extra DB reads, a monotonic async_set_max repair to propagate the correction to other workers, and an opt-in fail_closed_budget_enforcement flag that rejects unverifiable spends with 503 instead of admitting.

• Stale counter floor: _authoritative_floor_spend reads the DB row (or spend-log aggregate for window counters) when the counter would admit the request and looks stale-low; result is cached 5s per counter per worker and repaired into Redis via atomic Lua GET/compare/SET.
• Reconcile reseed: post-call reservation reconcile now reseeds from the DB floor instead of deleting the counter and failing open when the counter is missing or an adjustment would underflow, closing the fail-open that left budgets unenforced after a Redis reload.
• Fail-closed strict mode: general_settings.fail_closed_budget_enforcement (default off) closes the both-stale gap by always running the authoritative re-check, rejecting with 503 when neither Redis nor the DB can verify spend.

Confidence Score: 5/5

Safe to merge. The fix correctly addresses a real enforcement bypass without changing default behavior for healthy counters, per-reset keys, or soft-budget reads, and all new code paths are covered by unit tests that fail on unfixed code.

The implementation is consistent and well-bounded: the authoritative re-check is guarded by current < max_budget (already-rejected traffic never hits the DB), and by the 5s in-process cache (persistently stale counters cost at most one extra DB read per counter per worker per 5s window, not one per request). The monotonic async_set_max Lua script is atomic and correctly short-circuits the EXPIRE when TTL is 0. The SpendCounterReseed.from_db short-circuits immediately for end-user and tag prefixes, so those paths add no DB round-trips. The reconcile behavioral change (reseed-from-DB instead of delete-and-raise) is consistent with the stated goal and the updated tests correctly capture the new contract. No test assertions were weakened.

No files require special attention. litellm/proxy/proxy_server.py has the highest complexity (the new get_current_spend branching tree and its helpers) but the logic is thoroughly covered by the 12 new unit tests in test_spend_counters.py.

Important Files Changed

Filename	Overview
litellm/caching/redis_cache.py	Adds async_set_max: an atomic Lua-based monotonic write that only raises a counter, never lowers it. TTL defaults to 60s (matching async_increment behavior) and is conditionally applied only when ARGV[2] > 0, preserving correct TTL semantics. Circuit-breaker guard applied.
litellm/proxy/proxy_server.py	Core of the fix: get_current_spend gains an optional max_budget parameter that triggers a re-read against the authoritative DB floor before admitting an under-budget request. Adds _authoritative_floor_spend (5s in-process cache, falls back to spend-logs aggregates for window counters), _repair_stale_spend_counter (monotonic Redis write via async_set_max), reseed_spend_counter_from_db (reseeds on reconcile inconsistency), and the opt-in fail_closed_budget_enforcement 503 path.
litellm/proxy/auth/auth_checks.py	All get_current_spend call sites updated to pass max_budget; end-user and tag callers also set fallback_authoritative=True. Window budget callers pass window_entity_type, window_entity_id, and window_start so stale-low window counters can be floored against spend-log aggregates.
litellm/proxy/spend_tracking/budget_reservation.py	Replaces fail-and-delete on inconsistent reservation counter with reseed-from-DB for post-call reconcile, and keeps fail-closed behavior (raise) for pre-call resize. _ensure_counter_can_apply_adjustment refactored to _counter_can_apply_adjustment returning bool.
litellm/proxy/auth/user_api_key_auth.py	One-line change: max_budget=team_member_budget threaded through to get_current_spend for the in-auth team member budget check.
tests/test_litellm/proxy/proxy_server/test_spend_counters.py	Adds 12 new unit tests covering all new code paths: stale-low floor, healthy fast path, no-budget bypass, post-reset admit, 5s cache coalescing, end-user/tag fallback, window spend-log floor, fail-closed 503, fail-closed pass-through for authoritative fallback, and strict-mode both-stale gap.
tests/test_litellm/proxy/test_budget_reservation.py	Three tests renamed and updated to verify new reconcile behavior (reseed-from-DB instead of delete-and-raise). Changes correctly reflect the intentional behavioral shift from fail-open deletion to fail-safe reseeding.
tests/test_litellm/proxy/test_proxy_server.py	Test updated to verify the new reseed-to-DB-floor behavior (counter becomes 0.6, not 0.25 direct-increment). Legitimate behavioral change, not a weakened assertion.
tests/test_litellm/proxy/auth/test_auth_checks.py	All mock get_current_spend signatures updated to accept max_budget=None, **kwargs so they stay compatible with the new call-site signature. No assertion logic changed.
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py	Mock signature updated to accept the new max_budget kwarg. No logic changes.
tests/test_litellm/proxy/auth/test_multi_budget_windows.py	Mock signatures updated. No assertion changes.
tests/test_litellm/proxy/test_litellm_pre_call_utils.py	Mock signatures updated to accept new kwargs. No assertion changes.

_{Reviews (7): Last reviewed commit: "fix(proxy): enforce budgets against auth..." | Re-trigger Greptile}

[yassin-berriai]

Good catch on the per-request DB read. Addressed in the latest commit: the authoritative re-check now caches the DB spend in-process for a short window (SPEND_DB_FLOOR_CACHE_TTL_SECONDS, 5s) via _authoritative_floor_spend, so a persistently stale counter (or a key whose auth-cache spend lags a reset for the 60s auth-cache TTL) drives at most one DB read per counter per window instead of one per request, and concurrent requests on a pod share that cached value. I kept the read-time floor rather than reseeding the Redis counter on purpose: an INCRBYFLOAT-based heal has no cross-pod coordination and can double-apply across replicas, which would over-count and over-block; the bounded read keeps enforcement correct without that risk. Added test_get_current_spend_floor_caches_db_read to pin the single-read behavior.

@greptileai

[greptile-apps]

The caching approach in _authoritative_floor_spend is sound. A few observations on the implementation:

1. Cache key namespace is clean — spend_db_floor:{counter_key} shares the same in_memory_cache as the spend counter but the prefix prevents collision.

2. Asyncio race is benign — two concurrent coroutines can both see cached is None before the first await from_db completes and both make the DB call. In asyncio's single-threaded model this is bounded (not unbounded like with threads) and both calls return the same authoritative value, so the extra read is harmless. Worth a comment if you want, but not a bug.

3. Post-reset false-positive window — if the floor cache holds a stale-high value (set just before a budget reset fires), requests can be blocked for up to 5s after the counter resets to 0. The test test_get_current_spend_floor_admits_after_reset covers the case where the cache is cold at reset time, but not the warm-cache path. That 5s window is an acceptable trade-off and you've documented it, just worth being explicit about in the comments.

4. fallback_spend as the staleness signal — using the auth-object spend (from the 60s auth cache) as the floor heuristic is correct and avoids a DB read for healthy counters. The only edge: if fallback_spend is also stale-low for some independent reason (e.g., a bug in the auth cache), the guard could silently pass. But that's outside the scope of this fix and fallback_spend still acts as the reseed baseline in the existing _read_spend_counter_estimate path, so it's no worse than before.

The logic is correct, the hot path stays off the DB for healthy keys, and the new test pins the single-read guarantee. LGTM.

[veria-ai]

PR overview

This pull request updates the proxy budget-enforcement path so cross-pod spend counters are checked against a more authoritative database-backed view when cached or Redis-backed spend data may be stale.

One budget-enforcement gap remains open: if Redis has a stale spend value that is below the configured budget but above an older cached fallback value, the code can skip the authoritative database floor and allow an already over-budget user to continue making LLM requests. That creates a bounded budget-bypass/resource-consumption risk until the counter is repaired or refreshed. Three prior issues have already been addressed, so the remaining risk is narrowed to this stale-counter edge case.

Open issues (1)

• Medium: Budget bypass when stale Redis is above cached fallback — litellm/proxy/proxy_server.py:2105

Fixed/addressed: 3 · PR risk: 6/10

[yassin-berriai]

@greptileai please re-review the latest commit and post an updated confidence score; the per-request DB read concern from the first review was addressed with the in-process floor cache (_authoritative_floor_spend).

[yassin-berriai]

Addressed the open issue (stale counter not repaired) in the latest commit.

When the authoritative DB spend comes back higher than the counter, get_current_spend now writes that value back via _repair_stale_spend_counter, so a counter left stale-low by a Redis restart is corrected for every worker rather than only used for the current request's check. This closes the bypass where a worker whose own cached spend is also stale could keep admitting against the stale counter; once any worker repairs it, all workers read the corrected total from Redis.

The repair uses an idempotent SET, not an increment, so concurrent repairs across pods converge to the same DB value instead of stacking, and it only runs when the counter reads strictly below the DB spend, so it can only raise a stale counter toward the recorded spend and never lowers a healthy one. Covered by test_get_current_spend_floors_stale_low_counter_against_db (asserts the write-back) and test_get_current_spend_floor_admits_after_reset (asserts no write when the counter already matches the DB).

[yassin-berriai]

Addressed the remaining open issue (repair could overwrite concurrent increments) in the latest commit.

The counter repair is now monotonic and atomic. Instead of a plain SET, it goes through RedisCache.async_set_max, which runs a single Lua GET/compare/SET: it only writes when the new value is greater than the stored value, so the GET-compare-SET cannot interleave with a concurrent increment, and a slightly-stale repair can neither lower the counter nor clobber an increment that a racing request already applied. The in-memory copy uses the same raise-only guard with no await between the read and the write, so it is atomic within the worker. The repair still only runs when the counter reads strictly below the authoritative DB spend, so it can only move the counter up toward recorded spend.

[yassin-berriai]

@greptileai please re-review the latest commit (HEAD); the repair was made monotonic/atomic via RedisCache.async_set_max to address the concurrent-increment concern.

[yassin-berriai]

Addressed the remaining open issue (unsupported counter floors) in the latest commit.

The authoritative floor now covers every hard-budget counter type rather than only the DB-row ones:

• key / user / team / team-member / org: floor against the DB row (unchanged).
• per-window key/team (spend:key:*:window:*, spend:team:*:window:*): floor against aggregated spend logs via SpendCounterReseed.window_from_spend_logs, with window context (entity type/id and window start) threaded from the multi-budget checks. Bounded by the same in-process cache, and the corrected value is repaired back monotonically.
• end-user / tag (spend:end_user:*, spend:tag:*): these have no DB row, so enforcement floors against the caller's fallback_spend, which the auth path already loads fresh from the end-user / tag object. No extra DB read, and a stale-low counter can no longer admit them past budget.

So a tag, end-user, or window that is already over budget is rejected after a Redis reload of a stale-low value, the same as the primary counters. Covered by new tests in test_spend_counters.py (...floors_end_user_tag_against_fallback, ...floors_window_against_spend_logs).

[yassin-berriai]

@greptileai re-review latest commit please; extended the authoritative floor to end-user, tag, and per-window counters so no hard-budget counter type is bypassed by a stale Redis value.

[yassin-berriai]

On the window-counter GROUP BY: this is an accepted trade-off rather than an oversight. Primary counters have a cheap staleness signal (recorded_spend > counter) so healthy traffic skips the DB, but per-window counters have no equivalent cheap signal (their fallback is 0), so closing the same stale-counter bypass for window budgets requires the spend-logs aggregation. It is bounded to at most one aggregation per window counter per 5s per worker by the in-process cache, and window budgets are coarse (hourly/daily/weekly), so the load is low and the staleness is acceptable. Keeping it so per-window budgets get the same stale-counter protection as primary counters; the alternative (no window floor) leaves window budgets bypassable after a Redis reload.

[yassin-berriai]

Pushed two further hardening changes from a customer incident (Redis counter nil/TTL -2, a $100 key reached $218):

1. Reconcile no longer fails open. The post-call reservation reconcile used to delete the counter and continue when it found the counter missing or an adjustment that would drive it negative (the deletion left counters unenforced after a Redis reload). It now reseeds from the DB's lagging authoritative floor instead of deleting; the monotonic set-max only ever raises a stale-low counter, and the read-time floor converges to the true total as the spend buffer flushes. The pre-call admission resize path deliberately keeps failing closed.

2. Opt-in general_settings.fail_closed_budget_enforcement (default false): rejects with 503 when spend can be verified against neither Redis nor the DB, instead of admitting on an unverifiable budget. Default behavior unchanged.

@greptileai please re-review the latest commit.

[veria-ai]

Medium: Budget bypass when stale Redis is above cached fallback

fallback_spend comes from cached auth/team/user/tag objects, so it can lag the DB independently of the Redis spend counter. If Redis restores 50, the cached token still says 40, and the DB row is already over a 100 budget at 120, this branch skips the DB floor and a user who is already over budget can keep making LLM requests until another path repairs the counter; consult the cached _authoritative_floor_spend whenever an enforced counter is below max_budget, or otherwise use a signal that does not depend on a potentially stale fallback being greater than the Redis value.

[yassin-berriai]

Latest commit closes the "stale Redis above cached fallback" gap and adds full live coverage.

The both-stale gap (counter stale-low AND the caller's cached fallback also stale-low, so the cheap fallback_spend > current guard skips the DB re-check) is real; live testing reproduced it on the team path. With fail_closed_budget_enforcement (strict mode) on, the authoritative re-check now runs for every budgeted entity, so a stale-low counter can no longer slip through regardless of the cached fallback; it stays bounded by the in-process cache. Default mode keeps the cheap guard for the hot path. Anyone who needs the hard guarantee enables the flag (which is also what arms the 503 fail-closed behavior).

Every enforcement path is now live-verified on a proxy (key, team, end-user, tag, per-window, Redis fully stopped -> enforced via Postgres, and the fail-closed 503 with Redis+Postgres both down). The matrix and transcripts are in the PR description.

@greptileai please re-review the latest commit.

[mateo-berri]

seems like fail_closed_budget_enforcement needs to be documented. Please re-request review after. Also, just merge internal staging to get the litellm on windows circleci fix to pass required ci

[yassin-berriai]

Note on the remaining Veria item ("budget bypass when the stale Redis counter is above the cached fallback"): this is the narrow race where a worker's Redis counter and its per-pod cached spend are both stale-low at the same time, so the cheap hot-path guard (fallback_spend > current) skips the DB re-check. Closing it unconditionally means running the authoritative DB read on every budgeted request, which is the hot-path cost the review flagged elsewhere; the two are the same tradeoff from opposite sides.

This is handled by design: fail_closed_budget_enforcement (strict mode) makes the authoritative re-check run for every budgeted entity, closing the gap entirely, and also arms the 503-on-unverifiable behavior. Deployments that want the hard guarantee enable that flag (the same flag the budget-safety ask here is about). Default mode keeps the cheap guard so healthy under-budget traffic stays off the DB. Both paths are live-verified and unit-pinned.

[yassin-berriai]

BerriAI/litellm-docs#371

[mateo-berri]

LGTM; thanks!