fix(proxy): seed cold spend counter from source cache with SET NX to prevent spend doubling

litellm/proxy/proxy_server.py

@@ -3837,11 +3837,8 @@ async def _prepare_spend_counter_increment(
     2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
        Falls back to the cached object's `.spend` via user_api_key_cache
        only if prisma is unavailable, since that value can lag the flusher.
-    3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
-       check-then-set race: if two pods cold-start simultaneously, both may see
-       the counter as absent and seed it. Using increment means the worst case
-       is over-counting (conservative, blocks slightly early) rather than
-       under-counting (would allow overspend).
+    3. Seed counter via SET NX under per-counter lock (matching SpendCounterReseed.coalesced)
+       to avoid double-counting base_spend when concurrent requests cold-start the counter.
     4. Increment is returned for the caller to apply via pipeline
     """
     with service_target(SPEND_COUNTERS_TARGET):
@@ -3944,9 +3941,34 @@ async def ensure_spend_counter_initialized(
         )
         if db_spend is None:
             # DB unavailable - fall back to in-process cache (may be stale).
-            base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key)
-            if base_spend > 0:
-                await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend)
+            async with await SpendCounterReseed._get_lock(counter_key):
+                if await _is_spend_counter_cache_warm(counter_key=counter_key):
+                    return
+                base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key)
+                if base_spend > 0:
+                    if spend_counter_cache.redis_cache is not None:
+                        seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
+                            key=counter_key,
+                            value=base_spend,
+                            nx=True,
+                        )
+                        if seeded:
+                            current_value = float(base_spend)
+                        else:
+                            cached: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
+                            current_value = float(cached) if cached is not None else float(base_spend)
+                        spend_counter_cache.in_memory_cache.set_cache(
+                            key=counter_key,
+                            value=current_value,
+                        )
+                        record_spend_counter_value(counter_key, current_value)
+                    else:
+                        cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+                        seeded_spend: Final = (
+                            max(base_spend, float(cached_spend)) if cached_spend is not None else base_spend
+                        )
+                        spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
+                        record_spend_counter_value(counter_key, seeded_spend)
 
 
 _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized

tests/unit/proxy/proxy_server/test_spend_counters.py

@@ -1326,14 +1326,57 @@ async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache(
 
     observed = {
         "source_cache_called": fake_user_cache.async_get_cache.called,
-        "seed_increment_called": fake_cache.redis_cache.async_increment.called,
+        "seed_set_cache_called": fake_cache.redis_cache.async_set_cache.called,
         "warm_check_done": fake_cache.redis_cache.async_get_cache.called,
     }
     assert normalize(observed) == {
         "source_cache_called": True,
-        "seed_increment_called": True,
+        "seed_set_cache_called": True,
         "warm_check_done": True,
     }
+    fake_cache.redis_cache.async_set_cache.assert_awaited_once_with(
+        key="spend:user:u",
+        value=7.0,
+        nx=True,
+    )
+    assert not fake_cache.redis_cache.async_increment.called
+
+
+@pytest.mark.asyncio
+async def test_ensure_spend_counter_initialized_concurrent_cold_seeds_do_not_double_spend(
+    monkeypatch,
+):
+    redis_store: dict[str, float] = {}
+
+    async def _fake_async_get_cache(key, **kwargs):
+        return redis_store.get(key)
+
+    async def _fake_async_set_cache(key, value, nx=False, **kwargs):
+        if nx and key in redis_store:
+            return False
+        redis_store[key] = float(value)
+        return True
+
+    fake_cache = _make_spend_counter_cache()
+    fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_fake_async_get_cache)
+    fake_cache.redis_cache.async_set_cache = AsyncMock(side_effect=_fake_async_set_cache)
+
+    fake_user_cache = _make_user_api_key_cache(get_value={"spend": 6.0})
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+    monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
+    monkeypatch.setattr(ps, "prisma_client", None)
+    monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
+
+    counter_key = "spend:end_user:customer-1"
+    await asyncio.gather(
+        ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"),
+        ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"),
+        ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"),
+    )
+
+    assert redis_store[counter_key] == 6.0
+    assert fake_cache.redis_cache.async_increment.call_count == 0
+
 
 
 # ---------------------------------------------------------------------------