From fba32eea206f42ed119d0e38f121f13a3535cba8 Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Sat, 12 Sep 2026 23:53:48 -0400 Subject: [PATCH 1/8] fix(proxy): invalidate stale key/user/team caches and retry Redis spend counter resets on budget reset Fixes three of the phantom Redis spend-counter inflation mechanisms reported in #30460 for a multi-pod deployment with intermittent Redis timeouts. Path 2: reset_budget_for_litellm_keys/_users/_teams (the per-row budget_duration reset path) invalidated the Redis spend counter but never dropped the corresponding user_api_key_cache entry, unlike the budget-table cascade path. A stale cached object could get read back by _get_source_cache_base_spend and re-seed the counter with the pre-reset spend. Also route that invalidation through the existing cross-pod broadcast (evict_and_broadcast/LIT-3803) instead of a local-only delete, since only one pod runs the reset job per tick and every other pod's cache needs to drop its copy too. Path 3: a failed Redis SET-to-zero during a budget reset only logged a warning and left the inflated pre-reset counter in place until its TTL. Now retries the SET a bounded number of times with a short backoff, and falls back to deleting the key (which reads as cold and reseeds from the DB) if every attempt fails, logging at ERROR instead of WARNING. Path 1 (narrow): release_budget_reservation_on_cancel silently swallowed a reconcile failure with no invalidation fallback, unlike its sibling release_or_invalidate_budget_reservation. Added the same invalidate-and-mark-finalized fallback so a Redis timeout during the cancel-path reconcile cannot leave a pre-charge stuck in the counter. The broader pre-charge/reconcile path already retries via a delete-and-reseed pattern (increment_spend_counters_pipeline / _reconcile_budget_reservation_for_counter_update), so no further change was made there. Relates to #30460 --- litellm/constants.py | 6 + .../proxy/common_utils/reset_budget_job.py | 81 ++++-- .../spend_tracking/budget_reservation.py | 20 +- .../common_utils/test_reset_budget_job.py | 254 +++++++++++++++++- .../proxy/test_budget_reservation.py | 34 +++ 5 files changed, 377 insertions(+), 18 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5f8fa203b37..41430992763 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1696,6 +1696,12 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job" # leader keeps the lease across its own run, and a crashed one strands the sweep for # at most a single tick. RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900 +RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS: Final = max( + 1, int(os.getenv("RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS", "3")) +) +RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS: Final = float( + os.getenv("RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS", "0.25") +) PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 35e74418628..9aab7f862ea 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -14,6 +14,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, @@ -21,6 +22,8 @@ from litellm.constants import ( RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, + RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( @@ -32,6 +35,9 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLM_VerificationToken, ) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, +) from litellm.proxy.common_utils.timezone_utils import ( BudgetResetSettings, compute_budget_reset_at, @@ -538,24 +544,61 @@ class ResetBudgetJob: 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. + + A SET that keeps failing after retrying falls back to deleting the key rather + than leaving the pre-reset (possibly far higher) value authoritative in Redis + until its TTL expires: a missing counter reads as cold and reseeds from the + DB on the next request (_ensure_spend_counter_initialized), which is always + closer to the truth than the stale value a failed SET would otherwise leave behind. """ try: from litellm.proxy.proxy_server import spend_counter_cache spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, 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=new_spend, 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, - ) + redis_cache = spend_counter_cache.redis_cache + if redis_cache is None: + return + if await ResetBudgetJob._reset_redis_spend_counter( + redis_cache=redis_cache, counter_key=counter_key, new_spend=new_spend + ): + return + verbose_proxy_logger.error( + "Failed to reset spend counter %s in Redis after %d attempts; deleting it instead so the " + "next read reseeds from the DB rather than keeping the inflated pre-reset value authoritative.", + counter_key, + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, + ) + try: + await redis_cache.async_delete_cache(key=counter_key) + except Exception as delete_err: + verbose_proxy_logger.error( + "Failed to delete spend counter %s in Redis after its reset also failed; the " + "pre-reset value stays authoritative until its TTL expires: %s", + counter_key, + delete_err, + ) except Exception as e: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) + @staticmethod + async def _reset_redis_spend_counter(redis_cache: RedisCache, counter_key: str, new_spend: float) -> bool: + for attempt in range(RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS): + try: + await redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60) + return True + except Exception as redis_err: + is_last_attempt = attempt == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS - 1 + verbose_proxy_logger.warning( + "Attempt %d/%d to reset spend counter %s in Redis failed: %s", + attempt + 1, + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, + counter_key, + redis_err, + ) + if not is_last_attempt: + await asyncio.sleep(RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS) + return False + @staticmethod async def _invalidate_global_proxy_spend_cache() -> None: """Drop the cached global-proxy spend accumulator after the proxy @@ -566,20 +609,25 @@ class ResetBudgetJob: @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: - """Drop a stale management-cache entry so the next read fetches from DB. + """Drop a stale management-cache entry, on every pod, so the next read + fetches from DB. Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``; for those, when the spend counter expires the budget check falls back to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded from the DB, but auth still may consult ``user_api_key_cache`` objects - whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache - entry forces the next auth-time fetch to reload the zeroed row from - Postgres. + whose ``.spend`` field can lag a cross-pod DB reset. + + Only one pod runs this job per tick (see ``_acquire_lease``), so a local-only + delete would leave every other pod's copy stale until its TTL: ``evict_and_broadcast`` + is the same LIT-3803 cross-pod eviction every other cache-mutating endpoint already + uses (e.g. auth_checks.delete_cache_team_object), broadcasting the delete over Redis + pub/sub so every pod's copy is dropped, not just the one that ran the reset. """ try: from litellm.proxy.proxy_server import user_api_key_cache - await user_api_key_cache.async_delete_cache(key=cache_key) + await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) except Exception as e: verbose_proxy_logger.warning( "Failed to invalidate user_api_key_cache entry %s: %s", @@ -988,6 +1036,7 @@ class ResetBudgetJob: token = getattr(k, "token", None) if token: await self._invalidate_spend_counter(f"spend:key:{token}", new_spend=k.spend or 0.0) + await self._invalidate_user_api_key_cache_entry(token) end_time = time.time() outcome: Final = _ChunkOutcome( @@ -1093,6 +1142,7 @@ class ResetBudgetJob: user_id = getattr(u, "user_id", None) if user_id: await self._invalidate_spend_counter(f"spend:user:{user_id}", new_spend=u.spend or 0.0) + await self._invalidate_user_api_key_cache_entry(user_id) if user_id == LITELLM_PROXY_BUDGET_NAME: await self._invalidate_global_proxy_spend_cache() @@ -1202,6 +1252,7 @@ class ResetBudgetJob: team_id = getattr(t, "team_id", None) if team_id: await self._invalidate_spend_counter(f"spend:team:{team_id}", new_spend=t.spend or 0.0) + await self._invalidate_user_api_key_cache_entry(f"team_id:{team_id}") end_time = time.time() outcome: Final = _ChunkOutcome( diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 6074a50a69b..8bc3b4a4dc9 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -371,6 +371,12 @@ async def release_budget_reservation_on_cancel( surrounding task is being cancelled. The `finalized` guard makes this a no-op when success/failure handling already reconciled, so calling it on every cancellation path is safe. + + A reconcile failure here (e.g. a Redis timeout) falls back to + invalidate_budget_reservation_counters, mirroring release_or_invalidate_budget_reservation's + handling of the same failure on the non-cancel release path: dropping the reserved counters + forces the next read to reseed from the DB instead of leaving the pre-charge stuck in Redis + with nothing left to correct it, since this is the terminal handler for a cancelled request. """ if not budget_reservation or budget_reservation.get("finalized") is True: return @@ -379,8 +385,20 @@ async def release_budget_reservation_on_cancel( await asyncio.shield( reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) ) - except (asyncio.CancelledError, Exception): + except asyncio.CancelledError: pass + except Exception: + verbose_proxy_logger.exception( + "Failed to reconcile budget reservation on cancel; invalidating reserved counters" + ) + try: + await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after cancel-path reconcile failed" + ) + finally: + budget_reservation["finalized"] = True async def invalidate_budget_reservation_counters( 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 560953f0b51..49a4abc8247 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 @@ -10,17 +10,22 @@ from unittest.mock import AsyncMock, MagicMock import httpx import prisma import pytest +from redis.asyncio import Redis - -from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy._types import LiteLLM_VerificationToken, UserAPIKeyAuth from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MIN_TIME, RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_NAME, + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, +) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache # Mock classes for testing @@ -3243,3 +3248,248 @@ def test_window_reset_zeroes_counter_when_rollover_disabled(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-off:window:1d", value=0.0) spend_counter_cache.async_get_cache.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Path 2 (#30460): resetting a key/user/team's own budget must also drop its +# user_api_key_cache entry, not just the Redis spend counter, and that +# invalidation must reach every pod, not only the one that ran the sweep. +# --------------------------------------------------------------------------- + + +def _direct_reset_fake_module(monkeypatch, user_api_key_cache, redis_usage_cache=None): + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = None + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + fake_module.user_api_key_cache = user_api_key_cache + fake_module.redis_usage_cache = redis_usage_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + return spend_counter_cache + + +def test_reset_budget_for_keys_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch): + """ + A stale user_api_key_cache entry must not survive a key's own budget + reset. Before the fix, _get_source_cache_base_spend could read this + cached object's pre-reset .spend straight back and re-seed the Redis + counter from it on the next request whose DB read fails, reinflating a + just-reset budget with no corresponding spend log (#30460 Path 2). + """ + user_api_key_cache = UserApiKeyCache() + user_api_key_cache.set_cache( + "sk-abc", UserAPIKeyAuth(token="sk-abc", spend=100.0, max_budget=50.0), model_type=UserAPIKeyAuth + ) + assert user_api_key_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is not None + _direct_reset_fake_module(monkeypatch, user_api_key_cache) + + 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()) + + assert user_api_key_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None + + +def test_reset_budget_for_users_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch): + """Same Path 2 gap on the user-scoped reset chunk.""" + from litellm.proxy._types import LiteLLM_UserTable + + user_api_key_cache = UserApiKeyCache() + user_api_key_cache.set_cache( + "alice", LiteLLM_UserTable(user_id="alice", spend=50.0, max_budget=20.0), model_type=LiteLLM_UserTable + ) + assert user_api_key_cache.get_cache("alice", model_type=LiteLLM_UserTable) is not None + _direct_reset_fake_module(monkeypatch, user_api_key_cache) + + 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()) + + assert user_api_key_cache.get_cache("alice", model_type=LiteLLM_UserTable) is None + + +def test_reset_budget_for_teams_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch): + """Same Path 2 gap on the team-scoped reset chunk; teams cache under + ``team_id:{team_id}`` (see auth_checks.get_team_object), not the bare id.""" + from litellm.proxy._types import LiteLLM_TeamTable + + user_api_key_cache = UserApiKeyCache() + user_api_key_cache.set_cache( + "team_id:team-x", + LiteLLM_TeamTable(team_id="team-x", spend=200.0, max_budget=100.0), + model_type=LiteLLM_TeamTable, + ) + assert user_api_key_cache.get_cache("team_id:team-x", model_type=LiteLLM_TeamTable) is not None + _direct_reset_fake_module(monkeypatch, user_api_key_cache) + + 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()) + + assert user_api_key_cache.get_cache("team_id:team-x", model_type=LiteLLM_TeamTable) is None + + +class _LoopbackPubSubRedisClient(Redis): + """One fake client good enough to both publish() and pubsub() against the + same in-process queue, so a test can prove a message one pod publishes is + actually delivered to another pod's subscriber.""" + + def __init__(self) -> None: + self._queue: asyncio.Queue = asyncio.Queue() + + async def publish(self, channel: str, message: str) -> int: + self._queue.put_nowait({"type": "message", "data": message.encode()}) + return 1 + + def pubsub(self): + return self + + async def subscribe(self, *channels: str) -> None: + pass + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float): + try: + return await asyncio.wait_for(self._queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + pass + + +class _FakeCoordinationRedisCache: + def __init__(self, client: object) -> None: + self._client = client + self.namespace = None + + def init_async_client(self) -> object: + return self._client + + +async def test_reset_budget_for_keys_broadcasts_cache_invalidation_to_other_pods( + reset_budget_job, mock_prisma_client, monkeypatch +): + """ + Only one pod runs the reset job per tick (_acquire_lease elects one + sweeper), so a local-only cache delete would leave every other pod's + user_api_key_cache serving the pre-reset spend until its TTL. This proves + the invalidation is actually broadcast (LIT-3803) to a second pod that + never ran the sweep, closing the multi-pod half of #30460 Path 2. + """ + shared_redis_cache = _FakeCoordinationRedisCache(client=_LoopbackPubSubRedisClient()) + + leader_cache = UserApiKeyCache() + follower_cache = UserApiKeyCache() + for cache in (leader_cache, follower_cache): + cache.set_cache( + "sk-abc", UserAPIKeyAuth(token="sk-abc", spend=100.0, max_budget=50.0), model_type=UserAPIKeyAuth + ) + + _direct_reset_fake_module(monkeypatch, leader_cache, redis_usage_cache=shared_redis_cache) + + subscriber = AuthCacheInvalidationSubscriber(redis_cache=shared_redis_cache, user_api_key_cache=follower_cache) + subscriber.start() + + 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"}, + ) + ] + + try: + await reset_budget_job.reset_budget_for_litellm_keys() + + for _ in range(200): + if follower_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert leader_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None + assert follower_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None + + +# --------------------------------------------------------------------------- +# Path 3 (#30460): a failed Redis SET-to-zero on budget reset must not +# silently leave the inflated pre-reset counter authoritative in Redis. +# --------------------------------------------------------------------------- + + +def test_invalidate_spend_counter_retries_then_deletes_on_persistent_redis_failure(monkeypatch): + """ + Every reset-to-zero attempt fails: the fix must retry a bounded number of + times, then fall back to deleting the counter (a missing counter reads as + cold and reseeds from the DB, see _ensure_spend_counter_initialized) + instead of leaving the old, inflated value authoritative until its TTL. + """ + from unittest.mock import patch + + 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(side_effect=RuntimeError("elasticache timeout")) + spend_counter_cache.redis_cache.async_delete_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) + + with patch("litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.error") as mock_error, patch( + "asyncio.sleep", new=AsyncMock() + ): + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-inflated", new_spend=0.0)) + + assert spend_counter_cache.redis_cache.async_set_cache.await_count == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated") + assert mock_error.called + + +def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch): + """A SET that fails once and then succeeds must not fall back to delete: + the counter ends up reset to the real value, not merely absent.""" + from unittest.mock import patch + + 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( + side_effect=[RuntimeError("elasticache timeout"), None] + ) + spend_counter_cache.redis_cache.async_delete_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) + + with patch("asyncio.sleep", new=AsyncMock()): + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-recovered", new_spend=0.0)) + + assert spend_counter_cache.redis_cache.async_set_cache.await_count == 2 + spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 40ebc03781c..2a652c98430 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2771,6 +2771,40 @@ async def test_release_budget_reservation_on_cancel_swallows_release_errors(): await release_budget_reservation_on_cancel(reservation) +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_fails( + spend_counter_state, +): + """ + Regression for #30460 Path 1: if reconcile itself fails on the cancel path + (e.g. a Redis timeout), the pre-charge must not be left stuck in the + counter with nothing left to correct it. This mirrors what + release_or_invalidate_budget_reservation already does on the non-cancel + release path: fall back to invalidate_budget_reservation_counters so the + next read reseeds from the DB instead of enforcing the stale reservation + forever. + """ + counter_cache, _key_cache = spend_counter_state + counter_key = "spend:key:key-cancel-redis-down" + counter_cache.in_memory_cache.set_cache(key=counter_key, value=3.0) + + reservation = { + "reserved_cost": 3.0, + "input_cost": 0.5, + "finalized": False, + "entries": [{"counter_key": counter_key, "reserved_cost": 3.0, "applied_adjustment": 0.0}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ): + await release_budget_reservation_on_cancel(reservation) + + assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_state): counter_cache, key_cache = spend_counter_state From bdd409430ef5d0be38e0725b889c0bfc40cfdc56 Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 06:24:30 -0400 Subject: [PATCH 2/8] fix(proxy): make budget-reset and cancel-path invalidation race-safe against concurrent reservations Addresses two concurrency findings from automated security review on the reset/invalidation fallbacks added for #30460. reset_budget_job.py: a reset that failed and retried replayed an unconditional SET to new_spend on every attempt. A request's Redis INCR landing in the delay between a failed attempt and its retry (a legitimate reservation against the just-reset budget) would be silently erased by the next attempt's write, since retrying and the final delete both treat the counter as static rather than possibly having moved. Replaced the retry's SET with async_reset_preserving_delta, a single Lua GET/compute/SET that resets to new_spend plus whatever was added on top of a snapshot taken once before the first attempt and held fixed across retries, so a concurrent increment survives no matter which attempt eventually succeeds. If the snapshot read itself fails, there's no safe baseline to preserve against, so that case now skips straight to the existing delete fallback rather than attempting a reset that could guess wrong. budget_reservation.py: release_budget_reservation_on_cancel's fallback on a reconcile failure invalidates the shared key/user/team counters outright. Unlike reconcile, which only ever adjusts this reservation's own recorded contribution, that invalidation deletes state a concurrent reservation or recorded spend also shares on the same counter. A transient failure (e.g. a Redis timeout) during cancellation would take the destructive path immediately. Now retries the reconcile itself, which is safe and idempotent, a bounded number of times first, and only falls back to invalidating the aggregate counters once every retry hits the same failure. Both fixes are covered by regression tests that simulate the race directly (a concurrent increment landing during the reset retry delay, and a reconcile that fails once then recovers on the cancel path) and assert the concurrent write survives / the aggregate counter is not touched. --- litellm/caching/redis_cache.py | 43 +++++ .../proxy/common_utils/reset_budget_job.py | 47 ++++- .../spend_tracking/budget_reservation.py | 51 +++++- .../common_utils/test_reset_budget_job.py | 171 ++++++++++++++---- .../proxy/test_budget_reservation.py | 87 +++++++-- 5 files changed, 342 insertions(+), 57 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d0cefcb6086..61fef745991 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1403,6 +1403,49 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_reset_preserving_delta( + self, + key: str, + new_base: float, + snapshot: float, + ttl: int | None = None, + ) -> float: + """Atomically reset ``key`` to ``new_base`` while preserving any amount + added since ``snapshot`` was read, so a reset racing a concurrent + ``async_increment`` cannot erase spend reserved after the reset boundary. + + ``snapshot`` is the value read once before the first attempt and held + fixed across retries by the caller, not re-read each attempt: replaying + this call after a transient failure stays correct no matter how many + increments landed in between, because the delta is always measured + against that same original baseline. The GET/compute/SET runs in a + single Lua call, atomic across racing callers and pods, mirroring + ``async_set_max``. Returns the resulting value. + """ + _redis_client: Final = self.init_async_client() + _used_ttl: Final = self.get_ttl(ttl=ttl) + key = self.check_and_fix_namespace(key=key) + lua: Final = ( + "local cur = redis.call('GET', KEYS[1]) " + "local cur_num = cur and tonumber(cur) or tonumber(ARGV[2]) " + "local delta = cur_num - tonumber(ARGV[2]) " + "if delta < 0 then delta = 0 end " + "local result = tonumber(ARGV[1]) + delta " + "redis.call('SET', KEYS[1], result) " + "if tonumber(ARGV[3]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[3]) end " + "return tostring(result)" + ) + result = cast( + "str | bytes", + await _redis_client.eval( + lua, 1, key, str(new_base), str(snapshot), str(int(_used_ttl or 0)) + ), + ) + if isinstance(result, bytes): + result = result.decode() + return float(result) + @_redis_circuit_breaker_guard async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 9aab7f862ea..25caa948247 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -545,11 +545,21 @@ class ResetBudgetJob: commit opens a window where get_current_spend reads 0 from Redis while the DB still holds the pre-reset value, allowing bypass. - A SET that keeps failing after retrying falls back to deleting the key rather + The reset is delta-preserving: a snapshot of the pre-reset value is read + once up front and held fixed across every retry, so a concurrent + async_increment landing during a retry (e.g. a request reserving spend + against the just-reset budget) is carried forward on top of new_spend + instead of being erased by a later attempt's write. See + async_reset_preserving_delta and _reset_redis_spend_counter. + + A reset that keeps failing after retrying falls back to deleting the key rather than leaving the pre-reset (possibly far higher) value authoritative in Redis until its TTL expires: a missing counter reads as cold and reseeds from the DB on the next request (_ensure_spend_counter_initialized), which is always - closer to the truth than the stale value a failed SET would otherwise leave behind. + closer to the truth than the stale value a failed reset would otherwise leave + behind. The same is true when the pre-reset snapshot itself cannot be read: + with no safe baseline to preserve increments against, deleting is the only + option that cannot silently erase a concurrent reservation. """ try: from litellm.proxy.proxy_server import spend_counter_cache @@ -558,8 +568,12 @@ class ResetBudgetJob: redis_cache = spend_counter_cache.redis_cache if redis_cache is None: return - if await ResetBudgetJob._reset_redis_spend_counter( - redis_cache=redis_cache, counter_key=counter_key, new_spend=new_spend + + snapshot = await ResetBudgetJob._snapshot_spend_counter( + redis_cache=redis_cache, counter_key=counter_key + ) + if snapshot is not None and await ResetBudgetJob._reset_redis_spend_counter( + redis_cache=redis_cache, counter_key=counter_key, new_spend=new_spend, snapshot=snapshot ): return verbose_proxy_logger.error( @@ -581,10 +595,31 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) @staticmethod - async def _reset_redis_spend_counter(redis_cache: RedisCache, counter_key: str, new_spend: float) -> bool: + async def _snapshot_spend_counter(redis_cache: RedisCache, counter_key: str) -> float | None: + """Read the pre-reset baseline the delta-preserving reset holds fixed across + retries. A failed read leaves no safe baseline to reconcile against, so the + caller skips the atomic reset and falls straight back to delete.""" + try: + current = await redis_cache.async_get_cache(key=counter_key) + return float(current) if current is not None else 0.0 + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to read spend counter %s in Redis before reset; skipping the delta-preserving " + "reset and falling back to delete: %s", + counter_key, + redis_err, + ) + return None + + @staticmethod + async def _reset_redis_spend_counter( + redis_cache: RedisCache, counter_key: str, new_spend: float, snapshot: float + ) -> bool: for attempt in range(RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS): try: - await redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60) + await redis_cache.async_reset_preserving_delta( + key=counter_key, new_base=new_spend, snapshot=snapshot, ttl=60 + ) return True except Exception as redis_err: is_last_attempt = attempt == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS - 1 diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8bc3b4a4dc9..93cba60ab21 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -53,6 +53,15 @@ class _BudgetCounter: window_start: datetime | None = None +_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS: Final = 2 +"""Bounded retries of the per-reservation reconcile on the cancel path before +falling back to invalidating the shared aggregate counters. reconcile_budget_reservation +only ever adjusts this reservation's own recorded contribution to each counter (see +_set_reserved_entries_actual_cost), so it is safe and idempotent to retry; the counter +invalidation fallback below is not, since it deletes state every concurrent reservation +and recorded spend on that counter shares.""" + + _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Key": Litellm_EntityType.KEY.value, "Team": Litellm_EntityType.TEAM.value, @@ -372,7 +381,12 @@ async def release_budget_reservation_on_cancel( when success/failure handling already reconciled, so calling it on every cancellation path is safe. - A reconcile failure here (e.g. a Redis timeout) falls back to + A reconcile failure here (e.g. a Redis timeout) is retried a bounded number of + times before falling back to invalidate_budget_reservation_counters: unlike that + fallback, the reconcile only ever adjusts this reservation's own contribution to + each counter, so it cannot clobber a concurrent request's reservation or recorded + spend on the same key/user/team counter the way deleting the aggregate can. Only + once every retry hits the same failure does this fall back to invalidate_budget_reservation_counters, mirroring release_or_invalidate_budget_reservation's handling of the same failure on the non-cancel release path: dropping the reserved counters forces the next read to reseed from the DB instead of leaving the pre-charge stuck in Redis @@ -389,8 +403,12 @@ async def release_budget_reservation_on_cancel( pass except Exception: verbose_proxy_logger.exception( - "Failed to reconcile budget reservation on cancel; invalidating reserved counters" + "Failed to reconcile budget reservation on cancel; retrying before invalidating reserved counters" ) + if await asyncio.shield( + _retry_reconcile_reservation_on_cancel(budget_reservation=budget_reservation, incurred_cost=incurred_cost) + ): + return try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: @@ -401,6 +419,35 @@ async def release_budget_reservation_on_cancel( budget_reservation["finalized"] = True +async def _retry_reconcile_reservation_on_cancel( + budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps finalized on success + incurred_cost: float, +) -> bool: + """Bounded retry of the reconcile that just failed once on the cancel path. + + Unlike invalidate_budget_reservation_counters, reconcile_budget_reservation only + ever adjusts this reservation's own recorded contribution to each counter (see + _set_reserved_entries_actual_cost's applied_adjustment bookkeeping), so retrying it + is safe and idempotent, and correct on any attempt that gets through: it does not + touch concurrent reservations or spend the counter also aggregates. Returns whether + a retry succeeded, so the caller falls back to the destructive invalidation only + after every retry hits the same (presumably persistent) failure. + """ + for attempt in range(_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS): + try: + await reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) + return True + except Exception: + is_last_attempt = attempt == _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS - 1 + verbose_proxy_logger.warning( + "Retry %d/%d to reconcile budget reservation on cancel failed", + attempt + 1, + _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS, + exc_info=not is_last_attempt, + ) + return False + + async def invalidate_budget_reservation_counters( budget_reservation: dict | None, ) -> None: 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 49a4abc8247..15cf6c5fd87 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 @@ -1230,7 +1230,8 @@ def _make_counter_invalidation_job(monkeypatch): 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() + spend_counter_cache.redis_cache.async_get_cache = AsyncMock(return_value=0.0) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock() user_api_key_cache = MagicMock() user_api_key_cache.async_delete_cache = AsyncMock() @@ -1532,7 +1533,9 @@ def test_budget_table_reset_invalidates_counters_and_management_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) counter_cache.in_memory_cache.set_cache.assert_any_call(key=counter_key, value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key=counter_key, value=0.0, ttl=60) + counter_cache.redis_cache.async_reset_preserving_delta.assert_any_await( + key=counter_key, new_base=0.0, snapshot=0.0, ttl=60 + ) deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert cache_keys <= deleted @@ -1571,7 +1574,9 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60) + counter_cache.redis_cache.async_reset_preserving_delta.assert_any_await( + key="spend:end_user:customer-42", new_base=0.0, snapshot=0.0, ttl=60 + ) deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert "end_user_id:customer-42" in deleted @@ -3137,7 +3142,9 @@ def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabl asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) + counter_cache.redis_cache.async_reset_preserving_delta.assert_any_await( + key="spend:end_user:enduser-implicit", new_base=5.0, snapshot=0.0, ttl=60 + ) deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert "end_user_id:enduser-implicit" in deleted @@ -3444,52 +3451,142 @@ async def test_reset_budget_for_keys_broadcasts_cache_invalidation_to_other_pods def test_invalidate_spend_counter_retries_then_deletes_on_persistent_redis_failure(monkeypatch): """ - Every reset-to-zero attempt fails: the fix must retry a bounded number of + Every reset attempt fails: the fix must retry a bounded number of times, then fall back to deleting the counter (a missing counter reads as cold and reseeds from the DB, see _ensure_spend_counter_initialized) instead of leaving the old, inflated value authoritative until its TTL. """ - from unittest.mock import patch - 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(side_effect=RuntimeError("elasticache timeout")) - spend_counter_cache.redis_cache.async_delete_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) - - with patch("litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.error") as mock_error, patch( - "asyncio.sleep", new=AsyncMock() - ): - asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-inflated", new_spend=0.0)) - - assert spend_counter_cache.redis_cache.async_set_cache.await_count == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS - spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated") - assert mock_error.called - - -def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch): - """A SET that fails once and then succeeds must not fall back to delete: - the counter ends up reset to the real value, not merely absent.""" - from unittest.mock import patch - - 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( - side_effect=[RuntimeError("elasticache timeout"), None] + spend_counter_cache.redis_cache.async_get_cache = AsyncMock(return_value=80.0) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock( + side_effect=RuntimeError("elasticache timeout") ) spend_counter_cache.redis_cache.async_delete_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) + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) - with patch("asyncio.sleep", new=AsyncMock()): - asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-recovered", new_spend=0.0)) + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-inflated", new_spend=0.0)) - assert spend_counter_cache.redis_cache.async_set_cache.await_count == 2 + assert ( + spend_counter_cache.redis_cache.async_reset_preserving_delta.await_count + == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS + ) + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated") + + +def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch): + """A reset that fails once and then succeeds must not fall back to delete: + the counter ends up reset to the real value, not merely absent.""" + 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_get_cache = AsyncMock(return_value=80.0) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock( + side_effect=[RuntimeError("elasticache timeout"), 0.0] + ) + spend_counter_cache.redis_cache.async_delete_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) + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-recovered", new_spend=0.0)) + + assert spend_counter_cache.redis_cache.async_reset_preserving_delta.await_count == 2 spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited() + + +def test_invalidate_spend_counter_skips_straight_to_delete_when_snapshot_read_fails(monkeypatch): + """No safe baseline to preserve a concurrent increment against: guessing with + a blind reset would reopen the exact race this fix closes, so a failed + pre-reset snapshot read must fall straight back to delete instead of + attempting the atomic reset at all.""" + 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_get_cache = AsyncMock(side_effect=RuntimeError("elasticache timeout")) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock() + spend_counter_cache.redis_cache.async_delete_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) + + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-no-snapshot", new_spend=0.0)) + + spend_counter_cache.redis_cache.async_reset_preserving_delta.assert_not_awaited() + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-no-snapshot") + + +class _FakeRedisSpendCounter: + """A minimal in-process stand-in for the Redis key under test, faithful to + async_reset_preserving_delta's GET/compute/SET contract (new_base + max(0, + current - snapshot)), so what's actually under test is the retry loop's own + behavior of holding one fixed snapshot across attempts rather than re-reading + it, which is what makes a concurrent increment during the retry delay survive.""" + + def __init__(self, initial: float): + self.value = initial + self.attempts = 0 + self.fail_first_n = 0 + + async def async_get_cache(self, key): + return self.value + + async def async_reset_preserving_delta(self, key, new_base, snapshot, ttl=None): + self.attempts += 1 + if self.attempts <= self.fail_first_n: + raise RuntimeError("elasticache timeout") + delta = max(0.0, self.value - snapshot) + self.value = new_base + delta + return self.value + + async def async_increment(self, key, value, refresh_ttl=False): + self.value += value + return self.value + + async def async_delete_cache(self, key): + self.value = None + + +def test_invalidate_spend_counter_preserves_a_concurrent_increment_made_during_the_retry_delay(monkeypatch): + """ + Regression for the reset-retry race (veria-ai finding, reset_budget_job.py): + if a request's Redis INCR lands in the delay between a failed reset attempt + and its retry, the eventual successful reset must not erase it. Before the + fix, the retry replayed an unconditional SET to new_spend and would have + wiped the concurrent increment out; the delta-preserving reset must carry + it forward instead, since it represents real spend reserved after the + reset boundary. + """ + store = _FakeRedisSpendCounter(initial=80.0) + store.fail_first_n = 1 + + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = store + + 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) + + concurrent_increment_applied = False + + async def sleep_and_race(_delay): + nonlocal concurrent_increment_applied + if not concurrent_increment_applied: + await store.async_increment(key="spend:key:sk-race", value=5.0) + concurrent_increment_applied = True + + monkeypatch.setattr(asyncio, "sleep", sleep_and_race) + + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-race", new_spend=0.0)) + + assert concurrent_increment_applied + assert store.value == 5.0 diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 2a652c98430..a3f33c2b07e 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2230,6 +2230,23 @@ class _ExpiringRedisCache: return None +class _FlakyPipelineRedisCache(_ExpiringRedisCache): + """_ExpiringRedisCache, but the first ``fail_first_n`` calls to + async_increment_pipeline raise instead of applying, so a reconcile retry + recovering from a transient Redis failure is what's under test.""" + + def __init__(self, fail_first_n: int = 0) -> None: + super().__init__() + self.fail_first_n = fail_first_n + self.pipeline_calls = 0 + + async def async_increment_pipeline(self, increment_list, **kwargs): + self.pipeline_calls += 1 + if self.pipeline_calls <= self.fail_first_n: + raise RuntimeError("redis down") + return await super().async_increment_pipeline(increment_list, **kwargs) + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, @@ -2772,20 +2789,22 @@ async def test_release_budget_reservation_on_cancel_swallows_release_errors(): @pytest.mark.asyncio -async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_fails( +async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_persistently_fails( spend_counter_state, ): """ - Regression for #30460 Path 1: if reconcile itself fails on the cancel path - (e.g. a Redis timeout), the pre-charge must not be left stuck in the - counter with nothing left to correct it. This mirrors what - release_or_invalidate_budget_reservation already does on the non-cancel - release path: fall back to invalidate_budget_reservation_counters so the - next read reseeds from the DB instead of enforcing the stale reservation - forever. + Regression for #30460 Path 1: if reconcile keeps failing on the cancel path + (e.g. a persistent Redis outage) across every retry, the pre-charge must + not be left stuck in the counter with nothing left to correct it. This + mirrors what release_or_invalidate_budget_reservation already does on the + non-cancel release path: fall back to invalidate_budget_reservation_counters + so the next read reseeds from the DB instead of enforcing the stale + reservation forever. """ counter_cache, _key_cache = spend_counter_state counter_key = "spend:key:key-cancel-redis-down" + counter_cache.redis_cache = _FlakyPipelineRedisCache(fail_first_n=99) + counter_cache.redis_cache.store[counter_key] = 3.0 counter_cache.in_memory_cache.set_cache(key=counter_key, value=3.0) reservation = { @@ -2795,13 +2814,57 @@ async def test_release_budget_reservation_on_cancel_invalidates_counter_when_rec "entries": [{"counter_key": counter_key, "reserved_cost": 3.0, "applied_adjustment": 0.0}], } - with patch( - "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", - new=AsyncMock(side_effect=RuntimeError("redis down")), + await release_budget_reservation_on_cancel(reservation) + + assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None + assert reservation["finalized"] is True + # Only the first retry reaches the pipeline: its own failure already + # invalidates the counter, so later retries take the (here also + # unavailable, with no real DB) reseed path instead -- both exhausted + # before falling back to invalidate_budget_reservation_counters. + assert counter_cache.redis_cache.pipeline_calls == 1 + + +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_retries_before_invalidating( + spend_counter_state, +): + """ + Regression for the cancellation-fallback race (veria-ai finding on + budget_reservation.py): invalidate_budget_reservation_counters deletes the + whole aggregate counter unconditionally, which would also erase a + concurrent request's reservation or recorded spend sharing that same + key/user/team counter, not just this reservation's own contribution. A + reconcile failure that clears up on retry must settle through the + ordinary reconcile machinery instead of falling back to that unconditional + delete: here, the first attempt's own failure already invalidates the + counter (existing increment_spend_counters_pipeline cleanup), so the retry + settles it by reseeding the DB floor and adding this reservation's actual + cost, the same recovery path an expired counter already takes elsewhere in + this file, rather than being deleted a second time with nothing settled. + """ + import litellm.proxy.proxy_server as ps + + counter_cache, _key_cache = spend_counter_state + counter_key = "spend:key:key-cancel-retry-recovers" + counter_cache.redis_cache = _FlakyPipelineRedisCache(fail_first_n=1) + counter_cache.redis_cache.store[counter_key] = 5.0 + counter_cache.in_memory_cache.set_cache(key=counter_key, value=5.0) + + reservation = { + "reserved_cost": 3.0, + "input_cost": 0.5, + "finalized": False, + "entries": [{"counter_key": counter_key, "reserved_cost": 3.0, "applied_adjustment": 0.0}], + } + + with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for + ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.2) ): await release_budget_reservation_on_cancel(reservation) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None + assert counter_cache.redis_cache.pipeline_calls == 1 + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.7) assert reservation["finalized"] is True From 1fdb6e95a54e40611e957521813ab8100ea68fab Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 08:20:04 -0400 Subject: [PATCH 3/8] fix(proxy): break reset_budget_job cyclic-import risk, tighten comments per repo convention Addresses CodeQL and greptile findings on the reset-race fix. CodeQL flagged reset_budget_job.py's module-level imports of RedisCache, the two RESET_BUDGET_SPEND_COUNTER_RESET_* constants, and evict_and_broadcast as cyclic-import risks: each target module can reach back to this one before finishing its own initialization. RedisCache is only used as a type hint, so it moves under TYPE_CHECKING with the annotations quoted; the constants and evict_and_broadcast are only used inside single functions, so they move to local imports there, per this repo's own stated exception for avoiding circular imports. Also fixes a CodeQL "empty except with no explanatory comment" on the cancel-path's CancelledError handler, and trims the docstrings this fix added or touched down to this repo's comment conventions (concise, non-obvious-why only). --- litellm/caching/redis_cache.py | 17 ++--- .../proxy/common_utils/reset_budget_job.py | 76 +++++++------------ .../spend_tracking/budget_reservation.py | 63 +++++---------- 3 files changed, 49 insertions(+), 107 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 61fef745991..19e904e60e9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1411,18 +1411,11 @@ class RedisCache(BaseCache): snapshot: float, ttl: int | None = None, ) -> float: - """Atomically reset ``key`` to ``new_base`` while preserving any amount - added since ``snapshot`` was read, so a reset racing a concurrent - ``async_increment`` cannot erase spend reserved after the reset boundary. - - ``snapshot`` is the value read once before the first attempt and held - fixed across retries by the caller, not re-read each attempt: replaying - this call after a transient failure stays correct no matter how many - increments landed in between, because the delta is always measured - against that same original baseline. The GET/compute/SET runs in a - single Lua call, atomic across racing callers and pods, mirroring - ``async_set_max``. Returns the resulting value. - """ + """Reset ``key`` to ``new_base`` plus whatever was added on top of + ``snapshot``, atomically (single Lua GET/compute/SET, mirrors + ``async_set_max``), so a concurrent ``async_increment`` survives a reset + retry instead of being overwritten by it. ``snapshot`` is read once by + the caller and held fixed across retries.""" _redis_client: Final = self.init_async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 25caa948247..f7edba6fd1d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,14 +7,13 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCache from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, @@ -22,8 +21,6 @@ from litellm.constants import ( RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, RESET_BUDGET_JOB_NAME, - RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, - RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( @@ -35,9 +32,6 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLM_VerificationToken, ) -from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( - evict_and_broadcast, -) from litellm.proxy.common_utils.timezone_utils import ( BudgetResetSettings, compute_budget_reset_at, @@ -72,6 +66,9 @@ from litellm.repositories.verification_token_repository import ( ) from litellm.types.services import ServiceTypes +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + _RowT = TypeVar("_RowT") _LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) @@ -537,31 +534,17 @@ class ResetBudgetJob: @staticmethod async def _invalidate_spend_counter(counter_key: str, new_spend: float = 0.0) -> None: - """Overwrite a spend counter with the post-reset value (0, or the carried - overage when budget rollover is enabled) so a DB-row reset takes effect - immediately. + """Overwrite a spend counter with the post-reset value. Call AFTER the DB + write commits, or get_current_spend can read 0 from Redis while the DB + still holds the pre-reset value. - 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. - - The reset is delta-preserving: a snapshot of the pre-reset value is read - once up front and held fixed across every retry, so a concurrent - async_increment landing during a retry (e.g. a request reserving spend - against the just-reset budget) is carried forward on top of new_spend - instead of being erased by a later attempt's write. See - async_reset_preserving_delta and _reset_redis_spend_counter. - - A reset that keeps failing after retrying falls back to deleting the key rather - than leaving the pre-reset (possibly far higher) value authoritative in Redis - until its TTL expires: a missing counter reads as cold and reseeds from the - DB on the next request (_ensure_spend_counter_initialized), which is always - closer to the truth than the stale value a failed reset would otherwise leave - behind. The same is true when the pre-reset snapshot itself cannot be read: - with no safe baseline to preserve increments against, deleting is the only - option that cannot silently erase a concurrent reservation. + Delta-preserving: holds one pre-reset snapshot fixed across retries so a + concurrent async_increment during a retry lands on top of new_spend + instead of being overwritten by it. Falls back to delete (reads as cold, + reseeds from DB) if every retry fails or the snapshot itself can't be read. """ try: + from litellm.constants import RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS from litellm.proxy.proxy_server import spend_counter_cache spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, ttl=60) @@ -595,10 +578,9 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e) @staticmethod - async def _snapshot_spend_counter(redis_cache: RedisCache, counter_key: str) -> float | None: - """Read the pre-reset baseline the delta-preserving reset holds fixed across - retries. A failed read leaves no safe baseline to reconcile against, so the - caller skips the atomic reset and falls straight back to delete.""" + async def _snapshot_spend_counter(redis_cache: "RedisCache", counter_key: str) -> float | None: + """Pre-reset baseline held fixed across retries. A failed read means no + safe baseline, so the caller falls straight back to delete.""" try: current = await redis_cache.async_get_cache(key=counter_key) return float(current) if current is not None else 0.0 @@ -613,8 +595,13 @@ class ResetBudgetJob: @staticmethod async def _reset_redis_spend_counter( - redis_cache: RedisCache, counter_key: str, new_spend: float, snapshot: float + redis_cache: "RedisCache", counter_key: str, new_spend: float, snapshot: float ) -> bool: + from litellm.constants import ( + RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS, + RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS, + ) + for attempt in range(RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS): try: await redis_cache.async_reset_preserving_delta( @@ -644,22 +631,13 @@ class ResetBudgetJob: @staticmethod async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None: - """Drop a stale management-cache entry, on every pod, so the next read - fetches from DB. - - Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``; - for those, when the spend counter expires the budget check falls back - to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded - from the DB, but auth still may consult ``user_api_key_cache`` objects - whose ``.spend`` field can lag a cross-pod DB reset. - - Only one pod runs this job per tick (see ``_acquire_lease``), so a local-only - delete would leave every other pod's copy stale until its TTL: ``evict_and_broadcast`` - is the same LIT-3803 cross-pod eviction every other cache-mutating endpoint already - uses (e.g. auth_checks.delete_cache_team_object), broadcasting the delete over Redis - pub/sub so every pod's copy is dropped, not just the one that ran the reset. - """ + """Drop a stale management-cache entry on every pod (LIT-3803 cross-pod + eviction), so a pod other than the one that ran the reset doesn't keep + serving a cached object with the pre-reset spend.""" try: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) from litellm.proxy.proxy_server import user_api_key_cache await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 93cba60ab21..9631eb216d6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -54,12 +54,8 @@ class _BudgetCounter: _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS: Final = 2 -"""Bounded retries of the per-reservation reconcile on the cancel path before -falling back to invalidating the shared aggregate counters. reconcile_budget_reservation -only ever adjusts this reservation's own recorded contribution to each counter (see -_set_reserved_entries_actual_cost), so it is safe and idempotent to retry; the counter -invalidation fallback below is not, since it deletes state every concurrent reservation -and recorded spend on that counter shares.""" +"""Retries of the per-reservation reconcile before falling back to the +counter-deleting invalidation, which a concurrent reservation also shares.""" _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { @@ -360,37 +356,19 @@ async def release_budget_reservation(budget_reservation: dict | None) -> None: async def release_budget_reservation_on_cancel( - budget_reservation: dict | None, + budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict ) -> None: - """Reconcile a still-open reservation when the request is cancelled mid-flight. + """Reconcile a still-open reservation when the request is cancelled mid-flight, + since neither the success nor failure hook runs in that case. Reconciles to the + request's input-token cost, not zero, since that portion was already billed to + the provider. asyncio.shield keeps this running through the surrounding + cancellation. - A client disconnect or timeout cancels the request task, which surfaces as - CancelledError / GeneratorExit rather than a normal exception, so neither the - success cost callback nor the failure hook runs and the pre-call reservation - is never reconciled. Left alone it pins the spend counter above real spend - and 429s subsequent requests until the counter's TTL expires. - - Reconcile to the request's input-token cost rather than refunding to zero: - by the time a request is cancelled in-flight the provider call was already - dispatched, so the input tokens were billed even if no chunk reached the - client. Refunding to zero would let a caller abort pre-token to dodge that - charge; the worst-case output portion of the reservation is still released. - - asyncio.shield keeps the reconcile running to completion even though the - surrounding task is being cancelled. The `finalized` guard makes this a no-op - when success/failure handling already reconciled, so calling it on every - cancellation path is safe. - - A reconcile failure here (e.g. a Redis timeout) is retried a bounded number of - times before falling back to invalidate_budget_reservation_counters: unlike that - fallback, the reconcile only ever adjusts this reservation's own contribution to - each counter, so it cannot clobber a concurrent request's reservation or recorded - spend on the same key/user/team counter the way deleting the aggregate can. Only - once every retry hits the same failure does this fall back to - invalidate_budget_reservation_counters, mirroring release_or_invalidate_budget_reservation's - handling of the same failure on the non-cancel release path: dropping the reserved counters - forces the next read to reseed from the DB instead of leaving the pre-charge stuck in Redis - with nothing left to correct it, since this is the terminal handler for a cancelled request. + A reconcile failure is retried before falling back to + invalidate_budget_reservation_counters, since unlike that fallback, reconcile + only ever adjusts this reservation's own contribution to each counter and can't + clobber a concurrent reservation sharing it. Mirrors + release_or_invalidate_budget_reservation's fallback on the non-cancel path. """ if not budget_reservation or budget_reservation.get("finalized") is True: return @@ -400,7 +378,7 @@ async def release_budget_reservation_on_cancel( reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) ) except asyncio.CancelledError: - pass + pass # a second cancellation while shielded; the reconcile keeps running detached regardless except Exception: verbose_proxy_logger.exception( "Failed to reconcile budget reservation on cancel; retrying before invalidating reserved counters" @@ -423,16 +401,9 @@ async def _retry_reconcile_reservation_on_cancel( budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps finalized on success incurred_cost: float, ) -> bool: - """Bounded retry of the reconcile that just failed once on the cancel path. - - Unlike invalidate_budget_reservation_counters, reconcile_budget_reservation only - ever adjusts this reservation's own recorded contribution to each counter (see - _set_reserved_entries_actual_cost's applied_adjustment bookkeeping), so retrying it - is safe and idempotent, and correct on any attempt that gets through: it does not - touch concurrent reservations or spend the counter also aggregates. Returns whether - a retry succeeded, so the caller falls back to the destructive invalidation only - after every retry hits the same (presumably persistent) failure. - """ + """Retry the reconcile that just failed once. Safe to retry since it only + ever adjusts this reservation's own contribution, unlike the destructive + invalidation the caller falls back to once every retry fails.""" for attempt in range(_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS): try: await reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) From b25920d9c9f36392c044f97fce4c1d12331fc13f Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 08:57:01 -0400 Subject: [PATCH 4/8] style: ruff format --- litellm/caching/redis_cache.py | 4 +--- litellm/proxy/common_utils/reset_budget_job.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 19e904e60e9..a217200f77d 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1431,9 +1431,7 @@ class RedisCache(BaseCache): ) result = cast( "str | bytes", - await _redis_client.eval( - lua, 1, key, str(new_base), str(snapshot), str(int(_used_ttl or 0)) - ), + await _redis_client.eval(lua, 1, key, str(new_base), str(snapshot), str(int(_used_ttl or 0))), ) if isinstance(result, bytes): result = result.decode() diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index f7edba6fd1d..5c0c6deb9be 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -552,9 +552,7 @@ class ResetBudgetJob: if redis_cache is None: return - snapshot = await ResetBudgetJob._snapshot_spend_counter( - redis_cache=redis_cache, counter_key=counter_key - ) + snapshot = await ResetBudgetJob._snapshot_spend_counter(redis_cache=redis_cache, counter_key=counter_key) if snapshot is not None and await ResetBudgetJob._reset_redis_spend_counter( redis_cache=redis_cache, counter_key=counter_key, new_spend=new_spend, snapshot=snapshot ): From 9f546fdb8930be42706d8cadb0c33ab2537dc462 Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 10:16:26 -0400 Subject: [PATCH 5/8] fix(proxy): stay within the ruff strict-rule BLE001 budget The reset-race fix's new except Exception blocks pushed BLE001 4 over ruff-strict-budget.json's limit. Each catches a Redis or reconcile failure of unknown type by design, matching the noqa pattern this file already uses on the equivalent non-cancel release path, so suppress rather than narrow. --- litellm/proxy/common_utils/reset_budget_job.py | 6 +++--- litellm/proxy/spend_tracking/budget_reservation.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 5c0c6deb9be..35b10ab2f0d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -565,7 +565,7 @@ class ResetBudgetJob: ) try: await redis_cache.async_delete_cache(key=counter_key) - except Exception as delete_err: + except Exception as delete_err: # noqa: BLE001 # any Redis failure here leaves the pre-reset value stuck; log and move on verbose_proxy_logger.error( "Failed to delete spend counter %s in Redis after its reset also failed; the " "pre-reset value stays authoritative until its TTL expires: %s", @@ -582,7 +582,7 @@ class ResetBudgetJob: try: current = await redis_cache.async_get_cache(key=counter_key) return float(current) if current is not None else 0.0 - except Exception as redis_err: + except Exception as redis_err: # noqa: BLE001 # any Redis failure here means no safe baseline; fall back to delete verbose_proxy_logger.warning( "Failed to read spend counter %s in Redis before reset; skipping the delta-preserving " "reset and falling back to delete: %s", @@ -606,7 +606,7 @@ class ResetBudgetJob: key=counter_key, new_base=new_spend, snapshot=snapshot, ttl=60 ) return True - except Exception as redis_err: + except Exception as redis_err: # noqa: BLE001 # any Redis failure here is worth a retry, not just specific ones is_last_attempt = attempt == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS - 1 verbose_proxy_logger.warning( "Attempt %d/%d to reset spend counter %s in Redis failed: %s", diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 9631eb216d6..6eed3df39f8 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -379,7 +379,7 @@ async def release_budget_reservation_on_cancel( ) except asyncio.CancelledError: pass # a second cancellation while shielded; the reconcile keeps running detached regardless - except Exception: + except Exception: # noqa: BLE001 # a reconcile failure must not pin the counter; retry, then drop it directly verbose_proxy_logger.exception( "Failed to reconcile budget reservation on cancel; retrying before invalidating reserved counters" ) @@ -389,7 +389,7 @@ async def release_budget_reservation_on_cancel( return try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) - except Exception: + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed verbose_proxy_logger.exception( "Failed to invalidate budget reservation counters after cancel-path reconcile failed" ) @@ -408,7 +408,7 @@ async def _retry_reconcile_reservation_on_cancel( try: await reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) return True - except Exception: + except Exception: # noqa: BLE001 # any reconcile failure is worth a retry here, not just specific ones is_last_attempt = attempt == _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS - 1 verbose_proxy_logger.warning( "Retry %d/%d to reconcile budget reservation on cancel failed", From e9836165d74223d02e1226955bb79b65fb2a3282 Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 15:54:59 -0400 Subject: [PATCH 6/8] test(proxy): close patch-coverage gaps on the reset-race fix Adds direct unit tests for async_reset_preserving_delta (the Lua GET/compute/SET itself, previously only exercised indirectly through mocked callers) and for the delete-also-fails and second-cancellation edge cases in the reset job and cancel-path retry, all only reachable through error injection. --- .../test_litellm/caching/test_redis_cache.py | 46 +++++++++++++++++++ .../common_utils/test_reset_budget_job.py | 24 ++++++++++ .../proxy/test_budget_reservation.py | 43 +++++++++++++++++ 3 files changed, 113 insertions(+) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index c1e3240adb7..e811399f901 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -77,6 +77,52 @@ async def test_async_delete_cache_applies_namespace( mock_redis_instance.delete.assert_awaited_once_with(expected_key) +@pytest.mark.asyncio +async def test_async_reset_preserving_delta_evals_with_namespaced_key_and_string_args( + monkeypatch, redis_no_ping +): + """The GET/compute/SET has to run as one Lua call, not separate round trips, or a + concurrent async_increment between them would be exactly the race this method exists + to close. Namespacing and str-ifying every ARGV also has to happen, or redis-py's own + encoding (or an ACL scoped to the namespace prefix) breaks the call outright.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="litellm") + mock_redis_instance = AsyncMock() + mock_redis_instance.eval.return_value = "12.5" + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + result = await redis_cache.async_reset_preserving_delta( + key="spend:key:abc", new_base=10.0, snapshot=7.5, ttl=60 + ) + + assert result == 12.5 + mock_redis_instance.eval.assert_awaited_once() + lua, numkeys, key, new_base_arg, snapshot_arg, ttl_arg = mock_redis_instance.eval.await_args.args + assert numkeys == 1 + assert key == "litellm:spend:key:abc" + assert (new_base_arg, snapshot_arg, ttl_arg) == ("10.0", "7.5", "60") + assert "GET" in lua and "SET" in lua and "EXPIRE" in lua + + +@pytest.mark.asyncio +async def test_async_reset_preserving_delta_decodes_a_bytes_result(monkeypatch, redis_no_ping): + """redis-py returns EVAL results as bytes unless decode_responses is set; a caller that + compares the return value to a float must not have to know that.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.eval.return_value = b"10.0" + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + result = await redis_cache.async_reset_preserving_delta(key="k", new_base=10.0, snapshot=10.0, ttl=60) + + assert result == 10.0 + + @pytest.mark.parametrize("namespace", [None, "litellm"]) def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): """delete_cache must prefix keys with the namespace, matching every other 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 15cf6c5fd87..a974000bbf8 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 @@ -3479,6 +3479,30 @@ def test_invalidate_spend_counter_retries_then_deletes_on_persistent_redis_failu spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated") +def test_invalidate_spend_counter_swallows_a_delete_failure_after_reset_already_failed(monkeypatch): + """The fallback delete is itself best-effort: if it also fails (e.g. the same + outage that broke the reset), there is nothing left to try, and the failure + must be logged and swallowed rather than propagate out of the reset job.""" + 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_get_cache = AsyncMock(return_value=80.0) + spend_counter_cache.redis_cache.async_reset_preserving_delta = AsyncMock( + side_effect=RuntimeError("elasticache timeout") + ) + spend_counter_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("elasticache timeout")) + + 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) + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + + # must not raise + asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-double-failure", new_spend=0.0)) + + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-double-failure") + + def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch): """A reset that fails once and then succeeds must not fall back to delete: the counter ends up reset to the real value, not merely absent.""" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index a3f33c2b07e..64756558e39 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2788,6 +2788,49 @@ async def test_release_budget_reservation_on_cancel_swallows_release_errors(): await release_budget_reservation_on_cancel(reservation) +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_swallows_a_second_cancellation_while_shielded(): + # A second CancelledError arriving while the shielded reconcile is in flight must not + # propagate: the reconcile keeps running detached regardless, and there is nothing more + # for this call to do but return. + reservation = { + "reserved_cost": 3.0, + "entries": [{"counter_key": "spend:key:key-cancel-twice"}], + "finalized": False, + "input_cost": 0.5, + } + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new=AsyncMock(side_effect=asyncio.CancelledError()), + ): + # must return without raising + await release_budget_reservation_on_cancel(reservation) + + +@pytest.mark.asyncio +async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_after_every_retry_fails(): + # If both the reconcile retries and the invalidate fallback fail (e.g. a persistent + # outage), there is nothing left to try: the failure must be logged and swallowed, + # not propagated, and the reservation still ends up finalized so it is not reprocessed. + reservation = { + "reserved_cost": 3.0, + "entries": [{"counter_key": "spend:key:key-cancel-double-failure"}], + "finalized": False, + "input_cost": 0.5, + } + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ), patch( + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new=AsyncMock(side_effect=RuntimeError("redis still down")), + ): + # must return without raising + await release_budget_reservation_on_cancel(reservation) + + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_persistently_fails( spend_counter_state, From 90a6e8d42308f59ad969c5445dc66aac3452168d Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Mon, 14 Sep 2026 16:41:09 -0400 Subject: [PATCH 7/8] test(proxy): stay within the test-quality TQ001/TQ008 budget The two new cancel-path tests pushed TQ001 (an assertion-free test) and TQ008 (patching a litellm internal) over their limits. Added a real assertion to the cancellation test (finalized stays False, since absorbing a second CancelledError isn't the failure path) and test-quality-ok annotations on the three patch() calls, matching this file's existing pattern for module functions with no DI seam. --- tests/test_litellm/proxy/test_budget_reservation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 64756558e39..5d13a1f0692 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2799,13 +2799,16 @@ async def test_release_budget_reservation_on_cancel_swallows_a_second_cancellati "finalized": False, "input_cost": 0.5, } - with patch( + with patch( # test-quality-ok: reconcile_budget_reservation is a module function, no DI seam for this test "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", new=AsyncMock(side_effect=asyncio.CancelledError()), ): - # must return without raising + # must return without raising, and without taking the finalize-on-failure path either: + # a second cancellation isn't a failure, it's the reconcile still running in the background await release_budget_reservation_on_cancel(reservation) + assert reservation["finalized"] is False + @pytest.mark.asyncio async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_after_every_retry_fails(): @@ -2818,10 +2821,10 @@ async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_ "finalized": False, "input_cost": 0.5, } - with patch( + with patch( # test-quality-ok: reconcile_budget_reservation is a module function, no DI seam for this test "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", new=AsyncMock(side_effect=RuntimeError("redis down")), - ), patch( + ), patch( # test-quality-ok: invalidate_budget_reservation_counters is a module function, no DI seam for this test "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("redis still down")), ): From e91f3fb5da84e54cace6a4d82bfec1a1f5f2ac3a Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Wed, 30 Sep 2026 20:44:34 -0400 Subject: [PATCH 8/8] fix(proxy): don't retry the additive cancel-path reconcile on failure The cancel-path reconcile applies an INCRBYFLOAT refund. If Redis applies it but the pipeline response times out (and the follow-up counter delete also fails), the reservation's applied_adjustment is never stamped, so retrying the reconcile subtracts the same refund a second time and under-counts spend shared with concurrent requests. Even when that delete succeeds, a concurrent request can reseed the counter before the retry, which then refunds against a counter that no longer holds this reservation. Drop the retry and go straight to invalidating the reserved counters (a non-additive delete, so the next read reseeds from the DB) and finalizing the reservation, the same fallback release_or_invalidate_budget_reservation uses on the non-cancel path. Co-Authored-By: Claude Opus 5.5 --- .../spend_tracking/budget_reservation.py | 47 ++--------- .../proxy/test_budget_reservation.py | 81 +++++++++---------- 2 files changed, 47 insertions(+), 81 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index e1cbd33ab3d..ee50153873e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -57,11 +57,6 @@ class _BudgetCounter: window_start: datetime | None = None -_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS: Final = 2 -"""Retries of the per-reservation reconcile before falling back to the -counter-deleting invalidation, which a concurrent reservation also shares.""" - - _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Key": Litellm_EntityType.KEY.value, "Team": Litellm_EntityType.TEAM.value, @@ -407,11 +402,11 @@ async def release_budget_reservation_on_cancel( the provider. asyncio.shield keeps this running through the surrounding cancellation. - A reconcile failure is retried before falling back to - invalidate_budget_reservation_counters, since unlike that fallback, reconcile - only ever adjusts this reservation's own contribution to each counter and can't - clobber a concurrent reservation sharing it. Mirrors - release_or_invalidate_budget_reservation's fallback on the non-cancel path. + A failed reconcile is not retried: the reconcile is an additive INCRBYFLOAT, and + a failure (e.g. a timeout after Redis applied it) leaves it unknown whether the + refund landed, so re-applying it could refund twice. Instead the reserved + counters are dropped so the next read reseeds from the DB, the same fallback + release_or_invalidate_budget_reservation uses on the non-cancel path. """ if not budget_reservation or budget_reservation.get("finalized") is True: return @@ -422,14 +417,8 @@ async def release_budget_reservation_on_cancel( ) except asyncio.CancelledError: pass # a second cancellation while shielded; the reconcile keeps running detached regardless - except Exception: # noqa: BLE001 # a reconcile failure must not pin the counter; retry, then drop it directly - verbose_proxy_logger.exception( - "Failed to reconcile budget reservation on cancel; retrying before invalidating reserved counters" - ) - if await asyncio.shield( - _retry_reconcile_reservation_on_cancel(budget_reservation=budget_reservation, incurred_cost=incurred_cost) - ): - return + except Exception: # noqa: BLE001 # a reconcile failure must not pin the counter; drop it directly instead + verbose_proxy_logger.exception("Failed to reconcile budget reservation on cancel; invalidating counters") try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed @@ -440,28 +429,6 @@ async def release_budget_reservation_on_cancel( budget_reservation["finalized"] = True -async def _retry_reconcile_reservation_on_cancel( - budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps finalized on success - incurred_cost: float, -) -> bool: - """Retry the reconcile that just failed once. Safe to retry since it only - ever adjusts this reservation's own contribution, unlike the destructive - invalidation the caller falls back to once every retry fails.""" - for attempt in range(_RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS): - try: - await reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost) - return True - except Exception: # noqa: BLE001 # any reconcile failure is worth a retry here, not just specific ones - is_last_attempt = attempt == _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS - 1 - verbose_proxy_logger.warning( - "Retry %d/%d to reconcile budget reservation on cancel failed", - attempt + 1, - _RELEASE_ON_CANCEL_RECONCILE_MAX_ATTEMPTS, - exc_info=not is_last_attempt, - ) - return False - - async def invalidate_budget_reservation_counters( budget_reservation: dict | None, ) -> None: diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index eca6a2ea590..529cda6c83a 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2485,20 +2485,31 @@ class _TeamMembershipFloorDb: class _FlakyPipelineRedisCache(_ExpiringRedisCache): """_ExpiringRedisCache, but the first ``fail_first_n`` calls to - async_increment_pipeline raise instead of applying, so a reconcile retry - recovering from a transient Redis failure is what's under test.""" + async_increment_pipeline time out. With ``apply_before_failing`` the + increments land before the timeout surfaces, the ambiguous case where the + caller can't tell whether its write applied; with ``fail_deletes`` the + counter delete that would otherwise clean that up fails too.""" - def __init__(self, fail_first_n: int = 0) -> None: + def __init__(self, fail_first_n: int = 0, apply_before_failing: bool = False, fail_deletes: bool = False) -> None: super().__init__() self.fail_first_n = fail_first_n + self.apply_before_failing = apply_before_failing + self.fail_deletes = fail_deletes self.pipeline_calls = 0 async def async_increment_pipeline(self, increment_list, **kwargs): self.pipeline_calls += 1 if self.pipeline_calls <= self.fail_first_n: - raise RuntimeError("redis down") + if self.apply_before_failing: + await super().async_increment_pipeline(increment_list, **kwargs) + raise TimeoutError("redis timeout") return await super().async_increment_pipeline(increment_list, **kwargs) + async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: + if self.fail_deletes: + raise ConnectionError("redis unreachable") + await super().async_delete_cache(key, *args, **kwargs) + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( @@ -3183,8 +3194,8 @@ async def test_release_budget_reservation_on_cancel_swallows_a_second_cancellati @pytest.mark.asyncio -async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_after_every_retry_fails(): - # If both the reconcile retries and the invalidate fallback fail (e.g. a persistent +async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_after_reconcile_fails(): + # If both the reconcile and the invalidate fallback fail (e.g. a persistent # outage), there is nothing left to try: the failure must be logged and swallowed, # not propagated, and the reservation still ends up finalized so it is not reprocessed. reservation = { @@ -3207,17 +3218,16 @@ async def test_release_budget_reservation_on_cancel_swallows_invalidate_failure_ @pytest.mark.asyncio -async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_persistently_fails( +async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_fails( spend_counter_state, ): """ - Regression for #30460 Path 1: if reconcile keeps failing on the cancel path - (e.g. a persistent Redis outage) across every retry, the pre-charge must - not be left stuck in the counter with nothing left to correct it. This - mirrors what release_or_invalidate_budget_reservation already does on the - non-cancel release path: fall back to invalidate_budget_reservation_counters - so the next read reseeds from the DB instead of enforcing the stale - reservation forever. + Regression for #30460 Path 1: if reconcile fails on the cancel path (e.g. a + Redis outage), the pre-charge must not be left stuck in the counter with + nothing left to correct it. This mirrors what + release_or_invalidate_budget_reservation already does on the non-cancel + release path: drop the counter so the next read reseeds from the DB instead + of enforcing the stale reservation until the TTL. """ counter_cache, _key_cache = spend_counter_state counter_key = "spend:key:key-cancel-redis-down" @@ -3234,40 +3244,30 @@ async def test_release_budget_reservation_on_cancel_invalidates_counter_when_rec await release_budget_reservation_on_cancel(reservation) + assert counter_key not in counter_cache.redis_cache.store assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None assert reservation["finalized"] is True - # Only the first retry reaches the pipeline: its own failure already - # invalidates the counter, so later retries take the (here also - # unavailable, with no real DB) reseed path instead -- both exhausted - # before falling back to invalidate_budget_reservation_counters. assert counter_cache.redis_cache.pipeline_calls == 1 @pytest.mark.asyncio -async def test_release_budget_reservation_on_cancel_retries_before_invalidating( +async def test_release_budget_reservation_on_cancel_does_not_refund_twice_when_reconcile_times_out_after_applying( spend_counter_state, ): """ - Regression for the cancellation-fallback race (veria-ai finding on - budget_reservation.py): invalidate_budget_reservation_counters deletes the - whole aggregate counter unconditionally, which would also erase a - concurrent request's reservation or recorded spend sharing that same - key/user/team counter, not just this reservation's own contribution. A - reconcile failure that clears up on retry must settle through the - ordinary reconcile machinery instead of falling back to that unconditional - delete: here, the first attempt's own failure already invalidates the - counter (existing increment_spend_counters_pipeline cleanup), so the retry - settles it by reseeding the DB floor and adding this reservation's actual - cost, the same recovery path an expired counter already takes elsewhere in - this file, rather than being deleted a second time with nothing settled. + Regression for the duplicate-refund race (veria-ai finding on + budget_reservation.py): the reconcile's INCRBYFLOAT lands in Redis but the + pipeline response times out, and the counter delete meant to clean that up + fails too, so the counter still holds the already-refunded value. Re-applying + the reconcile would subtract this reservation's refund a second time and eat + into a concurrent request's spend sharing the counter. The refund must land + exactly once. """ - import litellm.proxy.proxy_server as ps - counter_cache, _key_cache = spend_counter_state - counter_key = "spend:key:key-cancel-retry-recovers" - counter_cache.redis_cache = _FlakyPipelineRedisCache(fail_first_n=1) + counter_key = "spend:key:key-cancel-timeout-after-apply" + counter_cache.redis_cache = _FlakyPipelineRedisCache(fail_first_n=99, apply_before_failing=True, fail_deletes=True) + # 2.0 of a concurrent request's spend plus this request's 3.0 reservation counter_cache.redis_cache.store[counter_key] = 5.0 - counter_cache.in_memory_cache.set_cache(key=counter_key, value=5.0) reservation = { "reserved_cost": 3.0, @@ -3276,13 +3276,12 @@ async def test_release_budget_reservation_on_cancel_retries_before_invalidating( "entries": [{"counter_key": counter_key, "reserved_cost": 3.0, "applied_adjustment": 0.0}], } - with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for - ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.2) - ): - await release_budget_reservation_on_cancel(reservation) + await release_budget_reservation_on_cancel(reservation) + await release_budget_reservation_on_cancel(reservation) + # 5.0 - (3.0 - 0.5), applied once; a second application would leave 0.0 + assert counter_cache.redis_cache.store[counter_key] == pytest.approx(2.5) assert counter_cache.redis_cache.pipeline_calls == 1 - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.7) assert reservation["finalized"] is True