scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(proxy): restore AST-visible end-user budget check in custom auth

Browse source 
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.
This commit is contained in:
Taranum01 2026-09-13 21:13:47 +05:30
parent 4e433bc9a1
commit b939bb0b72
1 changed files with 16 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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