"""Async implementation of Redis checkpoint saver.""" from __future__ import annotations import asyncio import json import logging from collections import defaultdict from contextlib import asynccontextmanager from types import TracebackType from typing import ( Any, AsyncIterator, Dict, List, Optional, Sequence, Tuple, Type, Union, cast, ) import orjson from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, PendingWrite, get_checkpoint_id, ) from langgraph.constants import TASKS from redis.asyncio import Redis as AsyncRedis from redis.asyncio.cluster import RedisCluster as AsyncRedisCluster from redisvl.index import AsyncSearchIndex from redisvl.query import FilterQuery from redisvl.query.filter import Num, Tag from redisvl.redis.connection import RedisConnectionFactory from ulid import ULID from langgraph.checkpoint.redis.base import ( CHECKPOINT_PREFIX, CHECKPOINT_WRITE_PREFIX, REDIS_KEY_SEPARATOR, BaseRedisSaver, ) from langgraph.checkpoint.redis.key_registry import ( AsyncCheckpointKeyRegistry as AsyncKeyRegistry, ) from langgraph.checkpoint.redis.util import ( EMPTY_ID_SENTINEL, from_storage_safe_id, from_storage_safe_str, to_storage_safe_id, to_storage_safe_str, ) logger = logging.getLogger(__name__) class AsyncRedisSaver( BaseRedisSaver[Union[AsyncRedis, AsyncRedisCluster], AsyncSearchIndex] ): """Async Redis implementation for checkpoint saver. Supports standard Redis URLs (redis://), SSL (rediss://), and Sentinel URLs (redis+sentinel://host:26379/service_name/db). """ _redis_url: str checkpoints_index: AsyncSearchIndex checkpoint_writes_index: AsyncSearchIndex _redis: Union[ AsyncRedis, AsyncRedisCluster ] # Support both standalone and cluster clients # Whether to assume the Redis server is a cluster; None triggers auto-detection cluster_mode: Optional[bool] = None _key_registry: Optional[AsyncKeyRegistry] = None # Track keys to avoid SCAN/KEYS # Instance-level cache (will be initialized in __init__) def __init__( self, redis_url: Optional[str] = None, *, redis_client: Optional[Union[AsyncRedis, AsyncRedisCluster]] = None, connection_args: Optional[Dict[str, Any]] = None, ttl: Optional[Dict[str, Any]] = None, checkpoint_prefix: str = CHECKPOINT_PREFIX, checkpoint_write_prefix: str = CHECKPOINT_WRITE_PREFIX, ) -> None: super().__init__( redis_url=redis_url, redis_client=redis_client, connection_args=connection_args, ttl=ttl, checkpoint_prefix=checkpoint_prefix, checkpoint_write_prefix=checkpoint_write_prefix, ) self.loop = asyncio.get_running_loop() # Instance-level cache for frequently used keys (limited size to prevent memory issues) self._key_cache: Dict[str, str] = {} self._key_cache_max_size = 1000 # Configurable limit self._separator = REDIS_KEY_SEPARATOR def configure_client( self, redis_url: Optional[str] = None, redis_client: Optional[Union[AsyncRedis, AsyncRedisCluster]] = None, connection_args: Optional[Dict[str, Any]] = None, ) -> None: """Configure the Redis client. Supports standard Redis URLs (redis://), SSL (rediss://), and Sentinel URLs (redis+sentinel://host:26379/service_name/db). """ self._owns_its_client = redis_client is None if redis_client is None: self._redis = RedisConnectionFactory.get_async_redis_connection( redis_url, **(connection_args or {}) ) else: self._redis = redis_client def create_indexes(self) -> None: """Create indexes without connecting to Redis.""" self.checkpoints_index = AsyncSearchIndex.from_dict( self.checkpoints_schema, redis_client=self._redis ) self.checkpoint_writes_index = AsyncSearchIndex.from_dict( self.writes_schema, redis_client=self._redis ) def _make_redis_checkpoint_key_cached( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str ) -> str: """Optimized key generation with caching.""" # Create cache key cache_key = f"ckpt:{thread_id}:{checkpoint_ns}:{checkpoint_id}" # Check cache first if cache_key in self._key_cache: return self._key_cache[cache_key] # Generate key using optimized string operations safe_thread_id = str(to_storage_safe_id(thread_id)) safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) safe_checkpoint_id = str(to_storage_safe_id(checkpoint_id)) # Use pre-computed prefix and join key = self._separator.join( [ self._checkpoint_prefix, safe_thread_id, safe_checkpoint_ns, safe_checkpoint_id, ] ) # Cache for future use (limit cache size to prevent memory issues) if len(self._key_cache) < self._key_cache_max_size: self._key_cache[cache_key] = key return key def _make_redis_checkpoint_writes_key_cached( self, thread_id: str, checkpoint_ns: str, checkpoint_id: str, task_id: str, idx: Optional[int], ) -> str: """Optimized writes key generation with caching.""" # Create cache key cache_key = f"write:{thread_id}:{checkpoint_ns}:{checkpoint_id}:{task_id}:{idx}" # Check cache first if cache_key in self._key_cache: return self._key_cache[cache_key] # Generate key using optimized string operations safe_thread_id = str(to_storage_safe_id(thread_id)) safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) safe_checkpoint_id = str(to_storage_safe_id(checkpoint_id)) # Build key components key_parts = [ self._checkpoint_write_prefix, safe_thread_id, safe_checkpoint_ns, safe_checkpoint_id, task_id, ] if idx is not None: key_parts.append(str(idx)) # Use pre-computed separator key = self._separator.join(key_parts) # Cache for future use (limit cache size) if len(self._key_cache) < 10000: self._key_cache[cache_key] = key return key async def __aenter__(self) -> AsyncRedisSaver: """Async context manager enter.""" await self.asetup() # Set client info once Redis is set up await self.aset_client_info() return self async def __aexit__( self, _exc_type: Optional[Type[BaseException]], _exc_val: Optional[BaseException], _exc_tb: Optional[TracebackType], ) -> None: """Async context manager exit.""" if self._owns_its_client: await self._redis.aclose() # RedisCluster doesn't have connection_pool attribute if getattr(self._redis, "connection_pool", None): coro = self._redis.connection_pool.disconnect() if coro: await coro # Prevent RedisVL from attempting to close the client # on an event loop in a separate thread. self.checkpoints_index._redis_client = None self.checkpoint_writes_index._redis_client = None async def asetup(self) -> None: """Set up the checkpoint saver.""" self.create_indexes() await self.checkpoints_index.create(overwrite=False) await self.checkpoint_writes_index.create(overwrite=False) # Detect cluster mode if not explicitly set await self._detect_cluster_mode() # Initialize key registry self._key_registry = AsyncKeyRegistry(self._redis) async def setup(self) -> None: # type: ignore[override] """Set up the checkpoint saver asynchronously. This method creates the necessary indices in Redis and detects cluster mode. It MUST be called before using the checkpointer. This async method follows the canonical pattern used by PostgreSQL and SQLite checkpointers in the LangGraph ecosystem. The type ignore is necessary because the base class defines a sync setup() method, but async checkpointers require an async setup() method to properly handle coroutines. Usage: await checkpointer.setup() """ await self.asetup() async def _detect_cluster_mode(self) -> None: """Detect if the Redis client is a cluster client by inspecting its class.""" if self.cluster_mode is not None: logger.info( f"Redis cluster_mode explicitly set to {self.cluster_mode}, skipping detection." ) return # Determine cluster mode based on client class if isinstance(self._redis, AsyncRedisCluster): logger.info("Redis client is a cluster client") self.cluster_mode = True else: logger.info("Redis client is a standalone client") self.cluster_mode = False async def _apply_ttl_to_keys( self, main_key: str, related_keys: Optional[list[str]] = None, ttl_minutes: Optional[float] = None, ) -> Any: """Apply Redis native TTL to keys asynchronously. Args: main_key: The primary Redis key related_keys: Additional Redis keys that should expire at the same time ttl_minutes: Time-to-live in minutes, overrides default_ttl if provided Use -1 to remove TTL (make keys persistent) Returns: Result of the Redis operation """ if ttl_minutes is None: # Check if there's a default TTL in config if self.ttl_config and "default_ttl" in self.ttl_config: ttl_minutes = self.ttl_config.get("default_ttl") if ttl_minutes is not None: # Special case: -1 means remove TTL (make persistent) if ttl_minutes == -1: # Apply PERSIST individually per key so that a single failure # does not prevent TTL removal on the remaining keys. all_keys = [main_key] + (related_keys or []) for key in all_keys: try: await self._redis.persist(key) except Exception: logger.warning("Failed to remove TTL from key: %s", key) return True # Regular TTL setting ttl_seconds = int(ttl_minutes * 60) # Apply TTL individually per key so that a single EXPIRE failure # (e.g. MOVED on Redis Enterprise proxy) does not prevent TTL # from being set on the remaining keys. all_keys = [main_key] + (related_keys or []) for key in all_keys: try: await self._redis.expire(key, ttl_seconds) except Exception: logger.warning("Failed to apply TTL to key: %s", key) return True async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from Redis asynchronously.""" thread_id = config["configurable"]["thread_id"] checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") # For values we store in Redis, we need to convert empty strings to the # sentinel value. storage_safe_thread_id = to_storage_safe_id(thread_id) storage_safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) if checkpoint_id and checkpoint_id != EMPTY_ID_SENTINEL: # Use direct key access instead of FT.SEARCH when checkpoint_id is known storage_safe_checkpoint_id = to_storage_safe_id(checkpoint_id) # Construct direct key for checkpoint data checkpoint_key = self._make_redis_checkpoint_key( storage_safe_thread_id, storage_safe_checkpoint_ns, storage_safe_checkpoint_id, ) # Create pipeline for efficient batch operations pipeline = self._redis.pipeline(transaction=False) # Add checkpoint data fetch to pipeline pipeline.json().get(checkpoint_key) # Add TTL check if refresh_on_read is enabled if self.ttl_config and self.ttl_config.get("refresh_on_read"): pipeline.ttl(checkpoint_key) # Execute pipeline to get checkpoint data and TTL pipeline_results = await pipeline.execute() checkpoint_data = pipeline_results[0] if not checkpoint_data: return None # Extract TTL if we fetched it current_ttl = None if self.ttl_config and self.ttl_config.get("refresh_on_read"): current_ttl = pipeline_results[1] # Create doc-like object from direct access doc = { "thread_id": checkpoint_data.get("thread_id", storage_safe_thread_id), "checkpoint_ns": checkpoint_data.get( "checkpoint_ns", storage_safe_checkpoint_ns ), "checkpoint_id": checkpoint_data.get( "checkpoint_id", storage_safe_checkpoint_id ), "parent_checkpoint_id": checkpoint_data.get( "parent_checkpoint_id", storage_safe_checkpoint_id ), "$.checkpoint": json.dumps(checkpoint_data.get("checkpoint", {})), "$.metadata": checkpoint_data.get( "metadata", "{}" ), # metadata is already a JSON string } else: # Try to get latest checkpoint using pointer latest_pointer_key = f"checkpoint_latest:{storage_safe_thread_id}:{storage_safe_checkpoint_ns}" checkpoint_key = await self._redis.get(latest_pointer_key) if not checkpoint_key: # No pointer means no checkpoints exist return None # Create pipeline for efficient operations pipeline = self._redis.pipeline(transaction=False) # Add checkpoint data fetch to pipeline pipeline.json().get(checkpoint_key) # Add TTL check if refresh_on_read is enabled if self.ttl_config and self.ttl_config.get("refresh_on_read"): pipeline.ttl(checkpoint_key) # Execute pipeline pipeline_results = await pipeline.execute() checkpoint_data = pipeline_results[0] if not checkpoint_data: # Pointer exists but checkpoint is missing - data inconsistency return None # Extract TTL if we fetched it current_ttl = None if self.ttl_config and self.ttl_config.get("refresh_on_read"): current_ttl = pipeline_results[1] # Create doc-like object from direct access doc = { "thread_id": checkpoint_data.get("thread_id", storage_safe_thread_id), "checkpoint_ns": checkpoint_data.get( "checkpoint_ns", storage_safe_checkpoint_ns ), "checkpoint_id": checkpoint_data.get("checkpoint_id"), "parent_checkpoint_id": checkpoint_data.get("parent_checkpoint_id"), "$.checkpoint": json.dumps(checkpoint_data.get("checkpoint", {})), "$.metadata": checkpoint_data.get( "metadata", "{}" ), # metadata is already a JSON string } doc_thread_id = from_storage_safe_id(doc["thread_id"]) doc_checkpoint_ns = from_storage_safe_str(doc["checkpoint_ns"]) doc_checkpoint_id = from_storage_safe_id(doc["checkpoint_id"]) doc_parent_checkpoint_id = from_storage_safe_id(doc["parent_checkpoint_id"]) # Lazy TTL refresh - only refresh if TTL is below threshold if self.ttl_config and self.ttl_config.get("refresh_on_read"): # If we didn't get TTL from pipeline (i.e., came from else branch), fetch it now if "current_ttl" not in locals(): # Get the checkpoint key checkpoint_key = self._make_redis_checkpoint_key( to_storage_safe_id(doc_thread_id), to_storage_safe_str(doc_checkpoint_ns), to_storage_safe_id(doc_checkpoint_id), ) current_ttl = await self._redis.ttl(checkpoint_key) # Always refresh TTL when refresh_on_read is enabled # This ensures all related keys maintain synchronized TTLs # Only refresh if key exists and has TTL (skip keys with no expiry) # TTL states: -2 = key doesn't exist, -1 = key exists but no TTL, 0 = expired, >0 = seconds remaining if current_ttl > 0: # Get write keys from registry instead of SCAN write_keys = [] if self._key_registry: write_keys = await self._key_registry.get_write_keys( doc_thread_id, doc_checkpoint_ns, doc_checkpoint_id ) # Apply TTL to checkpoint and write keys await self._apply_ttl_to_keys( checkpoint_key, write_keys if write_keys else None ) # Also refresh TTL on registry keys if they exist if self._key_registry and self.ttl_config: ttl_minutes = self.ttl_config.get("default_ttl") if ttl_minutes is not None: ttl_seconds = int(ttl_minutes * 60) # Registry TTL is handled per checkpoint await self._key_registry.apply_ttl( doc_thread_id, doc_checkpoint_ns, doc_checkpoint_id, ttl_seconds, ) # Fetch channel_values - pass channel_versions if we have them from direct access checkpoint_raw = ( doc.get("$.checkpoint") if isinstance(doc, dict) else getattr(doc, "$.checkpoint", None) ) if isinstance(checkpoint_raw, str): checkpoint_data_dict = json.loads(checkpoint_raw) else: checkpoint_data_dict = checkpoint_raw channel_versions_from_checkpoint = ( checkpoint_data_dict.get("channel_versions") if checkpoint_data_dict else None ) # Run channel_values, pending_sends, and pending_writes loads in parallel # Create list of coroutines to run tasks: List[Any] = [] # Always load channel values tasks.append( self.aget_channel_values( thread_id=doc_thread_id, checkpoint_ns=doc_checkpoint_ns, checkpoint_id=doc_checkpoint_id, channel_versions=channel_versions_from_checkpoint, ) ) # Conditionally load pending sends if parent exists if doc_parent_checkpoint_id: tasks.append( self._aload_pending_sends( thread_id=thread_id, checkpoint_ns=doc_checkpoint_ns, parent_checkpoint_id=doc_parent_checkpoint_id, ) ) # Always load pending writes tasks.append( self._aload_pending_writes(thread_id, checkpoint_ns, doc_checkpoint_id) ) # Execute all tasks in parallel - pending_sends is optional if doc_parent_checkpoint_id: results = await asyncio.gather(*tasks) channel_values: Dict[str, Any] = self._recursive_deserialize(results[0]) pending_sends: List[Tuple[str, Union[str, bytes]]] = results[1] pending_writes: List[PendingWrite] = results[2] else: # Only channel_values and pending_writes tasks results = await asyncio.gather(*tasks) channel_values = self._recursive_deserialize(results[0]) pending_sends = [] pending_writes = results[1] # Fetch and parse metadata raw_metadata = ( doc.get("$.metadata", "{}") if isinstance(doc, dict) else getattr(doc, "$.metadata", "{}") ) metadata_dict = ( json.loads(raw_metadata) if isinstance(raw_metadata, str) else raw_metadata ) # Ensure metadata matches CheckpointMetadata type sanitized_metadata = { k.replace("\u0000", ""): ( v.replace("\u0000", "") if isinstance(v, str) else v ) for k, v in metadata_dict.items() } metadata = cast(CheckpointMetadata, sanitized_metadata) config_param: RunnableConfig = { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": doc_checkpoint_id, } } # Handle both direct dict access and FT.SEARCH results checkpoint_data = doc["$.checkpoint"] if isinstance(checkpoint_data, dict): # Direct key access returns dict, convert to JSON string for consistency checkpoint_data = json.dumps(checkpoint_data) checkpoint_param = self._load_checkpoint( checkpoint_data, channel_values, pending_sends, ) # Build parent config if parent_checkpoint_id exists parent_config: RunnableConfig | None = None if doc_parent_checkpoint_id: parent_config = { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": doc_parent_checkpoint_id, } } return CheckpointTuple( config=config_param, checkpoint=checkpoint_param, metadata=metadata, parent_config=parent_config, pending_writes=pending_writes, ) async def alist( self, config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, # noqa: ARG002 limit: Optional[int] = None, ) -> AsyncIterator[CheckpointTuple]: """List checkpoints from Redis asynchronously.""" # Construct the filter expression filter_expression = [] if config: filter_expression.append( Tag("thread_id") == to_storage_safe_id(config["configurable"]["thread_id"]) ) if run_id := config["configurable"].get("run_id"): filter_expression.append(Tag("run_id") == to_storage_safe_id(run_id)) # Search for checkpoints with any namespace, including an empty # string, while `checkpoint_id` has to have a value. if checkpoint_ns := config["configurable"].get("checkpoint_ns"): filter_expression.append( Tag("checkpoint_ns") == to_storage_safe_str(checkpoint_ns) ) if checkpoint_id := get_checkpoint_id(config): filter_expression.append( Tag("checkpoint_id") == to_storage_safe_id(checkpoint_id) ) if filter: for k, v in filter.items(): if k == "source": filter_expression.append(Tag("source") == v) elif k == "step": filter_expression.append(Num("step") == v) elif k == "thread_id": filter_expression.append(Tag("thread_id") == to_storage_safe_id(v)) elif k == "run_id": filter_expression.append(Tag("run_id") == to_storage_safe_id(v)) else: raise ValueError(f"Unsupported filter key: {k}") if before: before_checkpoint_id = get_checkpoint_id(before) if before_checkpoint_id: try: before_ulid = ULID.from_str(before_checkpoint_id) before_ts = before_ulid.timestamp # Use numeric range query: checkpoint_ts < before_ts filter_expression.append(Num("checkpoint_ts") < before_ts) except Exception: # If not a valid ULID, ignore the before filter pass # Combine all filter expressions combined_filter = filter_expression[0] if filter_expression else "*" for expr in filter_expression[1:]: combined_filter &= expr # Construct the Redis query # Sort by checkpoint_id in descending order to get most recent checkpoints first query = FilterQuery( filter_expression=combined_filter, return_fields=[ "thread_id", "checkpoint_ns", "checkpoint_id", "parent_checkpoint_id", "$.checkpoint", "$.metadata", "has_writes", # Include has_writes to optimize pending_writes loading ], num_results=limit or 10000, sort_by=("checkpoint_id", "DESC"), ) # Execute the query asynchronously results = await self.checkpoints_index.search(query) # Pre-process all docs to collect batch query requirements all_docs_data = [] pending_sends_batch_keys = [] pending_writes_batch_keys = [] for doc in results.docs: # Extract all attributes once doc_dict = doc.__dict__ if hasattr(doc, "__dict__") else {} thread_id = from_storage_safe_id(doc["thread_id"]) checkpoint_ns = from_storage_safe_str(doc["checkpoint_ns"]) checkpoint_id = from_storage_safe_id(doc["checkpoint_id"]) parent_checkpoint_id = from_storage_safe_id(doc["parent_checkpoint_id"]) # Get channel values from inline checkpoint data (already returned by FT.SEARCH) checkpoint_data = doc_dict.get("$.checkpoint") or getattr( doc, "$.checkpoint", None ) if checkpoint_data: # Parse checkpoint to extract inline channel_values if isinstance(checkpoint_data, list) and checkpoint_data: checkpoint_data = checkpoint_data[0] # Use orjson for faster parsing checkpoint_dict = ( checkpoint_data if isinstance(checkpoint_data, dict) else orjson.loads(checkpoint_data) ) channel_values = self._recursive_deserialize( checkpoint_dict.get("channel_values", {}) ) else: # If checkpoint data is missing, the document is corrupted # Set empty channel values rather than attempting a fallback channel_values = {} # Collect batch keys for pending_sends if parent_checkpoint_id and parent_checkpoint_id != "None": batch_key = (thread_id, checkpoint_ns, parent_checkpoint_id) pending_sends_batch_keys.append(batch_key) # Collect batch keys for pending_writes checkpoint_has_writes = doc_dict.get("has_writes") or getattr( doc, "has_writes", False ) # Convert string "False" to boolean false if needed (optimize for common case) if checkpoint_has_writes == "true": checkpoint_has_writes = True elif checkpoint_has_writes == "false" or checkpoint_has_writes == "False": checkpoint_has_writes = False if checkpoint_has_writes: batch_key = (thread_id, checkpoint_ns, checkpoint_id) pending_writes_batch_keys.append(batch_key) # Store processed doc data for final iteration all_docs_data.append( { "doc": doc, "doc_dict": doc_dict, "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": checkpoint_id, "parent_checkpoint_id": parent_checkpoint_id, "checkpoint_data": checkpoint_data, "checkpoint_dict": checkpoint_dict if checkpoint_data else None, "channel_values": channel_values, "has_writes": checkpoint_has_writes, } ) # Load pending_sends for all parent checkpoints at once pending_sends_map = {} if pending_sends_batch_keys: pending_sends_map = await self._abatch_load_pending_sends( pending_sends_batch_keys ) # Load pending_writes for all checkpoints with writes at once pending_writes_map = {} if pending_writes_batch_keys: pending_writes_map = await self._abatch_load_pending_writes( pending_writes_batch_keys ) # Process the results using pre-loaded batch data for doc_data in all_docs_data: thread_id = doc_data["thread_id"] checkpoint_ns = doc_data["checkpoint_ns"] checkpoint_id = doc_data["checkpoint_id"] parent_checkpoint_id = doc_data["parent_checkpoint_id"] # Get pending_sends from batch results pending_sends: List[Tuple[str, Union[str, bytes]]] = [] if parent_checkpoint_id: batch_key = (thread_id, checkpoint_ns, parent_checkpoint_id) pending_sends = pending_sends_map.get(batch_key, []) # Fetch and parse metadata doc_dict = doc_data["doc_dict"] raw_metadata = doc_dict.get("$.metadata") or getattr( doc_data["doc"], "$.metadata", "{}" ) # Use orjson for faster parsing metadata_dict = ( orjson.loads(raw_metadata) if isinstance(raw_metadata, str) else raw_metadata ) # Only sanitize if null bytes detected (rare case) if any( "\u0000" in str(v) for v in metadata_dict.values() if isinstance(v, str) ): sanitized_metadata = { k.replace("\u0000", ""): ( v.replace("\u0000", "") if isinstance(v, str) else v ) for k, v in metadata_dict.items() } metadata = cast(CheckpointMetadata, sanitized_metadata) else: metadata = cast(CheckpointMetadata, metadata_dict) # Pre-create the config structure more efficiently config_param: RunnableConfig = { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": checkpoint_id, } } # Pass already parsed checkpoint_dict to avoid re-parsing checkpoint_param = self._load_checkpoint( ( doc_data["checkpoint_dict"] if doc_data["checkpoint_data"] else doc_data["doc"]["$.checkpoint"] ), doc_data["channel_values"], pending_sends, ) # Get pending_writes from batch results pending_writes: List[PendingWrite] = [] if doc_data["has_writes"]: batch_key = (thread_id, checkpoint_ns, checkpoint_id) pending_writes = pending_writes_map.get(batch_key, []) # Build parent config if parent_checkpoint_id exists parent_config: RunnableConfig | None = None if parent_checkpoint_id: parent_config = { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": parent_checkpoint_id, } } yield CheckpointTuple( config=config_param, checkpoint=checkpoint_param, metadata=metadata, parent_config=parent_config, pending_writes=pending_writes, ) async def aput( self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, new_versions: ChannelVersions, stream_mode: str = "values", ) -> RunnableConfig: """Store a checkpoint to Redis with proper transaction handling. This method ensures that all Redis operations are performed atomically using Redis transactions. In case of interruption (asyncio.CancelledError), the transaction will be aborted, ensuring consistency. Args: config: The config to associate with the checkpoint checkpoint: The checkpoint data to store metadata: Additional metadata to save with the checkpoint new_versions: New channel versions as of this write stream_mode: The streaming mode being used (values, updates, etc.) Returns: Updated configuration after storing the checkpoint Raises: asyncio.CancelledError: If the operation is cancelled/interrupted """ configurable = config["configurable"].copy() run_id = configurable.pop("run_id", metadata.get("run_id")) thread_id = configurable.pop("thread_id") checkpoint_ns = configurable.pop("checkpoint_ns") # Get checkpoint_id from config - this will be parent if saving a child config_checkpoint_id = configurable.pop("checkpoint_id", None) # For backward compatibility with thread_ts thread_ts = configurable.pop("thread_ts", "") # Determine the checkpoint ID # This follows the original logic but with clearer parent handling checkpoint_id = config_checkpoint_id or thread_ts or checkpoint.get("id", "") # If checkpoint has its own ID that's different from what we'd use, # and we have a config checkpoint_id, then config checkpoint_id is the parent parent_checkpoint_id = None if ( checkpoint.get("id") and config_checkpoint_id and checkpoint.get("id") != config_checkpoint_id ): parent_checkpoint_id = config_checkpoint_id checkpoint_id = checkpoint["id"] # For values we store in Redis, we need to convert empty strings to the # sentinel value. storage_safe_thread_id = to_storage_safe_id(thread_id) storage_safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) storage_safe_checkpoint_id = to_storage_safe_id(checkpoint_id) copy = checkpoint.copy() next_config = { "configurable": { "thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": checkpoint_id, } } # Store checkpoint data with cluster-aware handling try: # Store checkpoint data WITH inline channel values # Extract timestamp from checkpoint_id (ULID) checkpoint_ts = None if checkpoint_id: try: from ulid import ULID ulid_obj = ULID.from_str(checkpoint_id) checkpoint_ts = ulid_obj.timestamp # milliseconds since epoch except Exception: # If not a valid ULID, use current time import time checkpoint_ts = time.time() * 1000 checkpoint_data = { "thread_id": storage_safe_thread_id, "run_id": to_storage_safe_id(run_id) if run_id else "", "checkpoint_ns": storage_safe_checkpoint_ns, "checkpoint_id": storage_safe_checkpoint_id, "parent_checkpoint_id": ( to_storage_safe_id(parent_checkpoint_id) if parent_checkpoint_id else "" ), "checkpoint_ts": checkpoint_ts, "checkpoint": self._dump_checkpoint(copy), "metadata": self._dump_metadata(metadata), "has_writes": False, # Track if this checkpoint has pending writes } # store at top-level for filters in list() if all(key in metadata for key in ["source", "step"]): checkpoint_data["source"] = metadata["source"] checkpoint_data["step"] = metadata["step"] # Prepare checkpoint key checkpoint_key = self._make_redis_checkpoint_key_cached( thread_id, checkpoint_ns, checkpoint_id, ) # Calculate TTL in seconds if configured ttl_seconds = None if self.ttl_config and "default_ttl" in self.ttl_config: ttl_seconds = int(self.ttl_config["default_ttl"] * 60) # Store checkpoint with TTL in a single operation using SearchIndex await self.checkpoints_index.load( [checkpoint_data], keys=[checkpoint_key], ttl=ttl_seconds, # RedisVL applies TTL in its internal pipeline ) # For test compatibility: ensure TTL operations are visible to mocks if ( self.cluster_mode and self.ttl_config and "default_ttl" in self.ttl_config and ttl_seconds is not None ): # In cluster mode, also call expire directly so tests can verify try: await self._redis.expire(checkpoint_key, ttl_seconds) except Exception: logger.warning( "Failed to apply TTL to checkpoint key: %s", checkpoint_key ) # Update latest checkpoint pointer latest_pointer_key = f"checkpoint_latest:{storage_safe_thread_id}:{storage_safe_checkpoint_ns}" await self._redis.set(latest_pointer_key, checkpoint_key) # Apply TTL to latest pointer key as well (best-effort) if ttl_seconds is not None: try: await self._redis.expire(latest_pointer_key, ttl_seconds) except Exception: logger.warning( "Failed to apply TTL to latest pointer key: %s", latest_pointer_key, ) return next_config except asyncio.CancelledError: # Handle cancellation/interruption based on stream mode if stream_mode in ("values", "messages"): # For these modes, we want to ensure any partial state is committed # to allow resuming the stream later try: # Store minimal checkpoint data checkpoint_data = { "thread_id": storage_safe_thread_id, "run_id": to_storage_safe_id(run_id) if run_id else "", "checkpoint_ns": storage_safe_checkpoint_ns, "checkpoint_id": storage_safe_checkpoint_id, "parent_checkpoint_id": ( to_storage_safe_id( str(checkpoint.get("parent_checkpoint_id", "")) ) if checkpoint.get("parent_checkpoint_id") else "" ), "checkpoint": self._dump_checkpoint(copy), "metadata": self._dump_metadata( { **metadata, "interrupted": True, "stream_mode": stream_mode, } ), "has_writes": False, # Track if this checkpoint has pending writes } # Prepare checkpoint key checkpoint_key = self._make_redis_checkpoint_key( storage_safe_thread_id, storage_safe_checkpoint_ns, storage_safe_checkpoint_id, ) if self.cluster_mode: # For cluster mode, execute operation directly await self._redis.json().set( # type: ignore[misc] checkpoint_key, "$", checkpoint_data ) else: # For non-cluster mode, use pipeline pipeline = self._redis.pipeline(transaction=False) pipeline.json().set(checkpoint_key, "$", checkpoint_data) await pipeline.execute() except Exception: # If this also fails, we just propagate the original cancellation pass # Re-raise the cancellation raise except Exception as e: # Re-raise other exceptions raise e async def aput_writes( self, config: RunnableConfig, writes: Sequence[Tuple[str, Any]], task_id: str, task_path: str = "", ) -> None: """Store intermediate writes linked to a checkpoint using Redis JSON. This method uses Redis pipeline without transaction to avoid lock contention during parallel test execution. Args: config (RunnableConfig): Configuration of the related checkpoint. writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. task_path (str): Path of the task creating the writes. Raises: asyncio.CancelledError: If the operation is cancelled/interrupted """ if not writes: return thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = config["configurable"]["checkpoint_id"] # Transform writes into appropriate format writes_objects = [] for idx, (channel, value) in enumerate(writes): type_, blob = self.serde.dumps_typed(value) write_obj = { "thread_id": to_storage_safe_id(thread_id), "checkpoint_ns": to_storage_safe_str(checkpoint_ns), "checkpoint_id": to_storage_safe_id(checkpoint_id), "task_id": task_id, "task_path": task_path, "idx": WRITES_IDX_MAP.get(channel, idx), "channel": channel, "type": type_, "blob": self._encode_blob( blob ), # Encode bytes to base64 string for Redis } writes_objects.append(write_obj) try: # Determine if this is an upsert case upsert_case = all(w[0] in WRITES_IDX_MAP for w in writes) created_keys = [] if self.cluster_mode: # For cluster mode, execute operations individually for write_obj in writes_objects: key = self._make_redis_checkpoint_writes_key_cached( thread_id, checkpoint_ns, checkpoint_id, task_id, write_obj["idx"], # type: ignore[arg-type] ) # Redis JSON.SET is an UPSERT by default await self._redis.json().set(key, "$", cast(Any, write_obj)) # type: ignore[misc] created_keys.append(key) # Apply TTL to newly created keys if ( created_keys and self.ttl_config and "default_ttl" in self.ttl_config ): await self._apply_ttl_to_keys( created_keys[0], created_keys[1:] if len(created_keys) > 1 else None, ) # Register write keys in the key registry for cluster mode if self._key_registry: write_keys = [] for write_obj in writes_objects: key = self._make_redis_checkpoint_writes_key_cached( thread_id, checkpoint_ns, checkpoint_id, task_id, write_obj["idx"], # type: ignore[arg-type] ) write_keys.append(key) if write_keys: # Use per-checkpoint sorted set registry zset_key = self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) # Add all write keys with their index as score for ordering zadd_mapping = {key: idx for idx, key in enumerate(write_keys)} await self._redis.zadd(zset_key, zadd_mapping) # type: ignore[arg-type] # Apply TTL to registry key if configured (best-effort) if self.ttl_config and "default_ttl" in self.ttl_config: ttl_seconds = int(self.ttl_config.get("default_ttl") * 60) try: await self._redis.expire(zset_key, ttl_seconds) except Exception: logger.warning( "Failed to apply TTL to write registry key: %s", zset_key, exc_info=True, ) else: # For non-cluster mode, use pipeline without transaction to avoid lock contention. # IMPORTANT: Only critical commands (JSON.SET, JSON.MERGE, ZADD) go in the # pipeline. EXPIRE (TTL) commands are applied separately afterward to avoid # pipeline failures on Redis Enterprise proxy, where mixed JSON module + # native commands in a single pipeline can cause EXPIRE to fail, aborting # the entire pipeline and losing interrupt writes. pipeline = self._redis.pipeline(transaction=False) # Add all write operations to the pipeline (critical) for write_obj in writes_objects: key = self._make_redis_checkpoint_writes_key_cached( thread_id, checkpoint_ns, checkpoint_id, task_id, write_obj["idx"], # type: ignore[arg-type] ) pipeline.json().set(key, "$", cast(Any, write_obj)) created_keys.append(key) # Update checkpoint to indicate it has writes (critical) checkpoint_key = "" if writes_objects: checkpoint_key = self._make_redis_checkpoint_key( thread_id, checkpoint_ns, checkpoint_id ) pipeline.json().merge(checkpoint_key, "$", {"has_writes": True}) # Registry operation (critical) zset_key = "" if self._key_registry and created_keys: zset_key = self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) zadd_mapping = {key: idx for idx, key in enumerate(created_keys)} pipeline.zadd(zset_key, zadd_mapping) # type: ignore[arg-type] # Execute critical commands with raise_on_error=False to get # per-command results instead of aborting on first failure results = await pipeline.execute(raise_on_error=False) # Check results for critical command failures merge_failed = False for result in results: if isinstance(result, Exception): err_str = str(result) if "JSON.MERGE" in err_str or "merge" in err_str.lower(): merge_failed = True else: raise result # Handle JSON.MERGE fallback for older Redis versions if merge_failed and checkpoint_key: try: checkpoint_data = await self._redis.json().get( # type: ignore[misc] checkpoint_key ) if isinstance( checkpoint_data, dict ) and not checkpoint_data.get("has_writes"): checkpoint_data["has_writes"] = True await self._redis.json().set( # type: ignore[misc] checkpoint_key, "$", checkpoint_data ) except Exception: pass # Apply TTL separately (best-effort — failures here don't lose # writes). Individual calls ensure partial success: if one key's # EXPIRE fails on RE proxy, the others still get TTL applied. if ( created_keys and self.ttl_config and "default_ttl" in self.ttl_config ): ttl_seconds = int(self.ttl_config["default_ttl"] * 60) for key in created_keys: try: await self._redis.expire(key, ttl_seconds) except Exception: logger.warning( "Failed to apply TTL to checkpoint write key: %s", key, ) if zset_key and self.ttl_config and "default_ttl" in self.ttl_config: try: ttl_seconds = int(self.ttl_config["default_ttl"] * 60) await self._redis.expire(zset_key, ttl_seconds) except Exception: logger.warning( "Failed to apply TTL to write registry key: %s", zset_key, exc_info=True, ) except asyncio.CancelledError: # Handle cancellation/interruption # Pipeline will be automatically discarded # Either all operations succeed or none do raise except Exception as e: # Re-raise other exceptions raise e def put_writes( self, config: RunnableConfig, writes: Sequence[Tuple[str, Any]], task_id: str, task_path: str = "", ) -> None: """Synchronous wrapper for aput_writes. Args: config (RunnableConfig): Configuration of the related checkpoint. writes (List[Tuple[str, Any]]): List of writes to store. task_id (str): Identifier for the task creating the writes. task_path (str): Path of the task creating the writes. """ return asyncio.run_coroutine_threadsafe( self.aput_writes(config, writes, task_id), self.loop ).result() def get_channel_values( self, thread_id: str, checkpoint_ns: str = "", checkpoint_id: str = "" ) -> Dict[str, Any]: """Retrieve channel_values using efficient FT.SEARCH with checkpoint_id (sync wrapper).""" try: if asyncio.get_running_loop() is self.loop: raise asyncio.InvalidStateError( "Synchronous calls to AsyncRedisSaver are only allowed from a " "different thread. From the main thread, use the async interface." "For example, use `await checkpointer.get_channel_values(...)`." ) except RuntimeError: pass return asyncio.run_coroutine_threadsafe( self.aget_channel_values( thread_id, checkpoint_ns, checkpoint_id, ), self.loop, ).result() def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from Redis. Args: config (RunnableConfig): The config to use for retrieving the checkpoint. Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. Raises: asyncio.InvalidStateError: If called from the wrong thread/event loop """ try: # check if we are in the main thread, only bg threads can block if asyncio.get_running_loop() is self.loop: raise asyncio.InvalidStateError( "Synchronous calls to AsyncRedisSaver are only allowed from a " "different thread. From the main thread, use the async interface." "For example, use `await checkpointer.aget_tuple(...)` or `await " "graph.ainvoke(...)`." ) except RuntimeError: pass return asyncio.run_coroutine_threadsafe( self.aget_tuple(config), self.loop ).result() def put( self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, new_versions: ChannelVersions, ) -> RunnableConfig: """Store a checkpoint to Redis. Args: config (RunnableConfig): The config to associate with the checkpoint. checkpoint (Checkpoint): The checkpoint to save. metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. new_versions (ChannelVersions): New channel versions as of this write. Returns: RunnableConfig: Updated configuration after storing the checkpoint. Raises: asyncio.InvalidStateError: If called from the wrong thread/event loop """ try: # check if we are in the main thread, only bg threads can block if asyncio.get_running_loop() is self.loop: raise asyncio.InvalidStateError( "Synchronous calls to AsyncRedisSaver are only allowed from a " "different thread. From the main thread, use the async interface." "For example, use `await checkpointer.aput(...)` or `await " "graph.ainvoke(...)`." ) except RuntimeError: pass return asyncio.run_coroutine_threadsafe( self.aput(config, checkpoint, metadata, new_versions), self.loop ).result() @classmethod @asynccontextmanager async def from_conn_string( cls, redis_url: Optional[str] = None, *, redis_client: Optional[Union[AsyncRedis, AsyncRedisCluster]] = None, connection_args: Optional[Dict[str, Any]] = None, ttl: Optional[Dict[str, Any]] = None, checkpoint_prefix: str = CHECKPOINT_PREFIX, checkpoint_write_prefix: str = CHECKPOINT_WRITE_PREFIX, ) -> AsyncIterator[AsyncRedisSaver]: async with cls( redis_url=redis_url, redis_client=redis_client, connection_args=connection_args, ttl=ttl, checkpoint_prefix=checkpoint_prefix, checkpoint_write_prefix=checkpoint_write_prefix, ) as saver: yield saver async def aget_channel_values( self, thread_id: str, checkpoint_ns: str = "", checkpoint_id: str = "", channel_versions: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Retrieve channel_values using efficient FT.SEARCH with checkpoint_id.""" storage_safe_thread_id = to_storage_safe_id(thread_id) storage_safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) storage_safe_checkpoint_id = to_storage_safe_id(checkpoint_id) # Get checkpoint with inline channel_values using single JSON.GET operation (MongoDB approach) checkpoint_key = self._make_redis_checkpoint_key_cached( thread_id, checkpoint_ns, checkpoint_id, ) # Single JSON.GET operation to retrieve checkpoint with inline channel_values checkpoint_data = await self._redis.json().get(checkpoint_key, "$.checkpoint") # type: ignore[misc] if not checkpoint_data: return {} # checkpoint_data[0] is already a deserialized dict, not a typed tuple checkpoint = checkpoint_data[0] return checkpoint.get("channel_values", {}) async def _aload_pending_sends( self, thread_id: str, checkpoint_ns: str = "", parent_checkpoint_id: str = "" ) -> List[Tuple[str, Union[str, bytes]]]: """Load pending sends for a parent checkpoint. Args: thread_id: The thread ID checkpoint_ns: The checkpoint namespace parent_checkpoint_id: The ID of the parent checkpoint Returns: List of (type, blob) tuples representing pending sends """ if not parent_checkpoint_id: return [] # FAST PATH: Try sorted set registry first if self._key_registry: try: # Check if parent checkpoint has any writes in the sorted set write_count = await self._key_registry.get_write_count( thread_id, checkpoint_ns, parent_checkpoint_id ) if write_count == 0: # No writes for parent checkpoint - return immediately return [] # Get exact write keys from the per-checkpoint registry write_keys = await self._key_registry.get_write_keys( thread_id, checkpoint_ns, parent_checkpoint_id ) # Filter for TASKS channel writes task_write_keys = [] for key in write_keys: # Keys contain channel info: checkpoint_write:thread:ns:checkpoint:task:idx # We need to check if it's a TASKS channel write # This is a simple heuristic - we might need to fetch to be sure if TASKS in key or "__pregel_tasks" in key: task_write_keys.append(key) if not task_write_keys: return [] # Fetch task writes using pipeline (safe for cluster mode) pipeline = self._redis.pipeline(transaction=False) for key in task_write_keys: pipeline.json().get(key) results = await pipeline.execute() # Extract pending sends and sort them pending_sends_with_sort_keys = [] for write_data in results: if write_data and write_data.get("channel") == TASKS: pending_sends_with_sort_keys.append( ( write_data.get("task_path", ""), write_data.get("task_id", ""), write_data.get("idx", 0), write_data.get("type", ""), write_data.get("blob", b""), ) ) # Sort by task_path, task_id, idx pending_sends_with_sort_keys.sort(key=lambda x: (x[0], x[1], x[2])) # Return just the (type, blob) tuples return [(item[3], item[4]) for item in pending_sends_with_sort_keys] except Exception: # If sorted set approach fails, fall back to FT.SEARCH pass # Fallback to FT.SEARCH logic parent_writes_query = FilterQuery( filter_expression=( (Tag("thread_id") == to_storage_safe_id(thread_id)) & (Tag("checkpoint_ns") == to_storage_safe_str(checkpoint_ns)) & (Tag("checkpoint_id") == to_storage_safe_id(parent_checkpoint_id)) & (Tag("channel") == TASKS) ), return_fields=["type", "$.blob", "task_path", "task_id", "idx"], num_results=100, ) res = await self.checkpoint_writes_index.search(parent_writes_query) # Sort results for deterministic order docs = sorted( res.docs, key=lambda d: ( getattr(d, "task_path", ""), getattr(d, "task_id", ""), getattr(d, "idx", 0), ), ) # Convert to expected format return [ (d.type.encode(), blob) for d in docs if (blob := getattr(d, "$.blob", getattr(d, "blob", None))) is not None ] async def _aload_pending_writes( self, thread_id: str, checkpoint_ns: str = "", checkpoint_id: str = "", ) -> List[PendingWrite]: if checkpoint_id is None: return [] # Early return if no checkpoint_id # FAST PATH: Try sorted set registry first if self._key_registry: try: # Check if this checkpoint has any writes in the sorted set write_count = await self._key_registry.get_write_count( thread_id, checkpoint_ns, checkpoint_id ) if write_count == 0: # No writes for this checkpoint - return immediately return [] # Get exact write keys from the per-checkpoint registry write_keys = await self._key_registry.get_write_keys( thread_id, checkpoint_ns, checkpoint_id ) # Fetch all writes efficiently using pipeline pipeline = self._redis.pipeline(transaction=False) for key in write_keys: pipeline.json().get(key) results = await pipeline.execute() # Build the writes dictionary writes_dict: Dict[Tuple[str, str], Dict[str, Any]] = {} for write_data in results: if write_data: task_id = write_data.get("task_id", "") idx = str(write_data.get("idx", 0)) writes_dict[(task_id, idx)] = { "task_id": task_id, "idx": idx, "channel": write_data.get("channel", ""), "type": write_data.get("type", ""), "blob": write_data.get("blob", b""), } # Deserialize and return pending_writes = BaseRedisSaver._load_writes(self.serde, writes_dict) return pending_writes except Exception: # If sorted set approach fails, fall back to FT.SEARCH pass # FALLBACK: Use search index instead of keys() to avoid CrossSlot errors # Note: All tag fields use sentinel values for consistency writes_query = FilterQuery( filter_expression=(Tag("thread_id") == to_storage_safe_id(thread_id)) & (Tag("checkpoint_ns") == to_storage_safe_str(checkpoint_ns)) & (Tag("checkpoint_id") == to_storage_safe_id(checkpoint_id)), return_fields=["task_id", "idx", "channel", "type", "$.blob"], num_results=1000, # Adjust as needed ) writes_results = await self.checkpoint_writes_index.search(writes_query) # Sort results by idx to maintain order sorted_writes = sorted(writes_results.docs, key=lambda x: getattr(x, "idx", 0)) # Build the writes dictionary from search results search_writes_dict: Dict[Tuple[str, str], Dict[str, Any]] = {} for doc in sorted_writes: task_id = str(getattr(doc, "task_id", "")) idx = str(getattr(doc, "idx", 0)) blob_data = getattr(doc, "$.blob", "") # Ensure blob is bytes for deserialization if isinstance(blob_data, str): blob_data = blob_data.encode("utf-8") search_writes_dict[(task_id, idx)] = { "task_id": task_id, "idx": idx, "channel": str(getattr(doc, "channel", "")), "type": str(getattr(doc, "type", "")), "blob": blob_data, } pending_writes = BaseRedisSaver._load_writes(self.serde, search_writes_dict) return pending_writes async def _abatch_load_pending_sends( self, batch_keys: List[Tuple[str, str, str]] ) -> Dict[Tuple[str, str, str], List[Tuple[str, Union[str, bytes]]]]: """Batch load pending sends for multiple parent checkpoints. Args: batch_keys: List of (thread_id, checkpoint_ns, parent_checkpoint_id) tuples Returns: Dict mapping batch_key -> list of (type, blob) tuples """ if not batch_keys: return {} results_map = {} # Group by thread_id and checkpoint_ns for efficient querying grouped_keys: Dict[Tuple[str, str], List[str]] = {} for thread_id, checkpoint_ns, parent_checkpoint_id in batch_keys: group_key = (thread_id, checkpoint_ns) if group_key not in grouped_keys: grouped_keys[group_key] = [] grouped_keys[group_key].append(parent_checkpoint_id) # Batch query for each group for (thread_id, checkpoint_ns), parent_checkpoint_ids in grouped_keys.items(): storage_safe_thread_id = to_storage_safe_id(thread_id) storage_safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) storage_safe_parent_checkpoint_ids = [ to_storage_safe_id(pid) for pid in parent_checkpoint_ids ] # Build filter for multiple parent checkpoint IDs thread_filter = Tag("thread_id") == storage_safe_thread_id ns_filter = Tag("checkpoint_ns") == storage_safe_checkpoint_ns channel_filter = Tag("channel") == TASKS # Create filter for multiple parent checkpoint IDs (Tag supports lists) checkpoint_filter = ( Tag("checkpoint_id") == storage_safe_parent_checkpoint_ids ) batch_query = FilterQuery( filter_expression=thread_filter & ns_filter & checkpoint_filter & channel_filter, return_fields=[ "checkpoint_id", "type", "$.blob", "task_path", "task_id", "idx", ], num_results=1000, # Increased limit for batch loading ) batch_results = await self.checkpoint_writes_index.search(batch_query) # Group results by parent checkpoint ID writes_by_checkpoint: Dict[str, List[Any]] = {} for doc in batch_results.docs: parent_checkpoint_id = from_storage_safe_id(doc.checkpoint_id) if parent_checkpoint_id not in writes_by_checkpoint: writes_by_checkpoint[parent_checkpoint_id] = [] writes_by_checkpoint[parent_checkpoint_id].append(doc) # Sort and format results for each parent checkpoint for parent_checkpoint_id in parent_checkpoint_ids: batch_key = (thread_id, checkpoint_ns, parent_checkpoint_id) writes = writes_by_checkpoint.get(parent_checkpoint_id, []) # Sort results by task_path, task_id, idx sorted_writes = sorted( writes, key=lambda x: ( getattr(x, "task_path", ""), getattr(x, "task_id", ""), getattr(x, "idx", 0), ), ) # Extract type and blob pairs # Handle both direct attribute access and JSON path access results_map[batch_key] = [ ( getattr(doc, "type", ""), getattr(doc, "$.blob", getattr(doc, "blob", b"")), ) for doc in sorted_writes ] return results_map async def _abatch_load_pending_writes( self, batch_keys: List[Tuple[str, str, str]] ) -> Dict[Tuple[str, str, str], List[PendingWrite]]: """Batch load pending writes for multiple checkpoints. Args: batch_keys: List of (thread_id, checkpoint_ns, checkpoint_id) tuples Returns: Dict mapping batch_key -> list of PendingWrite objects """ if not batch_keys: return {} results_map: Dict[Tuple[str, str, str], List[PendingWrite]] = {} # If we have a key registry, use it for efficient batch loading if self._key_registry: # First, collect all write keys for all checkpoints in parallel pipeline = self._redis.pipeline(transaction=False) # Add all ZCARD operations to pipeline to check write counts for thread_id, checkpoint_ns, checkpoint_id in batch_keys: zset_key = self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) pipeline.zcard(zset_key) # Execute all ZCARD operations at once write_counts = await pipeline.execute() # Now get the actual keys for checkpoints that have writes pipeline = self._redis.pipeline(transaction=False) checkpoints_with_writes = [] for i, (thread_id, checkpoint_ns, checkpoint_id) in enumerate(batch_keys): if write_counts[i] > 0: checkpoints_with_writes.append( (thread_id, checkpoint_ns, checkpoint_id) ) zset_key = self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) pipeline.zrange(zset_key, 0, -1) else: # No writes for this checkpoint batch_key = (thread_id, checkpoint_ns, checkpoint_id) results_map[batch_key] = [] if checkpoints_with_writes: # Get all write keys at once all_write_keys_results = await pipeline.execute() # Now fetch all the actual write data in a single pipeline pipeline = self._redis.pipeline(transaction=False) write_key_mapping = {} # Maps pipeline index to checkpoint info pipeline_index = 0 for i, (thread_id, checkpoint_ns, checkpoint_id) in enumerate( checkpoints_with_writes ): write_keys = all_write_keys_results[i] if write_keys: decoded_keys = [ key.decode() if isinstance(key, bytes) else key for key in write_keys ] for key in decoded_keys: pipeline.json().get(key) write_key_mapping[pipeline_index] = ( thread_id, checkpoint_ns, checkpoint_id, key, ) pipeline_index += 1 # Execute all JSON.GET operations at once if pipeline_index > 0: all_writes_data = await pipeline.execute() # Group results by checkpoint writes_by_checkpoint: Dict[ Tuple[str, str, str], Dict[Tuple[str, str], Dict[str, Any]] ] = {} for idx, write_data in enumerate(all_writes_data): if write_data: thread_id, checkpoint_ns, checkpoint_id, key = ( write_key_mapping[idx] ) batch_key = (thread_id, checkpoint_ns, checkpoint_id) if batch_key not in writes_by_checkpoint: writes_by_checkpoint[batch_key] = {} task_id = write_data.get("task_id", "") idx_val = str(write_data.get("idx", 0)) writes_by_checkpoint[batch_key][(task_id, idx_val)] = { "task_id": task_id, "idx": idx_val, "channel": write_data.get("channel", ""), "type": write_data.get("type", ""), "blob": write_data.get("blob", b""), } # Deserialize and store results for batch_key, writes_dict in writes_by_checkpoint.items(): results_map[batch_key] = BaseRedisSaver._load_writes( self.serde, writes_dict ) else: # Fallback to batch search # Group by thread_id and checkpoint_ns for efficient querying grouped_keys: Dict[Tuple[str, str], List[str]] = {} for thread_id, checkpoint_ns, checkpoint_id in batch_keys: group_key = (thread_id, checkpoint_ns) if group_key not in grouped_keys: grouped_keys[group_key] = [] grouped_keys[group_key].append(checkpoint_id) # Batch query for each group for (thread_id, checkpoint_ns), checkpoint_ids in grouped_keys.items(): storage_safe_thread_id = to_storage_safe_id(thread_id) storage_safe_checkpoint_ns = to_storage_safe_str(checkpoint_ns) storage_safe_checkpoint_ids = [ to_storage_safe_id(cid) for cid in checkpoint_ids ] # Build batch query thread_filter = Tag("thread_id") == storage_safe_thread_id ns_filter = Tag("checkpoint_ns") == storage_safe_checkpoint_ns checkpoint_filter = Tag("checkpoint_id") == storage_safe_checkpoint_ids batch_query = FilterQuery( filter_expression=thread_filter & ns_filter & checkpoint_filter, return_fields=[ "checkpoint_id", "task_id", "idx", "channel", "type", "$.blob", ], num_results=5000, # Increased limit for batch ) batch_results = await self.checkpoint_writes_index.search(batch_query) # Group results by checkpoint ID fallback_writes_by_checkpoint: Dict[ str, Dict[Tuple[str, str], Dict[str, Any]] ] = {} for doc in batch_results.docs: checkpoint_id = from_storage_safe_id(doc.checkpoint_id) if checkpoint_id not in fallback_writes_by_checkpoint: fallback_writes_by_checkpoint[checkpoint_id] = {} task_id = getattr(doc, "task_id", "") idx_str = str(getattr(doc, "idx", 0)) blob = getattr(doc, "$.blob", getattr(doc, "blob", b"")) fallback_writes_by_checkpoint[checkpoint_id][(task_id, idx_str)] = { "task_id": task_id, "idx": idx_str, "channel": getattr(doc, "channel", ""), "type": getattr(doc, "type", ""), "blob": blob, } # Process results for each checkpoint for checkpoint_id in checkpoint_ids: batch_key = (thread_id, checkpoint_ns, checkpoint_id) writes_dict = fallback_writes_by_checkpoint.get(checkpoint_id, {}) results_map[batch_key] = BaseRedisSaver._load_writes( self.serde, writes_dict ) return results_map async def adelete_thread(self, thread_id: str) -> None: """Delete all checkpoints and writes associated with a specific thread ID. Args: thread_id: The thread ID whose checkpoints should be deleted. """ storage_safe_thread_id = to_storage_safe_id(thread_id) # Delete all checkpoints for this thread checkpoint_query = FilterQuery( filter_expression=Tag("thread_id") == storage_safe_thread_id, return_fields=["checkpoint_ns", "checkpoint_id"], num_results=10000, # Get all checkpoints for this thread ) checkpoint_results = await self.checkpoints_index.search(checkpoint_query) # Collect all keys to delete keys_to_delete = [] checkpoint_namespaces = set() for doc in checkpoint_results.docs: checkpoint_ns = getattr(doc, "checkpoint_ns", "") checkpoint_id = getattr(doc, "checkpoint_id", "") # Track unique namespaces for latest pointer cleanup checkpoint_namespaces.add(checkpoint_ns) # Delete checkpoint key checkpoint_key = self._make_redis_checkpoint_key( storage_safe_thread_id, checkpoint_ns, checkpoint_id ) keys_to_delete.append(checkpoint_key) # Add latest checkpoint pointers to deletion list for checkpoint_ns in checkpoint_namespaces: latest_pointer_key = f"checkpoint_latest:{storage_safe_thread_id}:{to_storage_safe_str(checkpoint_ns)}" keys_to_delete.append(latest_pointer_key) # Channel values are stored inline — no separate blob keys to clean up. # Delete all writes for this thread writes_query = FilterQuery( filter_expression=Tag("thread_id") == storage_safe_thread_id, return_fields=["checkpoint_ns", "checkpoint_id", "task_id", "idx"], num_results=10000, ) writes_results = await self.checkpoint_writes_index.search(writes_query) for doc in writes_results.docs: checkpoint_ns = getattr(doc, "checkpoint_ns", "") checkpoint_id = getattr(doc, "checkpoint_id", "") task_id = getattr(doc, "task_id", "") idx = getattr(doc, "idx", 0) write_key = self._make_redis_checkpoint_writes_key( storage_safe_thread_id, checkpoint_ns, checkpoint_id, task_id, idx ) keys_to_delete.append(write_key) # Delete the registry sorted sets for each checkpoint if self._key_registry: # Get unique checkpoints from the results we already have processed_checkpoints = set() for doc in checkpoint_results.docs: checkpoint_ns = getattr(doc, "checkpoint_ns", "") checkpoint_id = getattr(doc, "checkpoint_id", "") checkpoint_key = (thread_id, checkpoint_ns, checkpoint_id) if checkpoint_key not in processed_checkpoints: processed_checkpoints.add(checkpoint_key) # Add the write registry key for this checkpoint zset_key = self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) keys_to_delete.append(zset_key) # Execute all deletions based on cluster mode if self.cluster_mode: # For cluster mode, delete keys individually for key in keys_to_delete: await self._redis.delete(key) else: # For non-cluster mode, use pipeline for efficiency pipeline = self._redis.pipeline() for key in keys_to_delete: pipeline.delete(key) await pipeline.execute() async def aprune( self, thread_ids: Sequence[str], *, strategy: str = "keep_latest", keep_last: Optional[int] = None, max_results: int = 10000, ) -> None: """Prune old checkpoints for the given threads per namespace. Retains the most-recent checkpoints **per checkpoint namespace** and removes the rest, along with their associated write keys and key-registry sorted sets. Each namespace (root ``""`` and any subgraph namespaces) is treated as an independent checkpoint chain. Channel values are stored inline within each checkpoint document, so they are automatically removed when the checkpoint document is deleted. Args: thread_ids: Thread IDs whose old checkpoints should be pruned. strategy: Pruning strategy. ``"keep_latest"`` retains only the most recent checkpoint per namespace (default). ``"delete"`` removes all checkpoints for the thread. keep_last: Optional override — number of recent checkpoints to retain per namespace. When provided, takes precedence over ``strategy``. Use ``keep_last=0`` to remove all checkpoints. max_results: Maximum number of checkpoints fetched from the index per thread in a single query. Defaults to 10 000. """ # Resolve keep_last from strategy if not explicitly provided if keep_last is None: if strategy == "delete": keep_last = 0 else: keep_last = 1 # Validate inputs if not thread_ids: raise ValueError("``thread_ids`` must be a non-empty sequence") if keep_last < 0: raise ValueError(f"``keep_last`` must be >= 0, got {keep_last}") if max_results < 1: raise ValueError(f"``max_results`` must be >= 1, got {max_results}") for thread_id in thread_ids: storage_safe_thread_id = to_storage_safe_id(thread_id) # Fetch all checkpoints for this thread across all namespaces checkpoint_query = FilterQuery( filter_expression=Tag("thread_id") == storage_safe_thread_id, return_fields=["checkpoint_ns", "checkpoint_id"], num_results=max_results, ) checkpoint_results = await self.checkpoints_index.search(checkpoint_query) if not checkpoint_results.docs: continue # Group by namespace — each namespace is an independent checkpoint chain # (root graph vs. subgraph checkpoints must be evicted independently). by_ns: Dict[str, list] = defaultdict(list) for doc in checkpoint_results.docs: ns = getattr(doc, "checkpoint_ns", "") by_ns[ns].append(doc) # Within each namespace sort newest-first (ULIDs are lex time-ordered) # and collect checkpoints that fall outside the keep_last window. # Track namespaces where every checkpoint is evicted so we can clean # up the checkpoint_latest:{thread}:{ns} pointer key too to_evict = [] fully_evicted_ns: set = set() for ns, ns_docs in by_ns.items(): ns_sorted = sorted( ns_docs, key=lambda d: getattr(d, "checkpoint_id", ""), reverse=True, ) ns_evicted = ns_sorted[keep_last:] to_evict.extend(ns_evicted) if len(ns_evicted) == len(ns_docs): # nothing left in this namespace fully_evicted_ns.add(ns) if not to_evict: continue keys_to_delete = [] for doc in to_evict: checkpoint_ns = getattr(doc, "checkpoint_ns", "") checkpoint_id = getattr(doc, "checkpoint_id", "") # Evict checkpoint document keys_to_delete.append( self._make_redis_checkpoint_key( storage_safe_thread_id, checkpoint_ns, checkpoint_id ) ) # Evict all write documents for this checkpoint writes_query = FilterQuery( filter_expression=( (Tag("thread_id") == storage_safe_thread_id) & (Tag("checkpoint_id") == checkpoint_id) ), return_fields=["checkpoint_ns", "checkpoint_id", "task_id", "idx"], num_results=max_results, ) writes_results = await self.checkpoint_writes_index.search(writes_query) for wdoc in writes_results.docs: keys_to_delete.append( self._make_redis_checkpoint_writes_key( storage_safe_thread_id, getattr(wdoc, "checkpoint_ns", ""), getattr(wdoc, "checkpoint_id", ""), getattr(wdoc, "task_id", ""), int(getattr(wdoc, "idx", 0)), ) ) # Evict key-registry sorted set for this checkpoint if self._key_registry: keys_to_delete.append( self._key_registry.make_write_keys_zset_key( thread_id, checkpoint_ns, checkpoint_id ) ) for ns in fully_evicted_ns: keys_to_delete.append( f"checkpoint_latest:{storage_safe_thread_id}:{ns}" ) if self.cluster_mode: for key in keys_to_delete: await self._redis.delete(key) else: pipeline = self._redis.pipeline() for key in keys_to_delete: pipeline.delete(key) await pipeline.execute()