diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 7b3c261036e..0c8d9345870 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5a39b8c610a..db8da1ba910 100644 --- a/litellm/proxy/proxy_server.py +++ b/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( diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 816f9ae72f4..3021470955e 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/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 diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index fb3de990deb..62638eb2510 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/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, }