[Fix] Align user and org budget spend checks with atomic counter pattern

[yuneng-berri]

Summary

• User personal budget checks in common_checks now read spend from the spend:user:{user_id} Redis counter via get_current_spend, matching the existing pattern used for key and team budget checks. The comparison was also corrected from < to >= to be consistent with how key budget enforcement works.
• Organization budget checks in _organization_max_budget_check now read spend from a spend:org:{org_id} Redis counter instead of the DB-loaded org_table.spend field.
• increment_spend_counters gains an org_id parameter and atomically increments two new counters (spend:user:{user_id}, spend:org:{org_id}) after each request, keeping them in sync with the existing key/team counters.
• The _PROXY_MaxBudgetLimiter pre-call hook is updated to read user_max_budget and user_id directly from user_api_key_dict (set at auth time) and look up current spend from the spend:user:{user_id} counter, replacing a stale cache lookup path that was effectively inoperative.

Testing

• Verified via live proxy: user at/above budget is blocked; user below budget passes; key budget enforcement unchanged.
• uv run pytest tests/test_litellm/proxy/auth/test_auth_checks.py tests/test_litellm/proxy/auth/test_organization_budget_enforcement.py — 78 passed
• uv run pytest tests/test_litellm/proxy/auth/test_multi_budget_windows.py tests/test_litellm/proxy/auth/test_team_member_budget.py — 10 passed

Type

🐛 Bug Fix
✅ Test

[greptile-apps]

Greptile Summary

This PR extends the atomic Redis counter budget-enforcement pattern (already used for API keys and teams) to user personal budgets and organization budgets. It also activates _PROXY_MaxBudgetLimiter, which was previously a no-op due to a stale cache-key lookup.

The previously flagged P1 concern — MaxBudgetLimiter enforcing personal budgets on team-key requests — is resolved: the hook now returns early when user_api_key_dict.team_id is not None, matching the exemption already present in common_checks.

Confidence Score: 5/5

Safe to merge — no P0/P1 issues remain; the prior P1 (team-key exemption in MaxBudgetLimiter) has been addressed.

All critical and major issues are resolved. The only remaining comment is a P2 clarification note about counter semantics (personal vs. total spend) that matches prior DB behaviour and is not a regression.

No files require special attention; proxy_server.py counter logic has a minor semantics note worth a code comment.

Important Files Changed

Filename	Overview
litellm/proxy/proxy_server.py	Adds org_id param to increment_spend_counters and two new counter increments (spend:user:{user_id}, spend:org:{org_id}). Logic mirrors existing key/team patterns; counter placement is after request so auth-time read uses the previous request's value — consistent with the rest of the counters.
litellm/proxy/auth/auth_checks.py	User personal-budget check now reads from spend:user:{user_id} Redis counter (Redis-first with DB fallback) and uses >= comparison, consistent with key/team enforcement. Org budget check similarly updated to read from spend:org:{org_id} counter.
litellm/proxy/hooks/max_budget_limiter.py	Replaces stale cache-key lookup (effectively always a no-op) with direct user_max_budget/user_id read from user_api_key_dict plus a Redis counter lookup. Adds the team-key exemption (team_id is not None → return) to match common_checks section 4.1, resolving the prior P1 concern.
litellm/proxy/hooks/proxy_track_cost_callback.py	Single-line addition to thread org_id through to increment_spend_counters. Change is minimal and correct.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Auth as auth_checks
    participant Redis as Redis spend counters
    participant Hook as MaxBudgetLimiter
    participant CB as CostCallback
    participant DB as Database

    Client->>Auth: API request
    Auth->>Redis: get_current_spend(spend:user:ID)
    Redis-->>Auth: user spend or DB fallback
    Auth->>Auth: over personal budget? raise BudgetExceededError
    Auth->>Redis: get_current_spend(spend:org:ID)
    Redis-->>Auth: org spend or DB fallback
    Auth->>Auth: over org budget? raise BudgetExceededError

    Auth-->>Hook: async_pre_call_hook
    Hook->>Hook: team key? skip check and return
    Hook->>Redis: get_current_spend(spend:user:ID)
    Redis-->>Hook: user spend
    Hook->>Hook: over budget? raise HTTP 429

    Hook-->>Client: proceed

    Client->>CB: success callback
    CB->>Redis: increment spend:key:TOKEN
    CB->>Redis: increment spend:team:TEAM
    CB->>Redis: increment spend:user:USER
    CB->>Redis: increment spend:org:ORG
    CB->>DB: update_database (async)

Loading

_{Reviews (2): Last reviewed commit: "fix: skip personal budget check in MaxBu..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 69.23077% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/hooks/max_budget_limiter.py	42.85%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!