From 5b7f7125c2c198d417ea08875aacec69c02ec7e3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:23:51 +0000 Subject: [PATCH] feat(redis): add native seconds transport expiry commands --- litellm/redis/transport/client.py | 6 +- litellm/redis/transport/contract.py | 12 ++ litellm/redis/transport/redis_py.py | 49 +++++++ tests/unit/redis/stores/test_pod_lock.py | 6 + tests/unit/redis/transport/test_redis_py.py | 142 +++++++++++++++++++- 5 files changed, 213 insertions(+), 2 deletions(-) diff --git a/litellm/redis/transport/client.py b/litellm/redis/transport/client.py index 23f455c7ed9..958c7971597 100644 --- a/litellm/redis/transport/client.py +++ b/litellm/redis/transport/client.py @@ -37,7 +37,9 @@ class AsyncRedisCommands(Protocol): def mget(self, keys: Sequence[str]) -> Awaitable[object]: ... - def set(self, name: str, value: Arg, *, nx: bool, px: int | None) -> Awaitable[object]: ... + def set( + self, name: str, value: Arg, *, nx: bool, px: int | None = None, ex: timedelta | None = None + ) -> Awaitable[object]: ... def delete(self, *names: str) -> Awaitable[object]: ... @@ -47,6 +49,8 @@ class AsyncRedisCommands(Protocol): def pexpire(self, name: str, time: int) -> Awaitable[object]: ... + def expire(self, name: str, time: int) -> Awaitable[object]: ... + def ttl(self, name: str) -> Awaitable[object]: ... def rpush(self, name: str, *values: Arg) -> Awaitable[object]: ... diff --git a/litellm/redis/transport/contract.py b/litellm/redis/transport/contract.py index 83c4e2b0664..983da44a89f 100644 --- a/litellm/redis/transport/contract.py +++ b/litellm/redis/transport/contract.py @@ -53,12 +53,18 @@ class RedisCommandBatch(Protocol): def delete(self, *keys: str) -> BatchReply[int | RedisFailure]: ... + def set_seconds( + self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False + ) -> BatchReply[bool | RedisFailure]: ... + def incrby(self, key: str, amount: int) -> BatchReply[int | RedisFailure]: ... def incrbyfloat(self, key: str, amount: float) -> BatchReply[float | RedisFailure]: ... def expire(self, key: str, ttl: timedelta) -> BatchReply[bool | RedisFailure]: ... + def expire_seconds(self, key: str, ttl: timedelta) -> BatchReply[bool | RedisFailure]: ... + def ttl(self, key: str) -> BatchReply[int | RedisFailure]: ... def rpush(self, key: str, *values: Arg) -> BatchReply[int | RedisFailure]: ... @@ -89,12 +95,18 @@ class RedisTransport(Protocol): async def delete(self, *keys: str) -> int | RedisFailure: ... + async def set_seconds( + self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False + ) -> bool | RedisFailure: ... + async def incrby(self, key: str, amount: int) -> int | RedisFailure: ... async def incrbyfloat(self, key: str, amount: float) -> float | RedisFailure: ... async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure: ... + async def expire_seconds(self, key: str, ttl: timedelta) -> bool | RedisFailure: ... + async def ttl(self, key: str) -> int | RedisFailure: ... async def rpush(self, key: str, *values: Arg) -> int | RedisFailure: ... diff --git a/litellm/redis/transport/redis_py.py b/litellm/redis/transport/redis_py.py index 227996ed6ff..1729b0f1ca2 100644 --- a/litellm/redis/transport/redis_py.py +++ b/litellm/redis/transport/redis_py.py @@ -61,6 +61,13 @@ class RedisPyTransport: px: Final = None if ttl is None else ttl_ms(ttl) return _bool_reply(await self._run("set", lambda: self.client().set(self._key(key), value, px=px, nx=nx))) + async def set_seconds( + self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False + ) -> bool | RedisFailure: + return _set_reply( + await self._run("set_seconds", lambda: self.client().set(self._key(key), value, ex=ttl, nx=nx)) + ) + async def delete(self, *keys: str) -> int | RedisFailure: namespaced: Final = tuple(self._key(k) for k in keys) return _int_reply(await self._run("delete", lambda: self.client().delete(*namespaced))) @@ -74,6 +81,11 @@ class RedisPyTransport: async def expire(self, key: str, ttl: timedelta) -> bool | RedisFailure: return _bool_reply(await self._run("expire", lambda: self.client().pexpire(self._key(key), ttl_ms(ttl)))) + async def expire_seconds(self, key: str, ttl: timedelta) -> bool | RedisFailure: + return _expire_reply( + await self._run("expire_seconds", lambda: self.client().expire(self._key(key), int(ttl.total_seconds()))) + ) + async def ttl(self, key: str) -> int | RedisFailure: return _int_reply(await self._run("ttl", lambda: self.client().ttl(self._key(key)))) @@ -228,6 +240,17 @@ class RedisPyCommandBatch: lambda: self.transport.delete(*keys), ) + def set_seconds( + self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False + ) -> BatchResult[bool | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + self.engine.note_present(key) + return self._declare( + lambda pipe: _one(pipe.set(name, value, ex=ttl, nx=nx)), + lambda replies: _set_reply(replies[0]), + lambda: self.transport.set_seconds(key, value, ttl=ttl, nx=nx), + ) + def incrby(self, key: str, amount: int) -> BatchResult[int | RedisFailure]: name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own self.engine.note_present(key) @@ -260,6 +283,14 @@ class RedisPyCommandBatch: lambda pipe: _one(pipe.ttl(name)), lambda replies: _int_reply(replies[0]), lambda: self.transport.ttl(key) ) + def expire_seconds(self, key: str, ttl: timedelta) -> BatchResult[bool | RedisFailure]: + name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own + return self._declare( + lambda pipe: _one(pipe.expire(name, ttl)), + lambda replies: _expire_reply(replies[0]), + lambda: self.transport.expire_seconds(key, ttl), + ) + def rpush(self, key: str, *values: Arg) -> BatchResult[int | RedisFailure]: name: Final = self.transport._key(key) # pyright: ignore[reportPrivateUsage] # the batch is the transport's own self.engine.note_present(key) @@ -491,6 +522,24 @@ def _bool_reply(reply: object) -> bool | RedisFailure: return reply is not None and reply is not False and reply != 0 +def _expire_reply(reply: object) -> bool | RedisFailure: + if isinstance(reply, (RedisFailure, bool)): + return reply + if type(reply) is int and reply in (0, 1): + return reply == 1 + return _decode_failure(reply) + + +def _set_reply(reply: object) -> bool | RedisFailure: + if isinstance(reply, RedisFailure): + return reply + if reply is None or reply is False: + return False + if reply is True or reply in (b"OK", "OK"): + return True + return _decode_failure(reply) + + def _int_reply(reply: object) -> int | RedisFailure: if isinstance(reply, RedisFailure): return reply diff --git a/tests/unit/redis/stores/test_pod_lock.py b/tests/unit/redis/stores/test_pod_lock.py index 24e9c298996..f538d7694fb 100644 --- a/tests/unit/redis/stores/test_pod_lock.py +++ b/tests/unit/redis/stores/test_pod_lock.py @@ -17,6 +17,9 @@ class ExpiredLockTransport: async def set(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> bool | RedisFailure: return not nx + async def set_seconds(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> NoReturn: + raise AssertionError("unexpected set_seconds") + async def mget(self, keys: Sequence[str]) -> tuple[bytes | None, ...] | RedisFailure: raise AssertionError("unexpected mget") @@ -35,6 +38,9 @@ class ExpiredLockTransport: async def expire(self, key: str, ttl: timedelta) -> NoReturn: raise AssertionError("unexpected expire") + async def expire_seconds(self, key: str, ttl: timedelta) -> NoReturn: + raise AssertionError("unexpected expire_seconds") + async def ttl(self, key: str) -> NoReturn: raise AssertionError("unexpected ttl") diff --git a/tests/unit/redis/transport/test_redis_py.py b/tests/unit/redis/transport/test_redis_py.py index 7cd0de37958..76d77a22cba 100644 --- a/tests/unit/redis/transport/test_redis_py.py +++ b/tests/unit/redis/transport/test_redis_py.py @@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass, replace from datetime import datetime, timedelta from typing import Final, Literal -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import fakeredis import pytest +from redis.asyncio.client import Pipeline, Redis +from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import ResponseError import litellm @@ -389,6 +391,8 @@ _WRONG_TYPE_COMMANDS: Final[tuple[_Command, ...]] = ( _OUTAGE_COMMANDS: Final[tuple[_Command, ...]] = ( lambda t: t.incrby("k", 1), lambda t: t.expire("k", timedelta(seconds=1)), + lambda t: t.expire_seconds("k", timedelta(seconds=1)), + lambda t: t.set_seconds("k", "v", ttl=timedelta(seconds=1)), lambda t: t.ttl("k"), lambda t: t.lpop("k"), lambda t: t.scan("k", limit=1), @@ -428,6 +432,139 @@ async def test_expire_sets_the_ttl_in_milliseconds_and_ttl_reads_it_back( assert await transport.ttl("absent") == -2 +def _typed_pipeline(client: "Redis[bytes]") -> "Pipeline[bytes]": + return client.pipeline(transaction=False) + + +def _seconds_commands(transport: RedisPyTransport, mode: str) -> RedisPyTransport | RedisPyCommandBatch: + if mode == "direct": + return transport + if mode == "alone": + backend: Final = replace(transport.batch().engine.backend, runs_alone=True) + return RedisPyCommandBatch(transport, PipelineBatch(backend, name="seconds")) + return transport.batch(transaction=mode == "transaction") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone")) +@pytest.mark.parametrize("seconds", (2, 1.9, 0.5, 0, -0.5, -1)) +async def test_native_expire_truncates_seconds_and_deletes_nonpositive_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, seconds: float +) -> None: + transport: Final = fake_redis_transport(raw, failure_threshold=1, namespace="ns") + commands: Final = _seconds_commands(transport, mode) + with patch("time.time", return_value=1700000000.0): + await raw.set("ns:k", "v") + assert await commands.expire_seconds("ns:k", timedelta(seconds=seconds)) is True + assert await commands.expire_seconds("absent", timedelta(seconds=seconds)) is False + if seconds >= 1: + assert await raw.pttl("ns:k") == int(seconds) * 1000 + assert await raw.get("ns:k") == b"v" + else: + with patch("time.time", return_value=1700000000.001): + assert await raw.get("ns:k") is None + assert await raw.exists("ns:ns:k") == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone")) +async def test_native_set_preserves_seconds_expiry_persistence_and_nx( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str +) -> None: + commands: Final = _seconds_commands(fake_redis_transport(raw, failure_threshold=1, namespace="ns"), mode) + with patch("time.time", return_value=1700000000.0): + assert await commands.set_seconds("k", b"1.25", ttl=timedelta(seconds=1.9), nx=True) is True + assert await commands.set_seconds("ns:k", "rejected", ttl=None, nx=True) is False + assert await raw.get("ns:k") == b"1.25" + assert await raw.pttl("ns:k") == 1000 + assert await commands.set_seconds("ns:k", "persist", ttl=None) is True + assert await raw.get("ns:k") == b"persist" + assert await raw.pttl("ns:k") == -1 + assert await raw.keys("*") == [b"ns:k"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("direct", "batch", "transaction", "alone")) +@pytest.mark.parametrize("seconds", (0.5, 0, -0.5, -1)) +async def test_native_set_rejection_keeps_value_and_connection_and_emits_one_event( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, seconds: float +) -> None: + events: Final = _RecordingEvents() + transport: Final = fake_redis_transport(raw, failure_threshold=1, events=events, namespace="ns") + commands: Final = _seconds_commands(transport, mode) + await raw.set("ns:k", "before") + with service_target("seconds_test"): + rejected: Final = await commands.set_seconds("k", "after", ttl=timedelta(seconds=seconds)) + assert isinstance(rejected, RedisFailure) and rejected.kind == "server" + assert await raw.get("ns:k") == b"before" + assert len(events.seen) == 1 + assert events.seen[0].outcome == ("failure" if mode in ("direct", "alone") else "success") + assert events.seen[0].target == "seconds_test" + assert await commands.expire_seconds("k", timedelta(seconds=30)) is True + await _settle() + assert len(events.seen) == 2 + assert events.seen[1].outcome == "success" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("direct", "batch")) +async def test_native_seconds_outage_opens_breaker_without_duplicate_events( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str +) -> None: + events: Final = _RecordingEvents() + commands: Final = _seconds_commands(fake_redis_transport(raw, failure_threshold=1, events=events), mode) + server.connected = False + failed: Final = await commands.expire_seconds("k", timedelta(seconds=30)) + assert isinstance(failed, RedisFailure) and failed.kind == "transport" + server.connected = True + refused: Final = await commands.set_seconds("k", "v", ttl=None) + assert isinstance(refused, RedisFailure) and refused.kind == "breaker_open" + await _settle() + assert tuple(event.outcome for event in events.seen) == ( + ("failure",) if mode == "direct" else ("failure", "failure") + ) + assert await raw.get("k") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("direct", "batch")) +@pytest.mark.parametrize("reply", (b"malformed", 2, 1.0, (True,), [True], {"unexpected": True})) +async def test_native_seconds_rejects_malformed_command_replies( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, mode: str, reply: object +) -> None: + transport: Final = fake_redis_transport(raw, failure_threshold=1) + pipeline: Final = _typed_pipeline(raw) + with ( + patch.object(raw, "expire", new=AsyncMock(return_value=reply)), + patch.object(raw, "set", new=AsyncMock(return_value=reply)), + patch.object(raw, "pipeline", return_value=pipeline), + patch.object(pipeline, "execute", return_value=[reply]), + ): + commands: Final = _seconds_commands(transport, mode) + expired: Final = await commands.expire_seconds("k", timedelta(seconds=30)) + written: Final = await commands.set_seconds("k", "v", ttl=None) + assert isinstance(expired, RedisFailure) and expired.kind == "decode" + assert isinstance(written, RedisFailure) and written.kind == "decode" + + +@pytest.mark.asyncio +async def test_native_expire_lost_reply_is_unknown_applied_without_retry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport: Final = fake_redis_transport(raw, failure_threshold=1) + expire: Final = raw.expire + await raw.set("k", "v") + + async def lose_reply(name: str, time: int) -> bool: + await expire(name, time) + raise RedisConnectionError("reply lost after expiry applied") + + with patch("time.time", return_value=1700000000.0), patch.object(raw, "expire", side_effect=lose_reply): + result: Final = await transport.expire_seconds("k", timedelta(seconds=30)) + assert isinstance(result, RedisFailure) and result.kind == "transport" + assert await raw.pttl("k") == 30000 + + @pytest.mark.asyncio async def test_list_commands_push_trim_and_pop_from_the_head( server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis @@ -654,10 +791,13 @@ async def test_request_batch_commands_share_the_round_trip_of_redis_batch_caller old_read = legacy.mget(["legacy"]) new_read = commands.get("moved") missing = commands.get("absent") + written = commands.set_seconds("seconds", "v", ttl=None) + expired = commands.expire_seconds("seconds", timedelta(seconds=30)) assert await old_read == {"legacy": "old"} assert await new_read == b"new" assert await missing is None + assert await written is True and await expired is True assert legacy.flushes == 1 assert legacy.read_as_missing("absent") is True assert legacy.read_as_missing("moved") is False