fix(caching): set the counter TTL atomically in async_increment

RedisCache.async_increment sent INCRBYFLOAT and then EXPIRE (or TTL plus
EXPIRE) as separate awaits. A task cancelled between the two, which happens
whenever the client disconnects mid request, left the counter without any
expiry, so a spend or rate limit counter could live forever in Redis

The increment and the TTL decision now run in one Lua call, following the
pattern async_increment_with_floor and async_set_max already use in this
file. The refresh_ttl semantics are unchanged: the flag re-arms the TTL on
every increment, otherwise only a key with no TTL gets one

Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com>
Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>

litellm/caching/redis_cache.py

@@ -98,7 +98,16 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
     "return count"
 )
 
+_INCREMENT_WITH_TTL_LUA: Final = (
+    "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]) "
+    "local ttl = tonumber(ARGV[2]) "
+    "if ttl > 0 and (ARGV[3] == '1' or redis.call('TTL', KEYS[1]) == -1) then "
+    "redis.call('EXPIRE', KEYS[1], ttl) end "
+    "return value"
+)
+
 _LUA_COUNT: Final = TypeAdapter(int)
+_LUA_FLOAT: Final = TypeAdapter(float)
 _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
 
 
@@ -1242,21 +1251,15 @@ class RedisCache(BaseCache):
         parent_otel_span: Span | None = None,
         refresh_ttl: bool = False,
     ) -> float:
-        from redis.asyncio import Redis
-
-        _redis_client: Final[Redis] = self.init_async_client()
+        _redis_client: Final = self._async_commands()
         start_time: Final = time.time()
         _used_ttl: Final = self.get_ttl(ttl=ttl)
         key = self.check_and_fix_namespace(key=key)
         try:
-            result: Final = await _redis_client.incrbyfloat(name=key, amount=value)
-            if _used_ttl is not None:
-                if refresh_ttl:
-                    await _redis_client.expire(key, _used_ttl)
-                else:
-                    current_ttl: Final = await _redis_client.ttl(key)
-                    if current_ttl == -1:
-                        await _redis_client.expire(key, _used_ttl)
+            raw_value: Final = await _redis_client.eval(
+                _INCREMENT_WITH_TTL_LUA, 1, key, value, _used_ttl or 0, "1" if refresh_ttl else "0"
+            )
+            result: Final = _LUA_FLOAT.validate_python(raw_value)
 
             ## LOGGING ##
             end_time = time.time()

tests/test_litellm/caching/test_redis_cache.py

@@ -1202,3 +1202,68 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new
     new_probe_release.set()
     assert await new_probe == "new probe"
     assert breaker._state == breaker.CLOSED
+
+
+class _SpyRedisCommands:
+    def __init__(self, eval_result: object) -> None:
+        self.commands: list[str] = []
+        self.eval_calls: list[tuple[object, ...]] = []
+        self._eval_result = eval_result
+
+    async def eval(self, script: str, numkeys: int, *keys_and_args: object) -> object:
+        self.commands.append("eval")
+        self.eval_calls.append((script, numkeys, *keys_and_args))
+        return self._eval_result
+
+    async def incrbyfloat(self, name: str, amount: float) -> float:
+        self.commands.append("incrbyfloat")
+        return amount
+
+    async def ttl(self, name: str) -> int:
+        self.commands.append("ttl")
+        return -1
+
+    async def expire(self, name: str, time: int) -> bool:
+        self.commands.append("expire")
+        return True
+
+
+class _SpyRedisCache(RedisCache):
+    def __init__(self, spy: _SpyRedisCommands, **kwargs: object) -> None:
+        super().__init__(**kwargs)
+        self.spy = spy
+
+    def init_async_client(self, *args: object, **kwargs: object) -> object:
+        return self.spy
+
+
+@pytest.mark.parametrize(("namespace", "expected_key"), [(None, "spend:key:abc"), ("ns", "ns:spend:key:abc")])
+@pytest.mark.asyncio
+async def test_redis_cache_async_increment_arms_ttl_in_the_same_command(
+    namespace, expected_key, monkeypatch, redis_no_ping
+):
+    """The increment and its TTL reach Redis as one server-side step."""
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    spy = _SpyRedisCommands(eval_result=b"1.25")
+    redis_cache = _SpyRedisCache(spy, namespace=namespace)
+
+    result = await redis_cache.async_increment(key="spend:key:abc", value=0.25, ttl=30)
+
+    assert result == 1.25
+    assert spy.commands == ["eval"]
+    script, numkeys, key, amount, ttl, refresh = spy.eval_calls[0]
+    assert "INCRBYFLOAT" in script and "EXPIRE" in script and "TTL" in script
+    assert (numkeys, key, amount, ttl, refresh) == (1, expected_key, 0.25, 30, "0")
+
+
+@pytest.mark.asyncio
+async def test_redis_cache_async_increment_refresh_ttl_sends_refresh_flag(monkeypatch, redis_no_ping):
+    monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
+    spy = _SpyRedisCommands(eval_result="2.5")
+    redis_cache = _SpyRedisCache(spy)
+
+    result = await redis_cache.async_increment(key="spend:team:t1", value=0.5, refresh_ttl=True)
+
+    assert result == 2.5
+    assert spy.commands == ["eval"]
+    assert spy.eval_calls[0][3:] == (0.5, 60, "1")