diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e19d04a260..2c8299e77a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -626,11 +626,17 @@ async def common_checks( # noqa: PLR0915 and user_object.max_budget is not None ): user_budget = user_object.max_budget - if user_budget < user_object.spend: + from litellm.proxy.proxy_server import get_current_spend + + user_spend = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + ) + if user_spend >= user_budget: raise litellm.BudgetExceededError( - current_cost=user_object.spend, + current_cost=user_spend, max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}", + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", ) ## 4.2 check team member budget, if team key @@ -3665,12 +3671,20 @@ async def _organization_max_budget_check( if org_max_budget is None or org_max_budget <= 0: return + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) + from litellm.proxy.proxy_server import get_current_spend + + org_spend = await get_current_spend( + counter_key=f"spend:org:{org_id}", + fallback_spend=org_table.spend or 0.0, + ) + # Check if organization spend exceeds max budget - if org_table.spend >= org_max_budget: + if org_spend >= org_max_budget: # Trigger budget alert call_info = CallInfo( token=valid_token.token, - spend=org_table.spend, + spend=org_spend, max_budget=org_max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, @@ -3686,9 +3700,9 @@ async def _organization_max_budget_check( ) raise litellm.BudgetExceededError( - current_cost=org_table.spend, + current_cost=org_spend, max_budget=org_max_budget, - message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}", + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", ) diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 4b59f603d3..7789fa6a34 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -21,20 +21,30 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): ): try: verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook") - cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id" - user_row = await cache.async_get_cache( - cache_key, parent_otel_span=user_api_key_dict.parent_otel_span + max_budget = user_api_key_dict.user_max_budget + user_id = user_api_key_dict.user_id + + if max_budget is None or user_id is None: + return + + # Personal budget applies only to non-team requests, matching + # the explicit team-key exemption in common_checks section 4.1. + if user_api_key_dict.team_id is not None: + return + + from litellm.proxy.proxy_server import get_current_spend + + curr_spend = await get_current_spend( + counter_key=f"spend:user:{user_id}", + fallback_spend=user_api_key_dict.user_spend or 0.0, ) - if user_row is None: # value not yet cached - return - max_budget = user_row["max_budget"] - curr_spend = user_row["spend"] - if max_budget is None: - return - - if curr_spend is None: - return + verbose_proxy_logger.debug( + "MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f", + user_id, + curr_spend, + max_budget, + ) # CHECK IF REQUEST ALLOWED if curr_spend >= max_budget: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ea9c92fec6..c9946f4e26 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -213,6 +213,7 @@ class _ProxyDBLogger(CustomLogger): team_id=team_id, user_id=user_id, response_cost=response_cost, + org_id=org_id, ) # update cache (fire-and-forget for backward compat: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8354a798b..0efa1d452d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1795,6 +1795,7 @@ async def increment_spend_counters( team_id: Optional[str], user_id: Optional[str], response_cost: Optional[float], + org_id: Optional[str] = None, ): """ Atomically increment spend counters for budget enforcement. @@ -1881,6 +1882,20 @@ async def increment_spend_counters( increment=response_cost, ) + if user_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:user:{user_id}", + source_cache_key=user_id, + increment=response_cost, + ) + + if org_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:org:{org_id}", + source_cache_key=f"org_id:{org_id}", + increment=response_cost, + ) + async def _init_and_increment_spend_counter( counter_key: str,