From 78309d9c0b13b3816866962baf08625debda7786 Mon Sep 17 00:00:00 2001 From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Date: Fri, 11 Sep 2026 10:55:13 +0000 Subject: [PATCH 1/2] 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 --- litellm/caching/redis_cache.py | 25 +++---- .../test_litellm/caching/test_redis_cache.py | 65 +++++++++++++++++++ 2 files changed, 79 insertions(+), 11 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2c36995c4f8..a01a099bcc9 100644 --- a/litellm/caching/redis_cache.py +++ b/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() diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcae33b976e..d27ab6fe393 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/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") From c12a3dbe52938435c5104b6576aedbc7f3051eee Mon Sep 17 00:00:00 2001 From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:50:39 +0000 Subject: [PATCH 2/2] fix(caching): keep the exact TTL semantics in the atomic increment The Lua script skipped EXPIRE for any non-positive TTL, while the old code only skipped it when get_ttl returned None and otherwise passed the value through, so EXPIRE 0 still deleted the key. The TTL now travels as an empty string for None and as the literal value otherwise, and the script only branches on that emptiness Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng --- litellm/caching/redis_cache.py | 8 +++---- .../test_litellm/caching/test_redis_cache.py | 22 +++++++++++++++++-- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index a01a099bcc9..1dd3711e36e 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -100,9 +100,8 @@ _INCREMENT_WITH_FLOOR_LUA: Final = ( _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 " + "if ARGV[2] ~= '' and (ARGV[3] == '1' or redis.call('TTL', KEYS[1]) == -1) then " + "redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2])) end " "return value" ) @@ -1256,8 +1255,9 @@ class RedisCache(BaseCache): _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: + ttl_arg: Final = "" if _used_ttl is None else str(_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" + _INCREMENT_WITH_TTL_LUA, 1, key, value, ttl_arg, "1" if refresh_ttl else "0" ) result: Final = _LUA_FLOAT.validate_python(raw_value) diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index d27ab6fe393..5aaa75177ab 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1253,7 +1253,7 @@ async def test_redis_cache_async_increment_arms_ttl_in_the_same_command( 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") + assert (numkeys, key, amount, ttl, refresh) == (1, expected_key, 0.25, "30", "0") @pytest.mark.asyncio @@ -1266,4 +1266,22 @@ async def test_redis_cache_async_increment_refresh_ttl_sends_refresh_flag(monkey assert result == 2.5 assert spy.commands == ["eval"] - assert spy.eval_calls[0][3:] == (0.5, 60, "1") + assert spy.eval_calls[0][3:] == (0.5, "60", "1") + + +@pytest.mark.parametrize( + ("ttl", "default_ttl", "expected_ttl_arg"), [(None, None, ""), (0, None, "0"), (None, 15, "15")] +) +@pytest.mark.asyncio +async def test_redis_cache_async_increment_forwards_ttl_exactly( + ttl, default_ttl, expected_ttl_arg, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + spy = _SpyRedisCommands(eval_result=b"0.75") + redis_cache = _SpyRedisCache(spy) + redis_cache.default_ttl = default_ttl + + result = await redis_cache.async_increment(key="spend:key:abc", value=0.75, ttl=ttl) + + assert result == 0.75 + assert spy.eval_calls[0][4] == expected_ttl_arg