diff --git a/litellm/constants.py b/litellm/constants.py index 012599ab6a..32e016cd7a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1429,6 +1429,9 @@ SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) +SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int( + os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) +) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int( os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60) ) # 1 minute diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py new file mode 100644 index 0000000000..bf60a087c6 --- /dev/null +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -0,0 +1,159 @@ +""" +Coalesced reseed of spend counters from the authoritative DB. + +When a Redis spend counter expires (or is missing on a fresh pod), enforcement +must read the current spend from somewhere. The in-process management cache +(`user_api_key_cache.team_membership.spend`, etc.) is per-pod and lags DB +writes from other pods, so trusting it allows budget bypass in multi-pod +deployments. This module reseeds from the authoritative DB instead. + +A per-counter singleflight lock collapses concurrent reseeds on the same pod +to one DB query per cold-cache window. The lock dict is bounded LRU to cap +memory in long-lived deployments. +""" + +import asyncio +from collections import OrderedDict +from typing import TYPE_CHECKING, ClassVar, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.proxy.utils import PrismaClient + + +class SpendCounterReseed: + """ + Reseeds spend counters from the authoritative DB and warms the cache, + coalesced via per-counter singleflight locks. + + Counter key prefixes map to DB tables: + spend:key:{token} -> LiteLLM_VerificationToken.spend + spend:team:{team_id} -> LiteLLM_TeamTable.spend + spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend + spend:user:{user_id} -> LiteLLM_UserTable.spend + spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + """ + + _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() + _registry_lock: ClassVar[Optional[asyncio.Lock]] = None + + @staticmethod + async def _get_lock(counter_key: str) -> asyncio.Lock: + if SpendCounterReseed._registry_lock is None: + SpendCounterReseed._registry_lock = asyncio.Lock() + async with SpendCounterReseed._registry_lock: + lock = SpendCounterReseed._locks.get(counter_key) + if lock is not None: + SpendCounterReseed._locks.move_to_end(counter_key) + return lock + lock = asyncio.Lock() + SpendCounterReseed._locks[counter_key] = lock + if len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: + SpendCounterReseed._locks.popitem(last=False) + return lock + + @staticmethod + async def from_db( + prisma_client: Optional["PrismaClient"], counter_key: str + ) -> Optional[float]: + """ + Read the authoritative spend for a counter from the DB. + + Returns the spend value (including 0.0) when the DB is reachable + and the row exists. Returns None when prisma is unavailable, the + row is missing, the key format is unrecognized, or the query + raises. Callers use None to fall back to a caller-supplied source. + """ + if prisma_client is None: + return None + # Per-window counters share prefixes with primary counters but + # don't correspond to a DB row. + if ":window:" in counter_key: + return None + try: + if counter_key.startswith("spend:key:"): + token = counter_key[len("spend:key:") :] + row = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": token} + ) + elif counter_key.startswith("spend:team_member:"): + suffix = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + row = await prisma_client.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + elif counter_key.startswith("spend:team:"): + team_id = counter_key[len("spend:team:") :] + row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + elif counter_key.startswith("spend:user:"): + user_id = counter_key[len("spend:user:") :] + row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) + elif counter_key.startswith("spend:org:"): + org_id = counter_key[len("spend:org:") :] + row = await prisma_client.db.litellm_organizationtable.find_unique( + where={"organization_id": org_id} + ) + else: + return None + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.from_db: failed for %s", counter_key + ) + return None + if row is None: + return None + return float(getattr(row, "spend", 0.0) or 0.0) + + @staticmethod + async def coalesced( + prisma_client: Optional["PrismaClient"], + spend_counter_cache: "DualCache", + counter_key: str, + ) -> Optional[float]: + """ + Reseed a cold spend counter from the DB and warm the cache, + coalesced via a per-counter lock so concurrent callers (read path + + write path) collapse to one DB query per cold-cache window. + + Returns the spend value (including 0.0 from a fresh budget reset) + when the DB read succeeds, or None when the DB is unavailable. + """ + lock = await SpendCounterReseed._get_lock(counter_key) + async with lock: + # Re-check after acquiring the lock - another waiter may have warmed it. + if spend_counter_cache.redis_cache is not None: + try: + val = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + if val is not None: + return float(val) + except Exception: + pass + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) + + db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key) + if db_spend is None: + return None + # Warm even when 0 so subsequent reads hit cache, not DB. + try: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=db_spend + ) + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.coalesced: failed to warm counter %s", + counter_key, + ) + return db_spend diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ca6895c2c..3982d096e6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -323,6 +323,7 @@ from litellm.proxy.container_endpoints.endpoints import router as container_rout from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config @@ -1775,7 +1776,8 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: Fallback chain: 1. Redis counter (cross-pod, authoritative) 2. In-memory counter (single-instance or Redis failure) - 3. Cached object's .spend from DB (cold start, no counter yet) + 3. Reseed from authoritative DB spend (counter expired, cross-pod stale) + 4. Caller-supplied fallback (DB unavailable, cold start) """ # 1. Try Redis first (cross-pod authoritative) if spend_counter_cache.redis_cache is not None: @@ -1795,7 +1797,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: if val is not None: return float(val) - # 3. Final fallback: cached object's spend from DB + # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. + db_spend = await SpendCounterReseed.coalesced( + prisma_client=prisma_client, + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + ) + if db_spend is not None: + return db_spend + + # 4. Caller-supplied fallback (DB unavailable). return fallback_spend @@ -1906,69 +1917,6 @@ async def increment_spend_counters( ) -async def _reseed_spend_from_db(counter_key: str) -> float: - """ - Read the authoritative spend for a missing counter from the DB. The - counter_key prefix encodes the table to query: - - spend:key:{token} -> LiteLLM_VerificationToken.spend - spend:team:{team_id} -> LiteLLM_TeamTable.spend - spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend - spend:user:{user_id} -> LiteLLM_UserTable.spend - spend:org:{org_id} -> LiteLLM_OrganizationTable.spend - - Returns 0.0 if prisma is unavailable, the row is missing, or the - key format is unrecognized. On failure, logs and returns 0.0 rather - than raising so the caller can still record the current increment. - """ - if prisma_client is None: - return 0.0 - # Per-window counters (spend:*:window:{duration}) share prefixes with - # primary counters but don't correspond to a DB row; their ambiguity - # would otherwise be silently parsed as a regular counter and miss. - if ":window:" in counter_key: - return 0.0 - try: - if counter_key.startswith("spend:key:"): - token = counter_key[len("spend:key:") :] - row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": token} - ) - elif counter_key.startswith("spend:team_member:"): - suffix = counter_key[len("spend:team_member:") :] - if ":" not in suffix: - return 0.0 - user_id, team_id = suffix.rsplit(":", 1) - row = await prisma_client.db.litellm_teammembership.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} - ) - elif counter_key.startswith("spend:org:"): - org_id = counter_key[len("spend:org:") :] - row = await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": org_id} - ) - else: - return 0.0 - except Exception: - verbose_proxy_logger.exception( - "Failed to reseed spend counter %s from DB", counter_key - ) - return 0.0 - if row is None: - return 0.0 - return float(getattr(row, "spend", 0.0) or 0.0) - - async def _init_and_increment_spend_counter( counter_key: str, source_cache_key: str, @@ -1980,9 +1928,9 @@ async def _init_and_increment_spend_counter( On first access per pod: 1. Check spend_counter_cache (in-memory -> Redis via DualCache) - 2. If not found, reseed from the DB (`_reseed_spend_from_db`). 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. + 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 @@ -1992,20 +1940,25 @@ async def _init_and_increment_spend_counter( """ current = await spend_counter_cache.async_get_cache(key=counter_key) if current is None: - base_spend = await _reseed_spend_from_db(counter_key) - if prisma_client is None: - # Best-effort fallback when prisma is unavailable (tests or - # early-startup paths). May be stale but avoids resetting to 0. + # Shares the per-counter lock with get_current_spend. + db_spend = await SpendCounterReseed.coalesced( + prisma_client=prisma_client, + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + ) + if db_spend is None: + # DB unavailable - fall back to in-process cache (may be stale). source = await user_api_key_cache.async_get_cache(key=source_cache_key) + base_spend: float = 0.0 if source is not None: if isinstance(source, dict): base_spend = source.get("spend", 0.0) or 0.0 else: base_spend = getattr(source, "spend", 0.0) or 0.0 - if base_spend > 0: - await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend - ) + if base_spend > 0: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=base_spend + ) await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd1abbb38..3349a138ee 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5034,8 +5034,7 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( async def test_reseed_spend_from_db_user_and_org_prefixes(): """User and org counters must reseed from their own DB tables, not fall through to 0.0 like the other counters do today.""" - import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _reseed_spend_from_db + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed user_row = MagicMock() user_row.spend = 17.0 @@ -5048,20 +5047,15 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): return_value=org_row ) - orig_prisma = ps.prisma_client - ps.prisma_client = fake_prisma - try: - assert await _reseed_spend_from_db("spend:user:alice") == 17.0 - fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with( - where={"user_id": "alice"} - ) + assert await SpendCounterReseed.from_db(fake_prisma, "spend:user:alice") == 17.0 + fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with( + where={"user_id": "alice"} + ) - assert await _reseed_spend_from_db("spend:org:acme") == 305.0 - fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( - where={"organization_id": "acme"} - ) - finally: - ps.prisma_client = orig_prisma + assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 + fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( + where={"organization_id": "acme"} + ) @pytest.mark.asyncio @@ -5069,19 +5063,391 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): """Window counters (spend:*:window:{duration}) share prefixes with primary counters but don't correspond to a DB row. The guard must short-circuit without querying the DB.""" - import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import _reseed_spend_from_db + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed fake_prisma = MagicMock() fake_prisma.db.litellm_verificationtoken.find_unique = AsyncMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock() - orig_prisma = ps.prisma_client + assert ( + await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h") + is None + ) + assert ( + await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d") + is None + ) + fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited() + fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_current_spend_reseeds_from_db_when_counter_missing(): + """ + When both the Redis and in-memory counters are missing, the enforcement + read path must reseed from the authoritative DB, not fall back to the + caller-supplied stale value. Otherwise, every Redis TTL expiry lets a + request through against a stale in-process `team_membership.spend`. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + recorded_warms: list = [] + + async def record_increment(key, value, ttl=None, **kwargs): + recorded_warms.append({"key": key, "value": value}) + return value + + fake_redis = AsyncMock() + fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_get_cache = AsyncMock(return_value=None) + counter_cache.redis_cache = fake_redis + + # DB has authoritative spend=362.0; caller hands us stale fallback=30.0 + # (the in-process team_membership.spend that hasn't caught up to DB). + db_row = MagicMock() + db_row.spend = 362.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache ps.prisma_client = fake_prisma try: - assert await _reseed_spend_from_db("spend:key:sk-abc:window:1h") == 0.0 - assert await _reseed_spend_from_db("spend:team:team-1:window:1d") == 0.0 - fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited() - fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() + spend = await get_current_spend( + counter_key="spend:team_member:user-1:team-1", + fallback_spend=30.0, + ) + assert spend == 362.0, ( + f"expected DB reseed to return 362.0, got {spend} " + f"(fallback would have returned 30.0 and caused bypass)" + ) + # Counter warmed so subsequent reads are fast + assert ("spend:team_member:user-1:team-1", 362.0) in [ + (w["key"], w["value"]) for w in recorded_warms + ] finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_uses_fallback_when_db_unavailable(): + """ + If prisma is unavailable and both counters are missing, the read path + must degrade to the caller-supplied fallback rather than raising. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + counter_cache.redis_cache = fake_redis + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = None # simulate prisma unavailable + try: + spend = await get_current_spend( + counter_key="spend:team_member:user-1:team-1", + fallback_spend=15.5, + ) + assert spend == 15.5 + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_coalesces_concurrent_reseeds(): + """ + When several concurrent calls hit a cold counter on the same pod, + only one DB query should fire. The rest should wait for the lock, + re-check the warmed counter, and return without hitting the DB. + """ + import asyncio as _asyncio + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-coalesce" + + # Track DB query calls and inject a small delay so the concurrent + # callers actually overlap in the lock-acquire window. + db_call_count = 0 + + async def slow_find_unique(**kwargs): + nonlocal db_call_count + db_call_count += 1 + await _asyncio.sleep(0.05) + row = MagicMock() + row.spend = 100.0 + return row + + fake_redis = AsyncMock() + redis_store: dict = {} + + async def redis_get(key, **_): + return redis_store.get(key) + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=slow_find_unique + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + results = await _asyncio.gather( + *[ + get_current_spend(counter_key=counter_key, fallback_spend=0.0) + for _ in range(5) + ] + ) + assert results == [100.0] * 5, f"all callers should see DB value, got {results}" + assert ( + db_call_count == 1 + ), f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_uses_db_zero_over_stale_fallback(): + """ + When DB returns spend=0 (e.g. just after a budget period reset), the + authoritative DB value must win over a stale non-zero fallback. The + fallback in production is the in-process team_membership.spend, which + can still hold the pre-reset value across pods. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + counter_cache.redis_cache = fake_redis + + db_row = MagicMock() + db_row.spend = 0.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend( + counter_key="spend:team_member:user-1:team-after-reset", + fallback_spend=42.0, + ) + assert ( + spend == 0.0 + ), f"DB authoritative 0 must override stale fallback 42, got {spend}" + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_concurrent_read_and_write_paths_share_one_db_query(): + """ + The read path (`get_current_spend`) and the write path + (`_init_and_increment_spend_counter`) both reseed cold counters from + the DB. They must share the per-counter lock so a concurrent pre-call + enforcement read and post-call increment for the same counter collapse + to one DB query, not two. + """ + import asyncio as _asyncio + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import ( + _init_and_increment_spend_counter, + get_current_spend, + ) + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-cross-path" + + db_call_count = 0 + + async def slow_find_unique(**kwargs): + nonlocal db_call_count + db_call_count += 1 + await _asyncio.sleep(0.05) + row = MagicMock() + row.spend = 50.0 + return row + + redis_store: dict = {} + + async def redis_get(key, **_): + return redis_store.get(key) + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=slow_find_unique + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma, orig_user = ( + ps.spend_counter_cache, + ps.prisma_client, + ps.user_api_key_cache, + ) + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + ps.user_api_key_cache = DualCache() + try: + results = await _asyncio.gather( + get_current_spend(counter_key=counter_key, fallback_spend=0.0), + _init_and_increment_spend_counter( + counter_key=counter_key, + source_cache_key="ignored", + increment=1.5, + ), + get_current_spend(counter_key=counter_key, fallback_spend=0.0), + ) + assert ( + db_call_count == 1 + ), f"expected 1 DB query for concurrent read+write+read, got {db_call_count}" + # Read-path callers see the warmed counter; the write path's + # increment may or may not have landed by then, so accept either + # the seeded value or seeded+increment. + assert results[0] in (50.0, 51.5), f"got {results[0]}" + assert results[2] in (50.0, 51.5), f"got {results[2]}" + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + ps.user_api_key_cache = orig_user + + +@pytest.mark.asyncio +async def test_reseed_locks_dict_is_bounded(): + """ + `SpendCounterReseed._locks` is an LRU bounded at + `SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE` to prevent unbounded growth in + long-lived deployments with high counter-key churn. Inserting more + than the cap evicts the oldest entries. + """ + import litellm.constants as constants + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + orig_locks = SpendCounterReseed._locks.copy() + SpendCounterReseed._locks.clear() + orig_max = constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE + constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = 5 + # The class reads the constant via module-level import, so patch the + # module-level name on the spend_counter_reseed module too. + import litellm.proxy.db.spend_counter_reseed as scr + + orig_module_max = scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE + scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = 5 + try: + for i in range(7): + await SpendCounterReseed._get_lock(f"spend:key:test-key-{i}") + assert ( + len(SpendCounterReseed._locks) == 5 + ), f"got {len(SpendCounterReseed._locks)}" + # Oldest two evicted + assert "spend:key:test-key-0" not in SpendCounterReseed._locks + assert "spend:key:test-key-1" not in SpendCounterReseed._locks + # Most recent retained + assert "spend:key:test-key-6" in SpendCounterReseed._locks + finally: + constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = orig_max + scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = orig_module_max + SpendCounterReseed._locks.clear() + SpendCounterReseed._locks.update(orig_locks) + + +@pytest.mark.asyncio +async def test_reseed_warms_cache_even_on_zero_db_spend(): + """ + When DB returns 0.0 (fresh entity / just after reset), the cache must + still be warmed so subsequent reads hit the cache instead of issuing + another DB query. Skipping the warm causes O(requests) DB load on + zero-spend entities. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-zero-warm" + redis_store: dict = {} + + async def redis_get(key, **_): + return redis_store.get(key) + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + db_call_count = 0 + + async def find_unique(**kwargs): + nonlocal db_call_count + db_call_count += 1 + row = MagicMock() + row.spend = 0.0 + return row + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + side_effect=find_unique + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + # First call: cold cache, hits DB, returns 0. + spend1 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + # Second call: cache should be warmed at 0, no second DB query. + spend2 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend1 == 0.0 and spend2 == 0.0 + assert ( + db_call_count == 1 + ), f"second read should hit warmed cache, got {db_call_count} DB queries" + assert redis_store.get(counter_key) == 0.0, "cache must be warmed at 0" + finally: + ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma