fix(proxy): preserve local spend adjustments during reseeding

litellm/proxy/db/spend_counter_reseed.py

@@ -104,6 +104,15 @@ class SpendCounterReseed:
                 SpendCounterReseed._locks.popitem(last=False)
             return lock
 
+    @staticmethod
+    async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None:
+        """Apply local deltas after an in-flight reseed establishes the spend balance."""
+        lock: Final = await SpendCounterReseed._get_lock(counter_key)
+        async with lock:
+            return await spend_counter_cache.async_increment_cache(
+                key=counter_key, value=increment, local_only=True, refresh_ttl=True
+            )
+
     @staticmethod
     async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
         """
@@ -245,16 +254,10 @@ class SpendCounterReseed:
                         value=current_value,
                     )
                 else:
-                    # Repair/reservations can populate the counter during the DB read.
-                    # Seed a floor without adding the database balance again.
-                    # No await between read/compare/write: atomic within this worker.
-                    cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
-                    current_value = float(db_spend)
-                    if cached is not None:
-                        current_value = max(current_value, float(cached))
-                    spend_counter_cache.in_memory_cache.set_cache(
-                        key=counter_key, value=current_value
-                    )
+                    cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+                    seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend
+                    spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
+                    return seeded_spend
             except Exception:
                 verbose_proxy_logger.exception(
                     "SpendCounterReseed.coalesced: failed to warm counter %s",
@@ -447,16 +450,12 @@ class SpendCounterReseed:
                         value=current_value,
                     )
                 else:
-                    # Repair/reservations can populate the counter during the DB read.
-                    # Seed a floor without adding the database balance again.
-                    # No await between read/compare/write: atomic within this worker.
-                    cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
-                    current_value = float(window_spend)
-                    if cached is not None:
-                        current_value = max(current_value, float(cached))
-                    spend_counter_cache.in_memory_cache.set_cache(
-                        key=counter_key, value=current_value
+                    cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+                    seeded_spend: Final = (
+                        max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend
                     )
+                    spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
+                    return seeded_spend
             except Exception:
                 verbose_proxy_logger.exception(
                     "SpendCounterReseed.coalesced_window: failed to warm counter %s",

litellm/proxy/proxy_server.py

@@ -3330,10 +3330,8 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float):
         )
         return current_value
 
-    return await spend_counter_cache.async_increment_cache(
-        key=counter_key,
-        value=increment,
-        refresh_ttl=True,
+    return await SpendCounterReseed.increment_in_memory(
+        spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment
     )
 
 
@@ -3356,10 +3354,8 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme
     redis_cache: Final = spend_counter_cache.redis_cache
     if redis_cache is None:
         for item in pending:
-            await spend_counter_cache.async_increment_cache(
-                key=item.counter_key,
-                value=item.increment,
-                refresh_ttl=True,
+            await SpendCounterReseed.increment_in_memory(
+                spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment
             )
         return
     ttl: Final = redis_cache.get_ttl()

tests/test_litellm/proxy/db/test_spend_counter_reseed.py

@@ -8,6 +8,7 @@ allowed to run: only when the row is missing or belongs to an older window.
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
 from datetime import datetime, timedelta, timezone
 from types import SimpleNamespace
 from typing import Final
@@ -79,6 +80,39 @@ def _row(window_start: datetime, spend: float) -> SimpleNamespace:
     return SimpleNamespace(window_start=window_start, spend=spend)
 
 
+class _PausedSpendTable:
+    def __init__(self, spend: float) -> None:
+        self.spend: Final = spend
+        self.read_started: Final = asyncio.Event()
+        self.resume_read: Final = asyncio.Event()
+
+    async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace:
+        self.read_started.set()
+        await self.resume_read.wait()
+        return _row(WINDOW_START, self.spend)
+
+
+async def _reseed_with_paused_table(
+    table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool
+) -> float | None:
+    prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table))
+    if window:
+        return await SpendCounterReseed.coalesced_window(
+            prisma_client=prisma,
+            spend_counter_cache=cache,
+            counter_key=counter_key,
+            entity_type="Team",
+            entity_id="team-1",
+            window_duration="1d",
+            window_start=WINDOW_START,
+        )
+    return await SpendCounterReseed.coalesced(
+        prisma_client=prisma,
+        spend_counter_cache=cache,
+        counter_key=counter_key,
+    )
+
+
 @pytest.mark.asyncio
 async def test_window_from_table_reads_row_by_primary_key():
     """The lookup must use the table's own entity_type values ("key"), not the
@@ -275,49 +309,59 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
 @pytest.mark.parametrize("window", [False, True], ids=["primary", "window"])
 @pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0])
 async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache(
-    monkeypatch: pytest.MonkeyPatch,
     window: bool,
     concurrent_spend: float,
-):
-    """A repair/reservation write may populate the counter while DB read runs.
+) -> None:
+    cache: Final = DualCache(in_memory_cache=InMemoryCache())
+    counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1"
+    db_spend: Final = 989.01459411
+    table: Final = _PausedSpendTable(db_spend)
+    reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window))
 
-    Reseeding must establish the larger value, not increment the concurrent
-    value by the same authoritative spend a second time.
-    """
-    cache = DualCache(in_memory_cache=InMemoryCache())
-    counter_key = (
-        "spend:team:team-1:window:1d" if window else "spend:user:user-1"
-    )
-    db_spend = 989.01459411
+    await asyncio.wait_for(table.read_started.wait(), timeout=5)
+    cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend)
+    table.resume_read.set()
+    result: Final = await asyncio.wait_for(reseed_task, timeout=5)
 
-    async def read_db(*args, **kwargs):
-        cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend)
-        return db_spend
-
-    if window:
-        monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db))
-        result = await SpendCounterReseed.coalesced_window(
-            prisma_client=None,
-            spend_counter_cache=cache,
-            counter_key=counter_key,
-            entity_type="Team",
-            entity_id="team-1",
-            window_duration="1d",
-            window_start=WINDOW_START,
-        )
-    else:
-        monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db))
-        result = await SpendCounterReseed.coalesced(
-            prisma_client=None,
-            spend_counter_cache=cache,
-            counter_key=counter_key,
-        )
-
-    expected = max(db_spend, concurrent_spend)
+    expected: Final = max(db_spend, concurrent_spend)
     assert cache.in_memory_cache.get_cache(key=counter_key) == expected
     assert result == expected
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"])
+@pytest.mark.parametrize("batch", [False, True], ids=["single_increment", "batch_increment"])
+@pytest.mark.parametrize("increment", [5.0, -5.0], ids=["charge", "refund"])
+async def test_cold_reseed_preserves_concurrent_local_increment(
+    monkeypatch: pytest.MonkeyPatch, window: bool, batch: bool, increment: float
+) -> None:
+    from litellm.proxy import proxy_server
+
+    cache: Final = DualCache(in_memory_cache=InMemoryCache())
+    counter_key: Final = (
+        f"spend:team:concurrent-{batch}-{increment}:window:1d"
+        if window
+        else f"spend:user:concurrent-{batch}-{increment}"
+    )
+    table: Final = _PausedSpendTable(100.0)
+    monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
+    reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window))
+    await asyncio.wait_for(table.read_started.wait(), timeout=5)
+
+    increment_task: Final = asyncio.create_task(
+        proxy_server._apply_spend_counter_increments(
+            pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),)
+        )
+        if batch
+        else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment)
+    )
+    await asyncio.sleep(0)
+    table.resume_read.set()
+    await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5)
+
+    assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment
+
+
 @pytest.mark.asyncio
 async def test_end_user_from_db_reads_the_end_user_row_by_user_id():
     prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0))
@@ -352,8 +396,7 @@ async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the
 @pytest.mark.asyncio
 async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error():
     assert (
-        await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42")
-        is None
+        await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") is None
     )
     assert (
         await SpendCounterReseed.end_user_from_db(