From a35cf6bcb0af701410f425cd9eeee04252e6b767 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 18:07:39 -0800 Subject: [PATCH] fix(proxy): raise a refused TTL command after a late set as main's add does When Redis applied SET NX GET but refused the EXPIRE after it, the late write fell back to main's add, which added the spend on top of the value the set had just written A refusal that redis-py names after the first pipeline command now raises, as main's add raises on the same EXPIRE, and the caller deletes the counter as on main --- litellm/caching/redis_cache.py | 16 ++++++++++++--- tests/unit/caching/test_redis_cache.py | 27 +++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f6223fb8d94..6ef99f5ffd1 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -143,6 +143,7 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( _LUA_COUNT: Final = TypeAdapter(int) _PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") +_PIPELINE_COMMAND_NUMBER: Final = re.compile(r"Command # ([0-9]+) \(") _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) @@ -153,6 +154,13 @@ def _as_incrbyfloat_writes(value: float) -> str: return "0" if text == "-0" else text +def _failed_after_the_first_command(error: Exception) -> bool: + """Whether redis-py names a pipeline command after the first as the one that failed, so the first one applied.""" + message: Final = error.args[0] if error.args else None + number: Final = _PIPELINE_COMMAND_NUMBER.match(message) if isinstance(message, str) else None + return number is not None and int(number.group(1)) > 1 + + def _plain_float(reply: object) -> float | None: """The finite float a stored value spells in plain decimal or exponent form, or None for anything else.""" text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply @@ -1494,8 +1502,9 @@ class RedisCache(BaseCache): async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: """SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or - None when the server refused a command rather than failed to answer, or the key holds a value this does not - read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing.""" + None when the server refused SET NX GET rather than failed to answer, or the key holds a value this does not + read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing. A refused TTL command + raises, as it does in main's add.""" from redis.exceptions import ClusterError, RedisError from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError @@ -1510,7 +1519,8 @@ class RedisCache(BaseCache): pipe.expire(key, ttl) replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics except RedisError as e: - if isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)): + failed_to_answer: Final = isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)) + if failed_to_answer or _failed_after_the_first_command(e): raise return None if replies[0] is None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2a2ac5f7063..e673f16e5a6 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1727,8 +1727,6 @@ class _KeyspacePipeline(_RaisingPipeline): self.commands.append(("SETNX", name, value)) def _write(self, command: str) -> None: - if command in self.keyspace.refused: - raise NoPermissionError(f"this user has no permissions to run the '{command.lower()}' command") if self.writes_left is not None: if self.writes_left == 0: raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.") @@ -1737,6 +1735,8 @@ class _KeyspacePipeline(_RaisingPipeline): def _run(self, command: tuple[object, ...]) -> object: name: Final = str(command[1]) held: Final = self.keyspace.values.get(name) + if command[0] in self.keyspace.refused: + raise NoPermissionError(f"this user has no permissions to run the '{str(command[0]).lower()}' command") if command[0] == "EXPIRE": return held is not None self._write(str(command[0])) @@ -1756,10 +1756,12 @@ class _KeyspacePipeline(_RaisingPipeline): if self.lost: raise RedisConnectionError("Connection closed by server.") replies: Final[list[object]] = [] - for command in self.commands: + for number, command in enumerate(self.commands, start=1): try: replies.append(self._run(command)) except RedisError as e: + words: str = " ".join(str(part) for part in command) + e.args = (f"Command # {number} ({words}) of pipeline caused error: {e.args}",) replies.append(e) first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) if first_error is not None: @@ -1801,6 +1803,25 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redi added.assert_called_once_with() +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_raises_a_refused_ttl_command_as_mains_add_does_without_adding_again( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("EXPIRE",))) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(NoPermissionError): + await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert keyspace.values == {"ns:spend:key:k": "6"} + assert keyspace.pipelines == 1 + added.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"]) async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add(