From fba32eea206f42ed119d0e38f121f13a3535cba8 Mon Sep 17 00:00:00 2001 From: Abhyuday Date: Sat, 12 Sep 2026 23:53:48 -0400 Subject: [PATCH] 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