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

@@ -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)}")

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()