fix: enforce per-model budget checks on cached admin and master key auth paths - #25368
fix: enforce per-model budget checks on cached admin and master key auth paths#25368jk-f5 wants to merge 2 commits into
Conversation
…uth paths The _user_api_key_auth_builder function has two early-return paths that skip the main budget-check block entirely: 1. Cached PROXY_ADMIN key (returns before line 1234 budget checks) 2. Master key (returns before line 1234 budget checks) Both paths call update_valid_token_with_end_user_params, which copies end_user_model_max_budget onto the token, but then return immediately without checking those budgets. This means per-model budget limits configured on end users (or on admin keys) are never enforced when requests come in via the master key or a cached admin key. This commit adds a _check_model_max_budget helper that runs both key-level and end-user-level per-model budget checks, and calls it on both early-return paths. The helper uses the same guard conditions and limiter calls as the existing budget-check block and _run_post_custom_auth_checks, and only runs on LLM API routes when a model is present in the request. The underlying limiter methods are in-memory cache lookups only, so there is no performance impact. Includes 16 new tests covering the helper (9 unit) and both auth paths (7 integration).
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes a per-model budget enforcement bypass on two early-return auth paths: cached Confidence Score: 5/5Safe to merge — targeted bug fix with no new P0/P1 issues found. The No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/auth/user_api_key_auth.py | Adds _check_model_max_budget helper and inserts calls on both early-return auth paths; guard conditions faithfully mirror the main budget-check block (including prisma_client is not None for key-level and omitting it for end-user level, matching the existing pattern). |
| tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py | New mock-only test file covering the helper directly and both early-return paths; tests budget exceeded propagation, non-LLM route skipping, missing model skipping, and positive/negative cases for both key-level and end-user budgets. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[_user_api_key_auth_builder] --> B{Cached PROXY_ADMIN key?}
B -- Yes --> C[update_valid_token_with_end_user_params]
C --> D[_check_model_max_budget NEW]
D --> E{LLM route AND model present?}
E -- No --> F[return valid_token]
E -- Yes --> G{model_max_budget set?}
G -- Yes --> H[is_key_within_model_budget]
G -- No / checked --> I{end_user_model_max_budget set?}
H --> I
I -- Yes --> J[is_end_user_within_model_budget]
I -- No / checked --> F
J --> F
B -- No --> K{Master key?}
K -- Yes --> L[_return_user_api_key_auth_obj]
L --> M[update_valid_token_with_end_user_params]
M --> N[_check_model_max_budget NEW]
N --> O[return _user_api_key_obj]
K -- No --> P[Normal key path - main budget-check block]
Reviews (2): Last reviewed commit: "fix: add missing prisma_client guard to ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if not RouteChecks.is_llm_api_route(route=route): | ||
| return | ||
|
|
||
| current_model = request_data.get("model", None) | ||
| if current_model is None: | ||
| return | ||
|
|
||
| # Key-level model_max_budget | ||
| max_budget_per_model = valid_token.model_max_budget | ||
| if ( | ||
| max_budget_per_model is not None | ||
| and isinstance(max_budget_per_model, dict) | ||
| and len(max_budget_per_model) > 0 | ||
| and valid_token.token is not None | ||
| ): | ||
| await model_max_budget_limiter.is_key_within_model_budget( | ||
| user_api_key_dict=valid_token, | ||
| model=current_model, | ||
| ) | ||
|
|
||
| # End-user model_max_budget | ||
| end_user_mmb = valid_token.end_user_model_max_budget | ||
| if ( | ||
| end_user_mmb is not None | ||
| and isinstance(end_user_mmb, dict) | ||
| and len(end_user_mmb) > 0 | ||
| and valid_token.end_user_id is not None | ||
| ): | ||
| await model_max_budget_limiter.is_end_user_within_model_budget( | ||
| end_user_id=valid_token.end_user_id, | ||
| end_user_model_max_budget=end_user_mmb, | ||
| model=current_model, | ||
| ) |
There was a problem hiding this comment.
prisma_client is not None guard for key-level budget check
The main budget-check block (lines 1416–1428) guards key-level model budget checks with and prisma_client is not None, but _check_model_max_budget omits this guard. In a no-database deployment, the cached PROXY_ADMIN or master-key path will now call is_key_within_model_budget unconditionally, while the normal key path still skips it. is_key_within_model_budget itself only touches dual_cache, so this likely won't crash — but the inconsistency means the two paths behave differently in db-less setups and could cause surprises if that assumption changes.
To match the existing guard, prisma_client would need to be passed into _check_model_max_budget as a parameter and checked in the key-level block before calling the limiter.
Rule Used: What: avoid backwards-incompatible changes without... (source)
| # Per-model budget check for cached PROXY_ADMIN keys | ||
| await _check_model_max_budget( | ||
| valid_token=valid_token, | ||
| request_data=request_data, | ||
| route=route, | ||
| model_max_budget_limiter=model_max_budget_limiter, | ||
| ) | ||
|
|
There was a problem hiding this comment.
Per the team's style guide, behavior changes in the critical auth path should be gated behind a user-controlled flag to avoid breaking existing deployments. Before this PR, master key and cached PROXY_ADMIN paths never enforced per-model budgets. Users who have model_max_budget configured on end users or admin keys, have exceeded those budgets, and rely on master/admin key access will start receiving 400s after upgrading with no warning.
This affects both insertion points — lines 1081–1088 (cached admin) and 1167–1174 (master key). The safer approach, matching the pattern from prior auth path changes, is to guard these calls behind a flag (e.g. litellm.enforce_model_budget_on_admin_paths) so existing deployments can opt in rather than being silently broken on upgrade.
Rule Used: What: avoid backwards-incompatible changes without... (source)
There was a problem hiding this comment.
This seems like a bad take to me. If someone deliberately configured model budgets and they're being silently bypassed, they'll want this fix.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
…onsistency with main budget-check block
|
Tip: Greploops — Automatically fix all review issues by running Use the Greptile plugin for Claude Code to query reviews, search comments, and manage custom context directly from your terminal. |
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. |
1 participant
Relevant issues
Fixes per-model budget enforcement bypass on cached PROXY_ADMIN and master key auth paths.
Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
🐛 Bug Fix
Changes
Problem
_user_api_key_auth_builderhas two early-return paths that skip the main budget-check block entirely:Cached PROXY_ADMIN key: When a PROXY_ADMIN key is found in cache,
update_valid_token_with_end_user_paramscopiesend_user_model_max_budgetonto the token, then returns immediately without checking per-model budgets.Master key When the request uses the master key,
update_valid_token_with_end_user_paramscopies end-user budget params onto the token, then returns immediately without checking per-model budgets.Both paths bypass the budget checks that call
model_max_budget_limiter.is_key_within_model_budgetandmodel_max_budget_limiter.is_end_user_within_model_budget.This means per-model budget limits configured on end users (or on admin keys with
model_max_budget) are never enforced when requests come in via the master key or a cached admin key.Fix
Added a
_check_model_max_budgethelper function that runs both key-level and end-user-level per-model budget checks. It uses the same guard conditions and limiter calls as the existing budget-check block and_run_post_custom_auth_checks:RouteChecks.is_llm_api_route(route)so non-LLM routes are not affectedrequest_data.get("model")being presentInserted calls to this helper on both early-return paths, after
update_valid_token_with_end_user_paramsbut before the return statement.Performance impact
None. The helper calls
is_key_within_model_budget/is_end_user_within_model_budget, which are in-memory cache lookups (dict lookups +DualCache.async_get_cache). The guard conditions mean the limiter methods are only called when budget dicts are actually populated. For requests without per-model budgets configured, the function returns after checkingmodel_max_budget is Noneorlen(...) == 0.