fix: custom auth budget issue

[Harshit28j]

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Feb 26, 2026 9:51am

[Harshit28j]

@greptile please review this PR

[greptile-apps]

Greptile Summary

This PR adds end-user per-model budget enforcement for custom auth flows in the LiteLLM proxy. Previously, custom auth functions returning UserAPIKeyAuth would bypass all post-authentication checks (budget enforcement, token expiry, team/user validation). The fix introduces _run_post_custom_auth_checks which replicates the relevant checks from the normal key-based auth path.

Key changes:

• New _run_post_custom_auth_checks function in user_api_key_auth.py: Runs end-user DB lookup, token expiry, key-level and end-user-level model budget checks, user/team object lookups, and common_checks after custom auth returns a UserAPIKeyAuth
• New end_user_model_max_budget field added to LiteLLM_VerificationTokenView and threaded through the metadata pipeline so spend tracking works for end users
• Fixed update_valid_token_with_end_user_params: Now only overwrites token fields when DB-derived values are not None, preventing DB lookups from silently clearing budget values set by custom auth functions
• New is_end_user_within_model_budget method in _PROXY_VirtualKeyModelMaxBudgetLimiter for checking end-user per-model spend, mirroring the existing virtual-key pattern
• Spend tracking extended in async_log_success_event to increment end-user per-model spend in the cache
• Comprehensive mock-only tests covering budget enforcement, spend tracking, and the no-override behavior

Confidence Score: 4/5

• This PR is safe to merge with minor caveats around edge case handling in the custom auth post-check path.
• Score reflects well-structured changes that follow existing patterns, good test coverage with mock-only tests, and proper handling of the DB-override issue. The _run_post_custom_auth_checks function correctly replicates the important checks from the normal path. Minor deduction because the custom auth path doesn't include some checks present in the normal path (e.g., virtual key max budget check, soft budget alerts), though this appears intentional since custom auth tokens may not have virtual key budgets.
• litellm/proxy/auth/user_api_key_auth.py — the new _run_post_custom_auth_checks function is the most critical addition and should be carefully reviewed for parity with the normal auth path.

Important Files Changed

Filename	Overview
litellm/proxy/auth/user_api_key_auth.py	Core change: adds _run_post_custom_auth_checks and _lookup_end_user_and_apply_budget to run budget/auth checks after custom auth returns. Also fixes update_valid_token_with_end_user_params to not override custom-auth-set values with None. Well-structured and follows existing patterns.
litellm/proxy/hooks/model_max_budget_limiter.py	Adds is_end_user_within_model_budget and _get_end_user_spend_for_model methods mirroring the existing virtual key pattern. Also extends async_log_success_event to track end-user per-model spend. Implementation follows existing patterns correctly.
litellm/proxy/_types.py	Adds end_user_model_max_budget: Optional[dict] to LiteLLM_VerificationTokenView. Remaining changes are formatting-only (auto-formatter adjustments).
litellm/proxy/litellm_pre_call_utils.py	Passes end_user_model_max_budget into request metadata so async_log_success_event can track spend. Remaining changes are import formatting only.
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py	Adds tests for is_end_user_within_model_budget, _get_end_user_spend_for_model, and end-user budget tracking in async_log_success_event. All properly mocked with no network calls.
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py	New test file covering _run_post_custom_auth_checks and update_valid_token_with_end_user_params. Tests verify custom auth budget values aren't overridden by None DB values, and that budget exceeded errors propagate correctly. All mocked, no network calls.

Sequence Diagram

sequenceDiagram
    participant Client
    participant CustomAuth as Custom Auth Function
    participant PostChecks as _run_post_custom_auth_checks
    participant DB as DB / Cache
    participant BudgetLimiter as ModelMaxBudgetLimiter

    Client->>CustomAuth: Request with API key
    CustomAuth-->>PostChecks: UserAPIKeyAuth (with end_user_id, end_user_model_max_budget)

    PostChecks->>DB: Look up end_user object (if end_user_id set)
    DB-->>PostChecks: end_user_object + budget table

    Note over PostChecks: update_valid_token_with_end_user_params<br/>(only overwrite if DB value is not None)

    PostChecks->>PostChecks: Check token expiry

    PostChecks->>BudgetLimiter: is_key_within_model_budget (key-level)
    BudgetLimiter-->>PostChecks: OK / BudgetExceededError

    PostChecks->>BudgetLimiter: is_end_user_within_model_budget (end-user-level)
    BudgetLimiter->>DB: Get cached spend for end_user + model
    DB-->>BudgetLimiter: current_spend
    BudgetLimiter-->>PostChecks: OK / BudgetExceededError

    PostChecks->>PostChecks: common_checks (team, user, org budgets)
    PostChecks-->>Client: UserAPIKeyAuth / Error

Loading

_{Last reviewed commit: 39af85d}

[greptile-apps]

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[Harshit28j]

@greptile please review this PR