From 6ae499ceb594106b7273143770b5e791d2f75cd8 Mon Sep 17 00:00:00 2001 From: suleimanelkhoury Date: Tue, 14 Apr 2026 11:43:32 +0200 Subject: [PATCH] fix(auth): implement end-user max budget checks in API key authentication Remove _check_end_user_budget functionality from auth_checks.py, which is being triggered each time get_end_user_object is called, causing all the models to be blocked (even the Zero Cost Ones), and move this functionality to max_budget_limiter.py, mimicing the max_model_budget_limiter.py functionality with handling exceeded model budget. This unifies all budget checks inside user_api_key_auth.py, make the code more streamlined, and with the combination with the previous commit solving None model cost values, it fixes the zero-cost models triggering budget exceeded. --- litellm/proxy/auth/auth_checks.py | 35 ----------------------- litellm/proxy/auth/user_api_key_auth.py | 30 +++++++++++++++++++ litellm/proxy/hooks/max_budget_limiter.py | 28 ++++++++++++++++++ 3 files changed, 58 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56958a88f6d..1fe99219138 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -866,35 +866,6 @@ async def _apply_default_budget_to_end_user( return end_user_obj -def _check_end_user_budget( - end_user_obj: LiteLLM_EndUserTable, - route: str, -) -> None: - """ - Check if end user is within their budget limit. - - Args: - end_user_obj: The end user object to check - route: The request route - - Raises: - litellm.BudgetExceededError: If end user has exceeded their budget - """ - if RouteChecks.is_info_route(route): - return - - if end_user_obj.litellm_budget_table is None: - return - - end_user_budget = end_user_obj.litellm_budget_table.max_budget - if end_user_budget is not None and end_user_obj.spend > end_user_budget: - raise litellm.BudgetExceededError( - current_cost=end_user_obj.spend, - max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}", - ) - - @log_db_metrics async def get_end_user_object( end_user_id: Optional[str], @@ -942,9 +913,6 @@ async def get_end_user_object( parent_otel_span=parent_otel_span, ) - # Check budget limits - _check_end_user_budget(end_user_obj=return_obj, route=route) - return return_obj # Fetch from database @@ -973,9 +941,6 @@ async def get_end_user_object( key="end_user_id:{}".format(end_user_id), value=_response.dict() ) - # Check budget limits - _check_end_user_budget(end_user_obj=_response, route=route) - return _response except Exception as e: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ebabf4cdb56..81e19eceada 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -318,6 +318,8 @@ def update_valid_token_with_end_user_params( valid_token.end_user_model_max_budget = end_user_params[ "end_user_model_max_budget" ] + if end_user_params.get("end_user_max_budget") is not None: + valid_token.end_user_max_budget = end_user_params["end_user_max_budget"] return valid_token @@ -1428,6 +1430,20 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 model=current_model, ) + # Check 5c. End user max budget + end_user_mb = valid_token.end_user_max_budget + 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, + ) + # Check 6: Additional Common Checks across jwt + key auth if valid_token.team_id is not None: try: @@ -1935,6 +1951,20 @@ async def _run_post_custom_auth_checks( model=current_model, ) + # 4b. Check end-user max_budget + end_user_mb = valid_token.end_user_max_budget + 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, + ) + # 5. Look up user object if user_id is set user_object = None if valid_token.user_id is not None: diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 4b59f603d3e..509d8ba4d7e 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -1,5 +1,6 @@ from fastapi import HTTPException +import litellm from litellm import verbose_logger from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -12,6 +13,33 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): def __init__(self): pass + async def is_end_user_within_budget( + self, + end_user_id: str, + end_user_max_budget: float, + end_user_spend: float, + route: str, + ) -> bool: + """ + Check if an end-user is within their overall max budget. + + Raises: + BudgetExceededError: If the end-user has exceeded overall budget. + """ + from litellm.proxy.auth.route_checks import RouteChecks + + if RouteChecks.is_info_route(route): + return True + + if end_user_spend > end_user_max_budget: + raise litellm.BudgetExceededError( + current_cost=end_user_spend, + max_budget=end_user_max_budget, + message=f"ExceededBudget: End User={end_user_id} over budget. Spend={end_user_spend}, Budget={end_user_max_budget}", + ) + + return True + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth,