fix(proxy): use atomic increment for global proxy spend to prevent false BudgetExceededError after reset

Replace the non-atomic read-modify-write on the global proxy spend
cache key with async_increment_cache, which uses Redis INCRBYFLOAT.
This prevents a stale pre-reset value from being resurrected when
ResetBudgetJob deletes the key while concurrent spend updates are in
flight.

Also change ResetBudgetJob to set the key to 0.0 instead of deleting
it, removing the cache-miss window entirely.

Adds regression tests covering the reset-then-increment scenario and
concurrent updates after reset.

litellm/proxy/common_utils/reset_budget_job.py

@@ -262,11 +262,21 @@ class ResetBudgetJob:
 
     @staticmethod
     async def _invalidate_global_proxy_spend_cache() -> None:
-        """Drop the cached global-proxy spend accumulator after the proxy
-        budget aggregate row is reset, so the next auth-time load reads the
-        zeroed row instead of a stale (potentially never-expiring) counter.
+        """Zero the cached global-proxy spend accumulator after the proxy
+        budget aggregate row is reset so atomic increments start from a known floor.
         """
-        await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY)
+        try:
+            from litellm.proxy.proxy_server import user_api_key_cache
+
+            await user_api_key_cache.async_set_cache(
+                key=GLOBAL_PROXY_SPEND_CACHE_KEY,
+                value=0.0,
+            )
+        except Exception as e:  # noqa: BLE001  # cache zero failure must not break budget reset
+            verbose_proxy_logger.warning(
+                "Failed to zero global proxy spend cache after budget reset: %s",
+                e,
+            )
 
     @staticmethod
     async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:

litellm/proxy/proxy_server.py

@@ -3079,7 +3079,7 @@ async def update_cache(
 
     ### UPDATE USER SPEND ###
     async def _update_user_cache():
-        ## UPDATE CACHE FOR USER ID + GLOBAL PROXY
+        ## UPDATE CACHE FOR USER ID
         if response_cost is None:
             return
         user_ids: Final = [user_id]
@@ -3110,14 +3110,6 @@ async def update_cache(
                         CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable),
                     )
                 )
-            ## UPDATE GLOBAL PROXY ##
-            global_proxy_spend: Final = await user_api_key_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
-            if global_proxy_spend is None:
-                # do nothing if not in cache
-                return
-            elif response_cost is not None and global_proxy_spend is not None:
-                increment: Final = global_proxy_spend + response_cost
-                values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment))
         except Exception as e:
             verbose_proxy_logger.warning(
                 "Spend tracking - failed to update user spend in cache. "
@@ -3129,6 +3121,29 @@ async def update_cache(
                 traceback.format_exc(),
             )
 
+    ### UPDATE GLOBAL PROXY SPEND ###
+    async def _update_global_proxy_spend() -> None:
+        if response_cost is None:
+            return
+        try:
+            # Atomic increment avoids the read-modify-write race with
+            # ResetBudgetJob and starts from 0 when the key is missing.
+            await user_api_key_cache.async_increment_cache(
+                key=GLOBAL_PROXY_SPEND_CACHE_KEY,
+                value=response_cost,
+                ttl=get_management_object_ttl(user_api_key_cache),
+                refresh_ttl=True,
+            )
+        except Exception as e:  # noqa: BLE001  # spend update must not fail the request
+            verbose_proxy_logger.warning(
+                "Spend tracking - failed to update global proxy spend in cache. "
+                "Budget enforcement may use stale spend values. "
+                "response_cost=%s - %s\n%s",
+                response_cost,
+                str(e),
+                traceback.format_exc(),
+            )
+
     ### UPDATE END-USER SPEND ###
     async def _update_end_user_cache():
         if end_user_id is None or response_cost is None:
@@ -3272,6 +3287,8 @@ async def update_cache(
     if user_id is not None:
         await _update_user_cache()
 
+    await _update_global_proxy_spend()
+
     if end_user_id is not None:
         await _update_end_user_cache()
 
@@ -3281,27 +3298,15 @@ async def update_cache(
     if tags is not None:
         await _update_tag_cache()
 
-    global_proxy_spend_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
-    local_object_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key)
-    shared_scalar_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key)
-
-    if local_object_updates:
+    if values_to_update_in_cache:
         asyncio.create_task(
             user_api_key_cache.async_set_cache_pipeline(
-                cache_list=list(local_object_updates),
+                cache_list=list(values_to_update_in_cache),
                 ttl=get_management_object_ttl(user_api_key_cache),
                 litellm_parent_otel_span=parent_otel_span,
                 local_only=True,
             )
         )
-    if shared_scalar_updates:
-        asyncio.create_task(
-            user_api_key_cache.async_set_cache_pipeline(
-                cache_list=list(shared_scalar_updates),
-                ttl=get_management_object_ttl(user_api_key_cache),
-                litellm_parent_otel_span=parent_otel_span,
-            )
-        )
 
 
 def run_ollama_serve():

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -1030,6 +1030,7 @@ def _make_counter_invalidation_job(monkeypatch):
 
     user_api_key_cache = MagicMock()
     user_api_key_cache.async_delete_cache = AsyncMock()
+    user_api_key_cache.async_set_cache = AsyncMock()
 
     fake_module = types.ModuleType("litellm.proxy.proxy_server")
     fake_module.spend_counter_cache = spend_counter_cache
@@ -1088,14 +1089,11 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock
     counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60)
 
 
-def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache(
+def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache(
     reset_budget_job, mock_prisma_client, monkeypatch
 ):
-    """Regression for LIT-4309: resetting the proxy-wide budget aggregate row
-    ("litellm-proxy-budget") must also drop the cached global-spend
-    accumulator ("{admin}:spend") that _global_proxy_budget_check enforces
-    against. Without the invalidation, the cached value survives the DB reset
-    and the global cap keeps blocking requests for the whole next window."""
+    """Resetting the proxy-wide budget aggregate row zeroes the cached
+    global-spend accumulator so the next period starts from a known floor."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
 
     now = datetime.now(timezone.utc)
@@ -1115,13 +1113,15 @@ def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache(
 
     asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
 
-    counter_cache.user_api_key_cache.async_delete_cache.assert_any_call(key="default_user_id:spend")
+    counter_cache.user_api_key_cache.async_set_cache.assert_any_call(
+        key="default_user_id:spend", value=0.0
+    )
 
 
 def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
     reset_budget_job, mock_prisma_client, monkeypatch
 ):
-    """The global-spend accumulator must only be dropped when the proxy
+    """The global-spend accumulator must only be zeroed when the proxy
     budget aggregate row itself resets, not on every user reset."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
 
@@ -1146,6 +1146,10 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
         call.kwargs.get("key") == "default_user_id:spend"
         for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list
     )
+    assert not any(
+        call.kwargs.get("key") == "default_user_id:spend"
+        for call in counter_cache.user_api_key_cache.async_set_cache.call_args_list
+    )
 
 
 def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch):

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -1270,3 +1270,65 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey
     )
 
     assert result is None
+
+
+@pytest.mark.asyncio
+async def test_update_cache_global_proxy_spend_starts_from_zero_after_reset(
+    monkeypatch,
+):
+    """Global proxy spend starts from zero after the cache key is reset."""
+    from litellm.caching.caching import DualCache
+
+    cache = DualCache(default_in_memory_ttl=300)
+    monkeypatch.setattr(ps, "user_api_key_cache", cache)
+
+    await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=50.0)
+    await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
+    await cache.async_set_cache(
+        key="user-lit",
+        value={"user_id": "user-lit", "spend": 0.0},
+    )
+    await ps.update_cache(
+        token=None,
+        user_id="user-lit",
+        end_user_id=None,
+        team_id=None,
+        response_cost=1.05,
+        parent_otel_span=None,
+    )
+
+    final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
+    assert final_value == 1.05
+
+
+@pytest.mark.asyncio
+async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset(
+    monkeypatch,
+):
+    """Concurrent spend updates after a reset accumulate from zero."""
+    from litellm.caching.caching import DualCache
+
+    cache = DualCache(default_in_memory_ttl=300)
+    monkeypatch.setattr(ps, "user_api_key_cache", cache)
+
+    await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=100.0)
+    await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
+    await cache.async_set_cache(
+        key="user-concurrent",
+        value={"user_id": "user-concurrent", "spend": 0.0},
+    )
+
+    async def _update(cost: float) -> None:
+        await ps.update_cache(
+            token=None,
+            user_id="user-concurrent",
+            end_user_id=None,
+            team_id=None,
+            response_cost=cost,
+            parent_otel_span=None,
+        )
+
+    await asyncio.gather(*[_update(0.01) for _ in range(50)])
+
+    final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
+    assert final_value == pytest.approx(0.5)

tests/test_litellm/proxy/test_proxy_server.py

@@ -5278,19 +5278,17 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared():
     ``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB
     reloads, so keeping it pod-local would let traffic spread across replicas
     exceed the proxy budget by roughly a factor of the replica count within a
-    cache TTL. Sharing this scalar is safe because it carries no limits or
-    permissions, so it cannot resurrect an invalidated auth blob.
+    cache TTL. The scalar is updated via an atomic cache increment, which
+    remains shared across pods and avoids resurrecting a stale pre-reset value.
     """
     from litellm.caching.caching import DualCache
 
     admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name
-    global_key = "{}:spend".format(admin_name)
+    global_key = f"{admin_name}:spend"
 
     async def fake_get(key, **kwargs):
         if key == "user-lit":
             return {"user_id": "user-lit", "spend": 1.0}
-        if key == global_key:
-            return 10.0
         return None
 
     original_cache = litellm.proxy.proxy_server.user_api_key_cache
@@ -5299,27 +5297,39 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared():
     try:
         with patch.object(cache, "async_get_cache", new=AsyncMock(side_effect=fake_get)):
             with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache:
-                await litellm.proxy.proxy_server.update_cache(
-                    token=None,
-                    user_id="user-lit",
-                    end_user_id=None,
-                    team_id=None,
-                    response_cost=5.0,
-                    parent_otel_span=None,
-                )
+                with patch.object(cache, "async_increment_cache", new=AsyncMock()) as mock_increment:
+                    await litellm.proxy.proxy_server.update_cache(
+                        token=None,
+                        user_id="user-lit",
+                        end_user_id=None,
+                        team_id=None,
+                        response_cost=5.0,
+                        parent_otel_span=None,
+                    )
 
-                pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-                if pending:
-                    await asyncio.wait(pending, timeout=5)
+                    pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+                    if pending:
+                        await asyncio.wait(pending, timeout=5)
 
-                calls = mock_set_cache.await_args_list
-                local_keys = [k for c in calls if c.kwargs.get("local_only") is True for k, _ in c.kwargs["cache_list"]]
-                shared_keys = [
-                    k for c in calls if c.kwargs.get("local_only") is not True for k, _ in c.kwargs["cache_list"]
-                ]
-                assert "user-lit" in local_keys
-                assert global_key not in local_keys
-                assert shared_keys == [global_key]
+                    calls = mock_set_cache.await_args_list
+                    local_keys = [k for c in calls if c.kwargs.get("local_only") is True for k, _ in c.kwargs["cache_list"]]
+                    shared_keys = [
+                        k for c in calls if c.kwargs.get("local_only") is not True for k, _ in c.kwargs["cache_list"]
+                    ]
+                    assert "user-lit" in local_keys
+                    assert global_key not in local_keys
+                    assert shared_keys == []
+
+                    increment_calls = [
+                        c
+                        for c in mock_increment.await_args_list
+                        if (c.args[0] if c.args else c.kwargs.get("key")) == global_key
+                    ]
+                    assert len(increment_calls) == 1
+                    call = increment_calls[0]
+                    value = call.args[1] if len(call.args) >= 2 else call.kwargs.get("value")
+                    assert value == 5.0
+                    assert call.kwargs.get("refresh_ttl") is True
     finally:
         setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)