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

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

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,