From 7c2e100ea660b65568bf7a64a71a3762696d04ae Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:38:22 +0000 Subject: [PATCH] fix(router): drain increment queue before awaiting the redis write A second sync entering _push_in_memory_increments_to_redis while the first was still writing could snapshot the same undrained queue and resubmit it, inflating spend; the competing prefix removals could also drop operations the remover never submitted. The queue is now swapped out synchronously before the first await, so a racing push sees an empty queue, and a failed write puts the operations back in front of anything queued mid-flight. Also adds a regression test for the double-apply and drops the comment block above the Lua script in favour of documenting the KEYS/ARGV contract on the method. --- litellm/router_strategy/budget_limiter.py | 21 +++--- .../router_strategy/test_budget_limiter.py | 65 +++++++++++++++++++ 2 files changed, 78 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 17d3b00d43f..f27d46af66e 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -43,9 +43,6 @@ from litellm.types.utils import GenericBudgetConfigType, StandardLoggingPayload DEFAULT_REDIS_SYNC_INTERVAL = 1 -# KEYS[1] = budget start time key, KEYS[2] = spend key -# ARGV[1] = current time, ARGV[2] = response cost, ARGV[3] = window length in seconds -# Returns {window start time, "1" if this caller opened the window else "0"} BUDGET_WINDOW_CLAIM_SCRIPT = """ local budget_start = redis.call('GET', KEYS[1]) local current_time = tonumber(ARGV[1]) @@ -437,6 +434,10 @@ class RouterBudgetLimiting(CustomLogger): expired start time resets the counter, everyone else racing across the same boundary increments it. + The script takes KEYS = [start_time_key, spend_key] and ARGV = + [current_time, response_cost, window_length_seconds], and returns + [window_start, "1" if this caller opened the window else "0"]. + Returns the window start time, or None when Redis is unavailable, so the caller can fall back to the local path. """ @@ -619,15 +620,19 @@ class RouterBudgetLimiting(CustomLogger): "Pushing Redis Increment Pipeline for queue: %s", self.redis_increment_operation_queue, ) - pending_increments = list(self.redis_increment_operation_queue) + pending_increments = self.redis_increment_operation_queue if len(pending_increments) == 0: return - await self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=pending_increments, - ) + self.redis_increment_operation_queue = [] - self.redis_increment_operation_queue = self.redis_increment_operation_queue[len(pending_increments) :] + try: + await self.dual_cache.redis_cache.async_increment_pipeline( + increment_list=pending_increments, + ) + except Exception: + self.redis_increment_operation_queue = pending_increments + self.redis_increment_operation_queue + raise except Exception as e: verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index 0d6f6176deb..88eccc288f3 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -242,6 +242,71 @@ async def test_push_in_memory_increments_waits_for_redis_and_keeps_new_increment assert [op["increment_value"] for op in limiter.redis_increment_operation_queue] == [0.07] +@pytest.mark.asyncio +async def test_concurrent_pushes_do_not_double_apply_increments(disable_budget_sync): + """ + A second sync racing while the first push is still writing must not resubmit + the same queued increments, otherwise the spend counter is inflated. + """ + redis_cache = FakeAtomicRedisCache() + limiter = _budget_limiter(redis_cache=redis_cache) + spend_key = "provider_spend:synthetic:1h" + + pipeline_started = asyncio.Event() + release_pipeline = asyncio.Event() + original_increment_pipeline = redis_cache.async_increment_pipeline + + async def blocking_increment_pipeline(increment_list, **kwargs): + pipeline_started.set() + await release_pipeline.wait() + return await original_increment_pipeline(increment_list, **kwargs) + + redis_cache.async_increment_pipeline = blocking_increment_pipeline # type: ignore[method-assign] + + await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.05, ttl=3600) + + first_push = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + await pipeline_started.wait() + second_push = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + await asyncio.sleep(0) + release_pipeline.set() + await asyncio.gather(first_push, second_push) + + assert len(redis_cache.increment_pipeline_calls) == 1 + assert [op["increment_value"] for op in redis_cache.increment_pipeline_calls[0]] == [0.05] + assert float(redis_cache.store[spend_key]) == pytest.approx(0.05) + assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_budget_window_claim_decodes_bytes_returned_by_redis(disable_budget_sync): + """ + redis-py can hand back bytes; the claim result must still be decoded and the + winning reset mirrored into the in-memory cache. + """ + + class BytesRedisCache(FakeAtomicRedisCache): + def async_register_script(self, script: str): + async def run_script(keys, args, client=None): + return [str(args[0]).encode("utf-8"), b"1"] + + return run_script + + limiter = _budget_limiter(redis_cache=BytesRedisCache()) + spend_key = "provider_spend:synthetic:1h" + + start_time = await limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key="provider_budget_start_time:synthetic", + current_time=1000.0, + response_cost=0.05, + ttl_seconds=3600, + ) + + assert start_time == 1000.0 + assert float(limiter.dual_cache.in_memory_cache.get_cache(spend_key)) == pytest.approx(0.05) + + @pytest.mark.asyncio async def test_push_in_memory_increments_retains_queue_when_redis_write_fails(disable_budget_sync): redis_cache = FakeAtomicRedisCache()