Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions tests/memory/test_redis_cluster_memory_service.py
Original file line number Diff line number Diff line change
@@ -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)
24 changes: 24 additions & 0 deletions tests/sessions/test_redis_cluster_session_service.py
Original file line number Diff line number Diff line change
@@ -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)
95 changes: 95 additions & 0 deletions tests/storage/test_redis_cluster_storage.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 2 additions & 0 deletions trpc_agent_sdk/memory/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -28,6 +29,7 @@
"EventTtl",
"InMemoryMemoryService",
"RedisMemoryService",
"RedisClusterMemoryService",
"MemStorageData",
"MemStorageEvent",
"SqlMemoryService",
Expand Down
31 changes: 31 additions & 0 deletions trpc_agent_sdk/memory/_redis_cluster_memory_service.py
Original file line number Diff line number Diff line change
@@ -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)
10 changes: 9 additions & 1 deletion trpc_agent_sdk/memory/_redis_memory_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions trpc_agent_sdk/sessions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -57,6 +58,7 @@
"SessionWithTTL",
"StateWithTTL",
"RedisSessionService",
"RedisClusterSessionService",
"Session",
"SessionSummarizer",
"SessionSummary",
Expand Down
32 changes: 32 additions & 0 deletions trpc_agent_sdk/sessions/_redis_cluster_session_service.py
Original file line number Diff line number Diff line change
@@ -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)
10 changes: 9 additions & 1 deletion trpc_agent_sdk/sessions/_redis_session_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
6 changes: 6 additions & 0 deletions trpc_agent_sdk/storage/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -45,6 +48,9 @@
"RedisExpire",
"RedisSession",
"RedisStorage",
"RedisClusterAsyncContextManager",
"RedisClusterClient",
"RedisClusterStorage",
"SqlAsyncContextManager",
"SqlCondition",
"SqlKey",
Expand Down
Loading
Loading