Upgrade AMS to Docket 0.18.2 cluster support (#216)

* Upgrade to Docket 0.18.2 so we can support Redis Cluster

Author: abrookins

.env.example

@@ -1,5 +1,6 @@
 # Redis connection
 REDIS_URL=redis://localhost:6379
+# For Redis Cluster with Docket 0.18.2+, use redis+cluster://host:port
 
 # Server port
 PORT=8000

README.md

@@ -54,6 +54,10 @@ expects a separate `agent-memory task-worker` process for non-blocking
 background tasks. The example above shows how to override this to use the
 asyncio backend for a single-container development setup.
 
+If you are connecting to a Redis Cluster and want Docket-backed workers, set
+`REDIS_URL` to a `redis+cluster://...` or `rediss+cluster://...` URL. AMS will
+translate that URL for its other Redis clients internally.
+
 **Production Deployment**:
 
 For production, run separate containers for the API and background workers:

agent_memory_server/cli.py

@@ -23,7 +23,11 @@
     migrate_add_memory_type_3,
     migrate_normalize_tag_separators_4,
 )
-from agent_memory_server.utils.redis import get_redis_conn
+from agent_memory_server.utils.redis import (
+    docket_stream_key,
+    get_redis_conn,
+    redis_url_for_docket,
+)
 
 
 logger = get_logger(__name__)
@@ -435,7 +439,7 @@ async def setup_and_run_task():
         # Initialize Docket client
         async with Docket(
             name=settings.docket_name,
-            url=settings.redis_url,
+            url=redis_url_for_docket(settings.redis_url),
         ) as docket:
             click.echo(f"Scheduling task {task_path} with arguments: {task_args}")
             await docket.add(task_func)(**task_args)
@@ -475,7 +479,7 @@ async def _ensure_stream_and_group():
         from redis.exceptions import ResponseError
 
         redis = await get_redis_conn()
-        stream_key = f"{settings.docket_name}:stream"
+        stream_key = docket_stream_key(settings.docket_name, settings.redis_url)
         group_name = "docket-workers"
 
         try:
@@ -493,7 +497,7 @@ async def _run_worker():
         await get_redis_conn()
         await Worker.run(
             docket_name=settings.docket_name,
-            url=settings.redis_url,
+            url=redis_url_for_docket(settings.redis_url),
             concurrency=concurrency,
             redelivery_timeout=timedelta(seconds=redelivery_timeout),
             tasks=["agent_memory_server.docket_tasks:task_collection"],

agent_memory_server/dependencies.py

@@ -8,6 +8,7 @@
 
 from agent_memory_server.config import settings
 from agent_memory_server.logging import get_logger
+from agent_memory_server.utils.redis import redis_url_for_docket
 
 
 logger = get_logger(__name__)
@@ -49,7 +50,7 @@ def run_in_thread():
                 async def schedule_task():
                     async with Docket(
                         name=settings.docket_name,
-                        url=settings.redis_url,
+                        url=redis_url_for_docket(settings.redis_url),
                     ) as docket:
                         # Schedule task in Docket's queue
                         await docket.add(func)(*args, **kwargs)

agent_memory_server/docket_tasks.py

@@ -26,6 +26,7 @@
     periodic_refresh_summary_views,
     refresh_summary_view,
 )
+from agent_memory_server.utils.redis import redis_url_for_docket
 
 
 logger = logging.getLogger(__name__)
@@ -58,7 +59,7 @@ async def register_tasks() -> None:
     # Initialize Docket client
     async with Docket(
         name=settings.docket_name,
-        url=settings.redis_url,
+        url=redis_url_for_docket(settings.redis_url),
     ) as docket:
         # Register all tasks
         for task in task_collection:

agent_memory_server/long_term_memory.py

@@ -46,7 +46,7 @@
     rerank_with_recency,
     update_memory_hash_if_text_changed,
 )
-from agent_memory_server.utils.redis import get_redis_conn
+from agent_memory_server.utils.redis import get_redis_conn, redis_url_for_docket
 from agent_memory_server.utils.tag_codec import encode_tag_values
 
 
@@ -249,7 +249,7 @@ async def schedule_trailing_extraction(
         try:
             async with Docket(
                 name=settings.docket_name,
-                url=settings.redis_url,
+                url=redis_url_for_docket(settings.redis_url),
             ) as docket:
                 # Schedule with a unique key per session
                 # If there's already a pending task for this session, the new one

agent_memory_server/memory_vector_db_factory.py

@@ -31,6 +31,7 @@
     MemoryVectorDatabase,
     RedisVLMemoryVectorDatabase,
 )
+from agent_memory_server.utils.redis import redis_url_for_redisvl
 
 
 logger = logging.getLogger(__name__)
