From 7656e26331ce43078c0bd44a0efca62d4b6a090b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 10:21:29 -0700 Subject: [PATCH 1/2] fix: align user and org spend checks with atomic counter pattern Brings user personal budget and organization budget enforcement in line with the existing key and team patterns, which already read spend from the atomic cross-pod Redis counter. --- litellm/proxy/auth/auth_checks.py | 28 +++++++++++++----- litellm/proxy/hooks/max_budget_limiter.py | 29 +++++++++++-------- .../proxy/hooks/proxy_track_cost_callback.py | 1 + litellm/proxy/proxy_server.py | 15 ++++++++++ 4 files changed, 54 insertions(+), 19 deletions(-) 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..4df28acc54 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -21,20 +21,25 @@ 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 + + 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, From c2b7c4bfcd25626d2042ae7595ebe15f0335b5cd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 10:38:08 -0700 Subject: [PATCH 2/2] fix: skip personal budget check in MaxBudgetLimiter for team-key requests --- litellm/proxy/hooks/max_budget_limiter.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 4df28acc54..7789fa6a34 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -27,6 +27,11 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): 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(