Merge 411c727595 into d23bec84c4

litellm/proxy/db/spend_counter_reseed.py

@@ -218,7 +218,15 @@ class SpendCounterReseed:
                         value=current_value,
                     )
                 else:
-                    await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True)
+                    # Re-check after the awaited DB read: another task (e.g. a
+                    # reservation reconcile calling reseed_spend_counter_from_db)
+                    # may have seeded the counter during the await. Incrementing
+                    # then would double it to 2x db_spend. get+set here has no
+                    # awaits, so it is atomic within the event loop.
+                    cached_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+                    if cached_val is not None:
+                        return float(cached_val)
+                    spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
             except Exception:
                 verbose_proxy_logger.exception(
                     "SpendCounterReseed.coalesced: failed to warm counter %s",

litellm/proxy/proxy_server.py

@@ -3103,7 +3103,36 @@ async def _ensure_spend_counter_initialized(
             # 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)
+                await _seed_spend_counter_if_absent(counter_key=counter_key, base_spend=base_spend)
+
+
+async def _seed_spend_counter_if_absent(counter_key: str, base_spend: float) -> None:
+    redis_value: Final = await _seed_redis_spend_counter_nx(counter_key=counter_key, base_spend=base_spend)
+    if redis_value is not None:
+        spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=redis_value)
+        return
+    if spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None:
+        spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=base_spend)
+
+
+async def _seed_redis_spend_counter_nx(counter_key: str, base_spend: float) -> float | None:
+    if spend_counter_cache.redis_cache is None:
+        return None
+    try:
+        seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
+            key=counter_key,
+            value=base_spend,
+            nx=True,
+        )
+        cached: Final = base_spend if seeded else await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
+    except Exception:  # noqa: BLE001  # any Redis failure falls back to in-memory seeding
+        verbose_proxy_logger.debug(
+            "Unable to seed Redis spend counter %s, falling back to in-memory",
+            counter_key,
+            exc_info=True,
+        )
+        return None
+    return float(cached) if cached is not None else base_spend
 
 
 async def _get_source_cache_base_spend(

tests/test_litellm/proxy/db/test_spend_counter_reseed.py

@@ -7,8 +7,10 @@ allowed to run: only when the row is missing or belongs to an older window.
 
 from __future__ import annotations
 
+import asyncio
 from datetime import datetime, timedelta, timezone
 from types import SimpleNamespace
+from typing import Final
 
 import pytest
 
@@ -16,6 +18,8 @@ from litellm.caching.dual_cache import DualCache
 from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
 
 WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
+DB_SPEND: Final = 2777.16
+COUNTER_KEY: Final = "spend:org:test-org"
 
 
 class _FakeWindowSpendTable:
@@ -248,3 +252,80 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
     assert result == 4.5
     assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5
     assert prisma.db.litellm_spendlogs.call_count == 0
+
+
+class _OrgRow:
+    spend: Final = DB_SPEND
+
+
+class _FakeOrgTable:
+    def __init__(self, read_started: asyncio.Event, read_release: asyncio.Event):
+        self.read_started: Final = read_started
+        self.read_release: Final = read_release
+
+    async def find_unique(self, where: dict[str, str]) -> _OrgRow:
+        self.read_started.set()
+        await self.read_release.wait()
+        return _OrgRow()
+
+
+class _FakeDB:
+    def __init__(self, table: _FakeOrgTable):
+        self.litellm_organizationtable: Final = table
+
+
+class _FakePrisma:
+    def __init__(self, table: _FakeOrgTable):
+        self.db: Final = _FakeDB(table)
+
+
+def _fake_prisma(read_started: asyncio.Event, read_release: asyncio.Event) -> _FakePrisma:
+    return _FakePrisma(_FakeOrgTable(read_started=read_started, read_release=read_release))
+
+
+@pytest.mark.asyncio
+async def test_coalesced_no_redis_does_not_double_when_seeded_during_db_read():
+    """
+    Regression test for the cold-seed race (LIT-5516): with no Redis, if a
+    concurrent task (e.g. reservation reconcile) seeds the in-memory counter
+    while coalesced() is awaiting the DB read, the blind increment doubled
+    the counter to 2x db_spend and falsely tripped budget enforcement.
+    """
+    cache: Final = DualCache(default_in_memory_ttl=60)
+    db_read_started: Final = asyncio.Event()
+    db_read_release: Final = asyncio.Event()
+
+    async def seed_during_db_read():
+        await db_read_started.wait()
+        cache.in_memory_cache.set_cache(key=COUNTER_KEY, value=DB_SPEND)
+        db_read_release.set()
+
+    result, _ = await asyncio.gather(
+        SpendCounterReseed.coalesced(
+            prisma_client=_fake_prisma(db_read_started, db_read_release),
+            spend_counter_cache=cache,
+            counter_key=COUNTER_KEY,
+        ),
+        seed_during_db_read(),
+    )
+
+    assert result == DB_SPEND
+    final_value: Final = cache.in_memory_cache.get_cache(key=COUNTER_KEY)
+    assert float(final_value) == DB_SPEND
+
+
+@pytest.mark.asyncio
+async def test_coalesced_no_redis_seeds_cold_counter():
+    cache: Final = DualCache(default_in_memory_ttl=60)
+    db_read_started: Final = asyncio.Event()
+    db_read_release: Final = asyncio.Event()
+    db_read_release.set()
+
+    result: Final = await SpendCounterReseed.coalesced(
+        prisma_client=_fake_prisma(db_read_started, db_read_release),
+        spend_counter_cache=cache,
+        counter_key=COUNTER_KEY,
+    )
+
+    assert result == DB_SPEND
+    assert float(cache.in_memory_cache.get_cache(key=COUNTER_KEY)) == DB_SPEND

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -1062,12 +1062,16 @@ async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache(
 
     observed = {
         "source_cache_called": fake_user_cache.async_get_cache.called,
+        "seed_set_nx_called": fake_cache.redis_cache.async_set_cache.called,
         "seed_increment_called": fake_cache.redis_cache.async_increment.called,
+        "in_memory_seeded": fake_cache.in_memory_cache.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_nx_called": True,
+        "seed_increment_called": False,
+        "in_memory_seeded": True,
         "warm_check_done": True,
     }