@@ -187,7 +188,10 @@ def create_redis_memory_vector_db(
     """
     try:
         schema = _build_redis_schema()
-        index = AsyncSearchIndex.from_dict(schema, redis_url=settings.redis_url)
+        index = AsyncSearchIndex.from_dict(
+            schema,
+            redis_url=redis_url_for_redisvl(settings.redis_url),
+        )
         return RedisVLMemoryVectorDatabase(index, embeddings)
     except Exception as e:
         logger.error(f"Error creating Redis memory vector database: {e}")

agent_memory_server/utils/redis.py

@@ -2,32 +2,118 @@
 
 import logging
 from typing import Any
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
 
 from redis.asyncio import Redis
+from redis.asyncio.cluster import RedisCluster
 
 from agent_memory_server.config import settings
 
 
 logger = logging.getLogger(__name__)
-_redis_pool: Redis | None = None
+_redis_pool: Redis | RedisCluster | None = None
+_CLUSTER_SCHEMES = {"redis+cluster", "rediss+cluster"}
 
 
-async def get_redis_conn(url: str = settings.redis_url, **kwargs) -> Redis:
+def _netloc_has_multiple_hosts(netloc: str) -> bool:
+    host_part = netloc.rsplit("@", 1)[-1]
+    return "," in host_part
+
+
+def is_redis_cluster_url(url: str) -> bool:
+    """Return True when the URL targets a Redis Cluster deployment."""
+    parsed = urlparse(url)
+    if parsed.scheme in _CLUSTER_SCHEMES:
+        return True
+
+    query = dict(parse_qsl(parsed.query, keep_blank_values=True))
+    if query.get("cluster", "").lower() == "true":
+        return True
+
+    return _netloc_has_multiple_hosts(parsed.netloc)
+
+
+def _strip_cluster_query(url: str) -> str:
+    parsed = urlparse(url)
+    query_items = [
+        (key, value)
+        for key, value in parse_qsl(parsed.query, keep_blank_values=True)
+        if key.lower() != "cluster"
+    ]
+    return urlunparse(parsed._replace(query=urlencode(query_items, doseq=True)))
+
+
+def redis_url_for_docket(url: str) -> str:
+    """Normalize a Redis URL for Docket's cluster-aware URL scheme."""
+    if not is_redis_cluster_url(url):
+        return url
+
+    parsed = urlparse(_strip_cluster_query(url))
+    if parsed.scheme in _CLUSTER_SCHEMES:
+        return urlunparse(parsed)
+
+    return urlunparse(parsed._replace(scheme=f"{parsed.scheme}+cluster"))
+
+
+def redis_url_for_redisvl(url: str) -> str:
+    """Normalize a Redis URL for RedisVL's cluster detection."""
+    if not is_redis_cluster_url(url):
+        return url
+
+    parsed = urlparse(_strip_cluster_query(url))
+    query_items = parse_qsl(parsed.query, keep_blank_values=True)
+    query_items.append(("cluster", "true"))
+    return urlunparse(
+        parsed._replace(
+            scheme=parsed.scheme.replace("+cluster", ""),
+            query=urlencode(query_items, doseq=True),
+        )
+    )
+
+
+def redis_url_for_async_redis(url: str) -> str:
+    """Normalize a Redis URL for redis-py's async standalone/cluster clients."""
+    if not is_redis_cluster_url(url):
+        return url
+
+    parsed = urlparse(_strip_cluster_query(url))
+    return urlunparse(parsed._replace(scheme=parsed.scheme.replace("+cluster", "")))
+
+
+def docket_prefix(name: str, redis_url: str) -> str:
+    """Return the Docket key prefix for the given deployment mode."""
+    if is_redis_cluster_url(redis_url):
+        return f"{{{name}}}"
+    return name
+
+
+def docket_stream_key(name: str, redis_url: str) -> str:
+    """Return the Docket stream key with cluster-safe hashing when needed."""
+    return f"{docket_prefix(name, redis_url)}:stream"
+
+
+async def get_redis_conn(
+    url: str = settings.redis_url, **kwargs
+) -> Redis | RedisCluster:
     """Get a Redis connection.
 
     Args:
         url: Redis connection URL, or None to use settings.redis_url
-        **kwargs: Additional arguments to pass to Redis.from_url
+        **kwargs: Additional arguments to pass to the Redis client
 
     Returns:
-        A Redis client instance
+        A Redis or RedisCluster client instance
     """
     global _redis_pool
 
     # Always use the existing _redis_pool if it's not None, regardless of the URL parameter
     # This ensures connection reuse and prevents multiple Redis connections
     if _redis_pool is None:
-        _redis_pool = Redis.from_url(url, **kwargs)
+        normalized_url = redis_url_for_async_redis(url)
+        if is_redis_cluster_url(url):
+            _redis_pool = RedisCluster.from_url(normalized_url, **kwargs)
+        else:
+            _redis_pool = Redis.from_url(normalized_url, **kwargs)
     return _redis_pool
 
 

docs/configuration.md

@@ -36,6 +36,10 @@ log_level: INFO
 REDIS_URL=redis://localhost:6379  # Redis connection string
 ```
 
+For Redis Cluster deployments, use `redis+cluster://host:port` or
+`rediss+cluster://host:port`. AMS translates that setting for Docket,
+RedisVL, and direct Redis connections internally.
+
 ### AI Model Configuration
 ```bash
 # Generation models for LLM tasks
@@ -306,7 +310,7 @@ compaction_every_minutes: 15
 ### High-Performance Setup
 ```yaml
 # config-performance.yaml
-redis_url: redis://redis-cluster:6379
+redis_url: redis+cluster://redis-cluster:6379
 fast_model: gpt-4o-mini
 slow_model: gpt-4o
 redisvl_indexing_algorithm: HNSW

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     "pydantic>=2.5.2",
     "pydantic-settings>=2.8.1",
     "python-dotenv>=1.0.0",
-    "pydocket>=0.6.3",
+    "pydocket>=0.18.2",
     "redisvl>=0.6.0",
     "structlog>=25.2.0",
     "tiktoken>=0.5.1",