scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

fix(proxy-auth): enforce end-user budget checks on master-key admin paths

Browse source 
Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
suleimanelkhoury 2026-04-28 14:28:02 +02:00
parent 30823250fe
commit 78845a276e
1 changed files with 57 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -324,6 +324,40 @@ def update_valid_token_with_end_user_params(
return valid_token
async def _enforce_end_user_budget_if_needed(
valid_token: UserAPIKeyAuth,
end_user_object: Optional[LiteLLM_EndUserTable],
request_data: dict,
route: str,
llm_router: Optional[Any],
proxy_logging_obj: ProxyLogging,
) -> None:
model = get_model_from_request(request_data, route)
skip_budget_checks = False
if model is not None and llm_router is not None:
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
if skip_budget_checks:
return
end_user_mb = valid_token.end_user_max_budget
if end_user_mb is None and end_user_object is not None:
budget_table = end_user_object.litellm_budget_table
end_user_mb = budget_table.max_budget if budget_table is not None else None
if (
end_user_mb is not None
and end_user_object is not None
and valid_token.end_user_id is not None
):
await proxy_logging_obj.max_budget_limiter.is_end_user_within_budget(
end_user_id=valid_token.end_user_id,
end_user_max_budget=end_user_mb,
end_user_spend=end_user_object.spend,
route=route,
)
# Reusable coordinator for global spend to prevent cache stampede
_global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]")
@ -1081,6 +1115,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
_end_user_object.object_permission
)
await _enforce_end_user_budget_if_needed(
valid_token=valid_token,
end_user_object=_end_user_object,
request_data=request_data,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
)
return valid_token
if (
@ -1146,6 +1189,20 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
route=route,
start_time=start_time,
)
_user_api_key_obj = update_valid_token_with_end_user_params(
valid_token=_user_api_key_obj, end_user_params=end_user_params
)
await _enforce_end_user_budget_if_needed(
valid_token=_user_api_key_obj,
end_user_object=_end_user_object,
request_data=request_data,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
)
asyncio.create_task(
_cache_key_object(
hashed_token=hash_token(master_key),
@ -1155,10 +1212,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
)
_user_api_key_obj = update_valid_token_with_end_user_params(
valid_token=_user_api_key_obj, end_user_params=end_user_params
)
return _user_api_key_obj
## IF it's not a master key