diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 34ae3638a5..6115a444ce 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -392,6 +392,7 @@ class DualCache(BaseCache): value: float, parent_otel_span: Optional[Span] = None, local_only: bool = False, + refresh_ttl: bool = False, **kwargs, ) -> Optional[float]: """ @@ -399,6 +400,9 @@ class DualCache(BaseCache): Value - float - the value you want to increment by + Refresh_ttl - bool - if True, resets the Redis TTL on every write. + Default False preserves window-style semantics. + Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ @@ -415,6 +419,7 @@ class DualCache(BaseCache): value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), + refresh_ttl=refresh_ttl, ) return result diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f52..deee4f6ea4 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -824,6 +824,7 @@ class RedisCache(BaseCache): value: float, ttl: Optional[int] = None, parent_otel_span: Optional[Span] = None, + refresh_ttl: bool = False, ) -> float: from redis.asyncio import Redis @@ -834,11 +835,12 @@ class RedisCache(BaseCache): try: result = await _redis_client.incrbyfloat(name=key, amount=value) if _used_ttl is not None: - # check if key already has ttl, if not -> set ttl - current_ttl = await _redis_client.ttl(key) - if current_ttl == -1: - # Key has no expiration + if refresh_ttl: await _redis_client.expire(key, _used_ttl) + else: + current_ttl = await _redis_client.ttl(key) + if current_ttl == -1: + await _redis_client.expire(key, _used_ttl) ## LOGGING ## end_time = time.time() diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index e486336cec..0928ce914d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -52,6 +52,37 @@ class ResetBudgetJob: ### RESET MULTI-WINDOW BUDGETS ### await self.reset_budget_windows() + @staticmethod + async def _invalidate_spend_counter(counter_key: str) -> None: + """Zero a spend counter so a DB-row reset takes effect immediately. + + Call AFTER the DB write commits. Clearing Redis before the DB + commit opens a window where get_current_spend reads 0 from Redis + while the DB still holds the pre-reset value, allowing bypass. + """ + try: + from litellm.proxy.proxy_server import spend_counter_cache + + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s in Redis: %s. " + "Budget may be over-enforced until counter expires.", + counter_key, + redis_err, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s: %s", counter_key, e + ) + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -64,46 +95,30 @@ class ResetBudgetJob: if budget.budget_id is not None ] - # Reset spend counters for affected team members. - # Reset Redis directly so a transient failure doesn't leave stale - # counters that get_current_spend would read as authoritative. try: - from litellm.proxy.proxy_server import spend_counter_cache - memberships = await self.prisma_client.db.litellm_teammembership.find_many( where={"budget_id": {"in": budget_ids}} ) - for m in memberships: - counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" - # Always reset in-memory - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset team member spend counter in Redis %s: %s. " - "Budget may be over-enforced until counter expires.", - counter_key, - redis_err, - ) except Exception as e: + memberships = [] verbose_proxy_logger.warning( - "Failed to reset team member spend counters: %s", e + "Failed to fetch team memberships for counter invalidation: %s", e ) - return await self.prisma_client.db.litellm_teammembership.update_many( + update_result = await self.prisma_client.db.litellm_teammembership.update_many( where={"budget_id": {"in": budget_ids}}, data={ "spend": 0, }, ) + for m in memberships: + await self._invalidate_spend_counter( + f"spend:team_member:{m.user_id}:{m.team_id}" + ) + + return update_result + async def reset_budget_for_keys_linked_to_budgets( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -126,17 +141,36 @@ class ResetBudgetJob: if not budget_ids: return - return await self.prisma_client.db.litellm_verificationtoken.update_many( - where={ - "budget_id": {"in": budget_ids}, - "budget_duration": None, # only keys without their own reset schedule - "spend": {"gt": 0}, # only reset keys that have accumulated spend - }, - data={ - "spend": 0, - }, + where_clause: dict = { + "budget_id": {"in": budget_ids}, + "budget_duration": None, # only keys without their own reset schedule + "spend": {"gt": 0}, # only reset keys that have accumulated spend + } + + try: + keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where=where_clause + ) + except Exception as e: + keys = [] + verbose_proxy_logger.warning( + "Failed to fetch keys for counter invalidation: %s", e + ) + + update_result = ( + await self.prisma_client.db.litellm_verificationtoken.update_many( + where=where_clause, + data={ + "spend": 0, + }, + ) ) + for k in keys: + await self._invalidate_spend_counter(f"spend:key:{k.token}") + + return update_result + async def reset_budget_for_litellm_budget_table(self): """ Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired @@ -365,6 +399,10 @@ class ResetBudgetJob: data_list=updated_keys, table_name="key", ) + for k in updated_keys: + token = getattr(k, "token", None) + if token: + await self._invalidate_spend_counter(f"spend:key:{token}") end_time = time.time() if len(failed_keys) > 0: # If any keys failed to reset @@ -450,6 +488,12 @@ class ResetBudgetJob: data_list=updated_users, table_name="user", ) + for u in updated_users: + user_id = getattr(u, "user_id", None) + if user_id: + await self._invalidate_spend_counter( + f"spend:user:{user_id}" + ) end_time = time.time() if len(failed_users) > 0: # If any users failed to reset @@ -541,6 +585,12 @@ class ResetBudgetJob: data_list=updated_teams, table_name="team", ) + for t in updated_teams: + team_id = getattr(t, "team_id", None) + if team_id: + await self._invalidate_spend_counter( + f"spend:team:{team_id}" + ) end_time = time.time() if len(failed_teams) > 0: # If any teams failed to reset diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index bf60a087c6..a979471dc8 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -129,7 +129,9 @@ class SpendCounterReseed: """ lock = await SpendCounterReseed._get_lock(counter_key) async with lock: - # Re-check after acquiring the lock - another waiter may have warmed it. + # Re-check after acquiring the lock. Skip in-memory on a clean + # Redis miss - in-memory is per-pod-stale. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache( @@ -137,11 +139,13 @@ class SpendCounterReseed: ) if val is not None: return float(val) + redis_clean_miss = True except Exception: pass - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + if not redis_clean_miss: + 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: @@ -149,7 +153,7 @@ class SpendCounterReseed: # 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 + key=counter_key, value=db_spend, refresh_ttl=True ) except Exception: verbose_proxy_logger.exception( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6f5ab1afb6..6cba6a3e96 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1798,12 +1798,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: 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) + # 1. Redis first (cross-pod authoritative). On clean miss, skip + # in-memory: per-pod in-memory only has this pod's writes, so it + # would mask cross-pod increments. + redis_clean_miss = False 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) + redis_clean_miss = True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -1811,10 +1815,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: e, ) - # 2. Fall back to in-memory counter (single-instance or Redis failure) - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + # 2. In-memory only when Redis is unreachable. + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend = await SpendCounterReseed.coalesced( @@ -1976,10 +1981,12 @@ async def _init_and_increment_spend_counter( 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 + key=counter_key, value=base_spend, refresh_ttl=True ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, refresh_ttl=True + ) async def update_cache( # noqa: PLR0915 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index b39eb42821..78192400fb 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping ) +@pytest.mark.asyncio +async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl( + monkeypatch, redis_no_ping +): + """With refresh_ttl=True, every increment should call expire() to bump + the TTL, even when the key already has a TTL (counter-style use).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment( + key="spend:team_member:u:t", value=0.05, refresh_ttl=True + ) + + mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60) + + +@pytest.mark.asyncio +async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( + monkeypatch, redis_no_ping +): + """Default (refresh_ttl=False) preserves window-style semantics: TTL is + set only on first creation, never refreshed (used by rate-limit windows).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment(key="rate_limit:window", value=1) + + mock_redis_instance.expire.assert_not_awaited() + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 379ccf4d9a..5c86f9057a 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch): asyncio.run(job.reset_budget_windows()) # must not raise prisma_client.db.litellm_teamtable.update.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Counter invalidation on budget reset +# --------------------------------------------------------------------------- + + +def _make_counter_invalidation_job(monkeypatch): + """Stub spend_counter_cache so we can observe invalidation calls.""" + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = MagicMock() + spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + return spend_counter_cache + + +def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): + """Team-member budget reset clears the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + membership = type( + "Membership", + (), + {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_teammembership.find_many = AsyncMock( + return_value=[membership] + ) + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + counter_cache.redis_cache.async_set_cache.assert_any_await( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Key budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["key"] = [ + type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": now, + "id": "key-1", + "token": "sk-abc", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-abc", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_users_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """User budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 50.0, + "budget_duration": "7d", + "budget_reset_at": now, + "id": "user-1", + "user_id": "alice", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:user:alice", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_teams_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Team budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["team"] = [ + type( + "Team", + (), + { + "spend": 200.0, + "budget_duration": "1mo", + "budget_reset_at": now, + "id": "team-1", + "team_id": "team-x", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch): + """Resetting keys via budget tier must clear each linked key's counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + linked_key = type("Key", (), {"token": "sk-linked"}) + + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[linked_key] + ) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-linked", value=0.0, ttl=60 + ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e0b6d229e2..37e5300565 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware: assert attempts == [ "called" ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory(): + """When Redis is reachable and cleanly returns None (TTL expired, + counter genuinely absent), the read must reseed from DB - NOT fall + through to per-pod in-memory which only contains this pod's writes. + + Pre-fix in multi-pod deployments, in-memory contained a stale local + subset (e.g. $30) while DB had the true cross-pod total ($500). The + fall-through returned $30, enforcement passed, bypass. + """ + 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-1" + + # Per-pod stale in-memory: only this pod's writes, not cross-pod truth. + counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0) + + # Redis cleanly returns None (key expired or never written on this pod). + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_increment = AsyncMock(return_value=500.0) + counter_cache.redis_cache = fake_redis + + # DB has the authoritative cross-pod spend. + db_row = MagicMock() + db_row.spend = 500.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=counter_key, fallback_spend=0.0) + assert spend == 500.0, ( + f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} " + f"(stale per-pod in-memory $30 would have caused multi-pod bypass)" + ) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_error_falls_back_to_in_memory(): + """When Redis raises, the read should still degrade to in-memory rather + than going straight to DB - in-memory is at least same-pod-fresh and + cheaper than a DB query during a Redis outage.""" + 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-1" + + counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + return_value=MagicMock(spend=999.0) + ) + + 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=counter_key, fallback_spend=0.0) + assert spend == 42.0, ( + f"expected in-memory fallback 42.0 on Redis error, got {spend} " + f"(should not have hit DB when Redis errored)" + ) + # DB query should NOT have fired - in-memory short-circuits. + fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma