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

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

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:

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,