From b939bb0b72722c6f59a95da1210e0a96e7a10e6d Mon Sep 17 00:00:00 2001 From: Taranum01 Date: Sun, 13 Sep 2026 21:13:47 +0530 Subject: [PATCH] fix(proxy): restore AST-visible end-user budget check in custom auth Two structural-test invariants broke during the upstream merge: 1. `test_master_key_auth_sets_via_virtual_key_marker` expected `_user_api_key_obj.via_virtual_key = True` after `update_valid_token_with_end_user_params`. The conflict resolution ate that line. 2. `test_custom_auth_also_skips_budget_checks_for_zero_cost_models` walks the AST of `_run_post_custom_auth_checks` and asserts that `is_end_user_within_model_budget` is called directly inside it, guarded by `skip_budget_checks`. The helper-extraction refactor moved that call into `_enforce_end_user_model_max_budget_checks`, which hides it from this function's AST. Keep `_enforce_end_user_model_max_budget_checks` for the main auth path and master-key path (where it is the only caller and the helper is the right factoring). Inline the check back into `_run_post_custom_auth_checks` to match upstream's structure and satisfy the structural invariant without rewriting tests. --- litellm/proxy/auth/user_api_key_auth.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 092994dbcba..1761ba1275f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1943,6 +1943,7 @@ async def _user_api_key_auth_builder( _user_api_key_obj = update_valid_token_with_end_user_params( valid_token=_user_api_key_obj, end_user_params=end_user_params ) + _user_api_key_obj.via_virtual_key = True await _maybe_enforce_master_key_end_user_model_max_budget( valid_token=_user_api_key_obj, @@ -3538,12 +3539,21 @@ async def _run_post_custom_auth_checks( ) # 4. Check end-user model_max_budget - await _enforce_end_user_model_max_budget_checks( - valid_token=valid_token, - request_data=request_data, - route=route, - request=request, - ) + end_user_mmb: Final = valid_token.end_user_model_max_budget + if ( + not skip_budget_checks + and end_user_mmb is not None + and isinstance(end_user_mmb, dict) + and len(end_user_mmb) > 0 + and current_models + and valid_token.end_user_id is not None + ): + for model_name in current_models: + 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=model_name, + ) # team / user / end_user / project context objects are fetched by # the centralized common_checks gate in user_api_key_auth after