Merge pull request #26182 from BerriAI/litellm_budget_spend_counter_alignment

[Fix] Align user and org budget spend checks with atomic counter pattern

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}",
         )
 
 

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:

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:

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,