diff --git a/tests/memory/test_redis_cluster_memory_service.py b/tests/memory/test_redis_cluster_memory_service.py new file mode 100644 index 00000000..0a76488f --- /dev/null +++ b/tests/memory/test_redis_cluster_memory_service.py @@ -0,0 +1,24 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Construction tests for RedisClusterMemoryService.""" + +from unittest.mock import MagicMock +from unittest.mock import patch + +from trpc_agent_sdk.memory import RedisClusterMemoryService + + +class TestRedisClusterMemoryService: + + @patch("trpc_agent_sdk.memory._redis_cluster_memory_service.RedisClusterStorage") + def test_uses_cluster_storage(self, storage_cls): + storage = MagicMock() + storage_cls.return_value = storage + + service = RedisClusterMemoryService(db_url="redis://seed:6379/0", is_async=True, max_connections=20) + + assert service._redis_storage is storage + storage_cls.assert_called_once_with(is_async=True, redis_url="redis://seed:6379/0", max_connections=20) diff --git a/tests/sessions/test_redis_cluster_session_service.py b/tests/sessions/test_redis_cluster_session_service.py new file mode 100644 index 00000000..7eb0db64 --- /dev/null +++ b/tests/sessions/test_redis_cluster_session_service.py @@ -0,0 +1,24 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Construction tests for RedisClusterSessionService.""" + +from unittest.mock import MagicMock +from unittest.mock import patch + +from trpc_agent_sdk.sessions import RedisClusterSessionService + + +class TestRedisClusterSessionService: + + @patch("trpc_agent_sdk.sessions._redis_cluster_session_service.RedisClusterStorage") + def test_uses_cluster_storage(self, storage_cls): + storage = MagicMock() + storage_cls.return_value = storage + + service = RedisClusterSessionService(db_url="redis://seed:6379/0", is_async=True, max_connections=20) + + assert service._redis_storage is storage + storage_cls.assert_called_once_with(is_async=True, redis_url="redis://seed:6379/0", max_connections=20) diff --git a/tests/storage/test_redis_cluster_storage.py b/tests/storage/test_redis_cluster_storage.py new file mode 100644 index 00000000..2312cae1 --- /dev/null +++ b/tests/storage/test_redis_cluster_storage.py @@ -0,0 +1,95 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Tests for the Redis Cluster storage adapter without a live cluster.""" + +from unittest.mock import MagicMock +from unittest.mock import patch + +from trpc_agent_sdk.storage import RedisClusterStorage +from trpc_agent_sdk.storage import RedisCommand +from trpc_agent_sdk.storage import RedisCondition + + +class _FakeClusterClient: + """Minimal async cluster client used to exercise the storage adapter.""" + + def __init__(self) -> None: + self.scan_calls: list[tuple[str, int]] = [] + self.closed = False + + def scan_iter(self, match: str, count: int): + self.scan_calls.append((match, count)) + + async def _iterate(): + # Duplicate keys are possible while a real cluster is resharding. + yield b"memory:app/user:one" + yield b"memory:app/user:two" + yield b"memory:app/user:one" + + return _iterate() + + async def type(self, key: str): + return "list" + + async def lrange(self, key: str, start: int, end: int): + return [f'{key}-event'] + + async def hgetall(self, key: str): + return {b"field": b"value"} + + async def close(self): + self.closed = True + + +class TestRedisClusterStorage: + + async def test_query_scans_all_cluster_keys_and_deduplicates(self): + storage = RedisClusterStorage(redis_url="redis://seed:6379/0", is_async=True) + client = _FakeClusterClient() + storage._redis_client = client + + async with storage.create_db_session() as conn: + results = await storage.query(conn, "memory:app/user:*", RedisCondition(limit=-1)) + + assert [key for key, _ in results] == ["memory:app/user:one", "memory:app/user:two"] + assert client.scan_calls == [("memory:app/user:*", storage._SCAN_COUNT)] + + async def test_keys_command_uses_scan_not_node_local_keys(self): + storage = RedisClusterStorage(redis_url="redis://seed:6379/0", is_async=True) + client = _FakeClusterClient() + + keys = await storage.execute_command(client, RedisCommand(method="keys", args=("session:app:*", ))) + + assert keys == ["memory:app/user:one", "memory:app/user:two"] + assert client.scan_calls == [("session:app:*", storage._SCAN_COUNT)] + + async def test_hgetall_normalizes_raw_redis_bytes(self): + storage = RedisClusterStorage(redis_url="redis://seed:6379/0", is_async=True, decode_responses=False) + client = _FakeClusterClient() + + result = await storage.execute_command(client, RedisCommand(method="hgetall", args=("state", ))) + + assert result == {"field": "value"} + + async def test_async_client_is_created_from_seed_url(self): + storage = RedisClusterStorage(redis_url="redis://seed:6379/0", is_async=True, max_connections=20) + client = MagicMock() + with patch("trpc_agent_sdk.storage._redis_cluster.AsyncRedisCluster") as client_cls: + client_cls.from_url.return_value = client + await storage.create_redis_engine() + + client_cls.from_url.assert_called_once_with("redis://seed:6379/0", decode_responses=True, max_connections=20) + assert storage._redis_client is client + + async def test_close_releases_cluster_client(self): + storage = RedisClusterStorage(redis_url="redis://seed:6379/0", is_async=True) + client = _FakeClusterClient() + storage._redis_client = client + + await storage.close() + + assert client.closed is True + assert storage._redis_client is None diff --git a/trpc_agent_sdk/memory/__init__.py b/trpc_agent_sdk/memory/__init__.py index cb78326a..9e6fd38b 100644 --- a/trpc_agent_sdk/memory/__init__.py +++ b/trpc_agent_sdk/memory/__init__.py @@ -16,6 +16,7 @@ from ._in_memory_memory_service import EventTtl from ._in_memory_memory_service import InMemoryMemoryService from ._redis_memory_service import RedisMemoryService +from ._redis_cluster_memory_service import RedisClusterMemoryService from ._sql_memory_service import MemStorageData from ._sql_memory_service import MemStorageEvent from ._sql_memory_service import SqlMemoryService @@ -28,6 +29,7 @@ "EventTtl", "InMemoryMemoryService", "RedisMemoryService", + "RedisClusterMemoryService", "MemStorageData", "MemStorageEvent", "SqlMemoryService", diff --git a/trpc_agent_sdk/memory/_redis_cluster_memory_service.py b/trpc_agent_sdk/memory/_redis_cluster_memory_service.py new file mode 100644 index 00000000..00c4682c --- /dev/null +++ b/trpc_agent_sdk/memory/_redis_cluster_memory_service.py @@ -0,0 +1,31 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Redis Cluster-backed long-term memory service.""" + +from __future__ import annotations + +from typing import Any + +from trpc_agent_sdk.storage import RedisClusterStorage + +from ._redis_memory_service import RedisMemoryService + + +class RedisClusterMemoryService(RedisMemoryService): + """Store and search cross-session memories in a Redis Cluster. + + Events retain the List-based format of :class:`RedisMemoryService`, but + cluster-wide memory lookups use ``SCAN`` across all primary nodes. This + avoids the incomplete results produced by a node-local ``KEYS`` command in + a sharded deployment. + + ``db_url`` identifies one cluster seed. Redis Cluster only supports + database 0; pass redis-py cluster options such as ``startup_nodes`` or + ``address_remap`` through keyword args when required by the deployment. + """ + + def _create_storage(self, db_url: str, is_async: bool, **kwargs: Any) -> RedisClusterStorage: + return RedisClusterStorage(is_async=is_async, redis_url=db_url, **kwargs) diff --git a/trpc_agent_sdk/memory/_redis_memory_service.py b/trpc_agent_sdk/memory/_redis_memory_service.py index 0c51dc05..c9d93303 100644 --- a/trpc_agent_sdk/memory/_redis_memory_service.py +++ b/trpc_agent_sdk/memory/_redis_memory_service.py @@ -43,7 +43,15 @@ def __init__( ): super().__init__(memory_service_config=memory_service_config, enabled=enabled) # Redis needs default TTL configuration - self._redis_storage = RedisStorage(is_async=is_async, redis_url=db_url, **kwargs) + self._redis_storage = self._create_storage(db_url=db_url, is_async=is_async, **kwargs) + + def _create_storage(self, db_url: str, is_async: bool, **kwargs: Any) -> RedisStorage: + """Create the backing storage. + + Subclasses override this factory to preserve memory behavior while + selecting a deployment-specific Redis client. + """ + return RedisStorage(is_async=is_async, redis_url=db_url, **kwargs) @override async def store_session(self, session: Session, agent_context: Optional[AgentContext] = None) -> None: diff --git a/trpc_agent_sdk/sessions/__init__.py b/trpc_agent_sdk/sessions/__init__.py index 33f41d5d..7f18a643 100644 --- a/trpc_agent_sdk/sessions/__init__.py +++ b/trpc_agent_sdk/sessions/__init__.py @@ -20,6 +20,7 @@ from ._in_memory_session_service import SessionWithTTL from ._in_memory_session_service import StateWithTTL from ._redis_session_service import RedisSessionService +from ._redis_cluster_session_service import RedisClusterSessionService from ._session import Session from ._session_summarizer import SessionSummarizer from ._session_summarizer import SessionSummary @@ -57,6 +58,7 @@ "SessionWithTTL", "StateWithTTL", "RedisSessionService", + "RedisClusterSessionService", "Session", "SessionSummarizer", "SessionSummary", diff --git a/trpc_agent_sdk/sessions/_redis_cluster_session_service.py b/trpc_agent_sdk/sessions/_redis_cluster_session_service.py new file mode 100644 index 00000000..1e19e905 --- /dev/null +++ b/trpc_agent_sdk/sessions/_redis_cluster_session_service.py @@ -0,0 +1,32 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Redis Cluster-backed session service.""" + +from __future__ import annotations + +from typing import Any + +from trpc_agent_sdk.storage import RedisClusterStorage + +from ._redis_session_service import RedisSessionService + + +class RedisClusterSessionService(RedisSessionService): + """Persist sessions in a Redis Cluster. + + The service preserves :class:`RedisSessionService` semantics while routing + single-key operations through redis-py's cluster client. Session listing + scans every primary node, so sessions are not missed when their keys occupy + different hash slots. + + ``db_url`` supplies one cluster seed, for example + ``redis://user:password@cluster-node-1:6379/0``. Redis Cluster only + supports database 0; additional redis-py cluster options, including + ``startup_nodes`` and ``address_remap``, may be provided as keyword args. + """ + + def _create_storage(self, db_url: str, is_async: bool, **kwargs: Any) -> RedisClusterStorage: + return RedisClusterStorage(is_async=is_async, redis_url=db_url, **kwargs) diff --git a/trpc_agent_sdk/sessions/_redis_session_service.py b/trpc_agent_sdk/sessions/_redis_session_service.py index 8751a75f..8bec47af 100644 --- a/trpc_agent_sdk/sessions/_redis_session_service.py +++ b/trpc_agent_sdk/sessions/_redis_session_service.py @@ -86,7 +86,15 @@ def __init__(self, # Default to store historical events for persistent backends. self._session_config.store_historical_events = True # Redis needs default TTL configuration - self._redis_storage = RedisStorage(is_async=is_async, redis_url=db_url, **kwargs) + self._redis_storage = self._create_storage(db_url=db_url, is_async=is_async, **kwargs) + + def _create_storage(self, db_url: str, is_async: bool, **kwargs: Any) -> RedisStorage: + """Create the backing storage. + + Subclasses override this factory to retain the session semantics while + selecting a different Redis deployment client, such as Redis Cluster. + """ + return RedisStorage(is_async=is_async, redis_url=db_url, **kwargs) @override async def create_session( diff --git a/trpc_agent_sdk/storage/__init__.py b/trpc_agent_sdk/storage/__init__.py index 6b839bf0..f06958b9 100644 --- a/trpc_agent_sdk/storage/__init__.py +++ b/trpc_agent_sdk/storage/__init__.py @@ -15,6 +15,9 @@ from ._redis import RedisExpire from ._redis import RedisSession from ._redis import RedisStorage +from ._redis_cluster import RedisClusterAsyncContextManager +from ._redis_cluster import RedisClusterClient +from ._redis_cluster import RedisClusterStorage from ._sql import SqlAsyncContextManager from ._sql import SqlCondition from ._sql import SqlKey @@ -45,6 +48,9 @@ "RedisExpire", "RedisSession", "RedisStorage", + "RedisClusterAsyncContextManager", + "RedisClusterClient", + "RedisClusterStorage", "SqlAsyncContextManager", "SqlCondition", "SqlKey", diff --git a/trpc_agent_sdk/storage/_redis_cluster.py b/trpc_agent_sdk/storage/_redis_cluster.py new file mode 100644 index 00000000..7b9db5b5 --- /dev/null +++ b/trpc_agent_sdk/storage/_redis_cluster.py @@ -0,0 +1,214 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Redis Cluster storage implementation. + +Unlike :class:`RedisStorage`, this adapter uses redis-py's native cluster +clients. Key-based commands are routed to their owning hash slot by the +client. Pattern queries use ``SCAN`` through ``RedisCluster.scan_iter`` so +they visit every primary node instead of issuing a node-local ``KEYS`` command. +""" + +from __future__ import annotations + +import inspect +import json +from typing import Any +from typing import Optional +from typing import Union + +from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster +from redis.cluster import RedisCluster as SyncRedisCluster + +from trpc_agent_sdk.log import logger + +from ._redis import RedisCommand +from ._redis import RedisCondition +from ._redis import RedisStorage + +RedisClusterClient = Union[AsyncRedisCluster, SyncRedisCluster] + + +class RedisClusterAsyncContextManager: + """Yield a shared Redis Cluster client without closing it per operation.""" + + def __init__(self, redis_storage: "RedisClusterStorage") -> None: + self._redis_storage = redis_storage + self._client: Optional[RedisClusterClient] = None + + async def __aenter__(self) -> RedisClusterClient: + self._client = await self._redis_storage.create_redis_session() + return self._client + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + # A cluster client owns pools for all discovered nodes. It must remain + # open for the storage lifetime, rather than being closed per request. + self._client = None + + +class RedisClusterStorage(RedisStorage): + """Storage adapter backed by redis-py's native Redis Cluster client. + + ``redis_url`` identifies one seed node. redis-py discovers the remaining + topology automatically; callers may alternatively pass ``startup_nodes`` + in ``kwargs``. Redis Cluster only supports logical database 0. + """ + + _SCAN_COUNT = 1000 + + def __init__(self, redis_url: str, is_async: bool = False, **kwargs: Any) -> None: + # State hashes are consumed as normal Python string dictionaries by the + # session services. Keep that invariant unless callers explicitly ask + # redis-py for raw responses. + kwargs.setdefault("decode_responses", True) + super().__init__(redis_url=redis_url, is_async=is_async, **kwargs) + self._redis_client: Optional[RedisClusterClient] = None + + async def create_redis_engine(self) -> None: + """Create the shared cluster-aware client lazily.""" + if self._redis_client: + return + try: + client_cls = AsyncRedisCluster if self._is_async else SyncRedisCluster + self._redis_client = client_cls.from_url(self._redis_url, **self._kwargs) + logger.debug("Redis Cluster client created successfully") + except Exception as ex: # pylint: disable=broad-except + raise ValueError(f"Failed to create Redis Cluster client for URL '{self._redis_url}'") from ex + + async def create_redis_session(self) -> RedisClusterClient: + """Return the shared client after lazily constructing it.""" + await self.create_redis_engine() + if not self._redis_client: + raise ValueError("Redis Cluster client not initialized") + return self._redis_client + + def create_db_session(self) -> RedisClusterAsyncContextManager: + """Create an async context that borrows the shared cluster client.""" + return RedisClusterAsyncContextManager(self) + + async def delete(self, conn: RedisClusterClient, key: str, conditions: Optional[RedisCondition] = None) -> None: + """Delete one key, letting redis-py route it to the correct slot.""" + ret = conn.delete(key) + if inspect.isawaitable(ret): + await ret + + async def query(self, conn: RedisClusterClient, key: str, conditions: RedisCondition) -> list[tuple[str, Any]]: + """Query keys across all cluster primaries using cursor-based scanning.""" + keys = await self._scan_keys(conn, key) + if conditions.limit > 0: + keys = keys[:conditions.limit] + + results: list[tuple[str, Any]] = [] + for redis_key in keys: + try: + key_type = await self.execute_command(conn, RedisCommand(method="type", args=(redis_key, ))) + key_type = self._decode_text(key_type) + + if key_type == "string": + value = await self.execute_command(conn, RedisCommand(method="get", args=(redis_key, ))) + if value is not None: + results.append((redis_key, self._deserialize_value(value))) + elif key_type == "hash": + hash_data = await self.execute_command(conn, RedisCommand(method="hgetall", args=(redis_key, ))) + if hash_data: + results.append((redis_key, hash_data)) + elif key_type == "list": + list_data = await self.execute_command(conn, RedisCommand(method="lrange", args=(redis_key, 0, -1))) + if list_data: + results.append((redis_key, list_data)) + elif key_type == "set": + set_data = await self.execute_command(conn, RedisCommand(method="smembers", args=(redis_key, ))) + if set_data: + results.append((redis_key, set_data)) + elif key_type == "zset": + zset_data = await self.execute_command( + conn, RedisCommand(method="zrange", args=(redis_key, 0, -1), kwargs={"withscores": True})) + if zset_data: + results.append((redis_key, zset_data)) + except Exception as ex: # pylint: disable=broad-except + logger.warning("Failed to query Redis Cluster key '%s': %s", redis_key, ex) + return results + + async def execute_command(self, conn: RedisClusterClient, command: RedisCommand) -> Any: + """Execute a command, replacing node-local ``KEYS`` with cluster ``SCAN``.""" + lower_method = command.method.lower() + if lower_method == "keys": + if not command.args: + raise ValueError("Redis KEYS command requires a match pattern") + return await self._scan_keys(conn, self._decode_text(command.args[0])) + + result = await super().execute_command(conn, command) + if lower_method == "hgetall" and isinstance(result, dict): + return {self._decode_text(k): self._decode_text(v) for k, v in result.items()} + return result + + async def close(self) -> None: + """Close all pools held by the shared Redis Cluster client.""" + if not self._redis_client: + return + try: + close_method = getattr(self._redis_client, "aclose", None) if self._is_async else None + if close_method is None: + close_method = self._redis_client.close + ret = close_method() + if inspect.isawaitable(ret): + await ret + except Exception as ex: # pylint: disable=broad-except + logger.info("Failed to close Redis Cluster client: %s", ex) + finally: + self._redis_client = None + + async def _scan_keys(self, conn: RedisClusterClient, match: str) -> list[str]: + """Return de-duplicated keys from every primary node in the cluster. + + ``RedisCluster.scan_iter`` starts its first scan on all nodes. This is + essential because ``KEYS`` without an explicit target node is scoped to + one default node in redis-py Cluster. + """ + iterator = conn.scan_iter(match=match, count=self._SCAN_COUNT) + if inspect.isawaitable(iterator): + iterator = await iterator + + keys: list[str] = [] + seen: set[str] = set() + + def _append(raw_key: Any) -> None: + normalized_key = self._decode_text(raw_key) + if normalized_key not in seen: + seen.add(normalized_key) + keys.append(normalized_key) + + if hasattr(iterator, "__aiter__"): + async for raw_key in iterator: + _append(raw_key) + else: + for raw_key in iterator: + _append(raw_key) + return keys + + @staticmethod + def _decode_text(value: Any) -> str: + """Decode Redis bytes while preserving normal string responses.""" + if isinstance(value, bytes): + return value.decode("utf-8") + return str(value) + + def _deserialize_value(self, value: Any) -> Any: + """Deserialize String values returned as bytes or decoded text.""" + if value is None: + return None + if isinstance(value, bytes): + try: + value_str = value.decode("utf-8") + except UnicodeDecodeError: + return value + elif isinstance(value, str): + value_str = value + else: + return value + try: + return json.loads(value_str) + except json.JSONDecodeError: + return value_str