fix: Use watches on compare-and-set pipelines

Author: abrookins

langgraph/checkpoint/redis/aio.py

@@ -10,6 +10,7 @@
 from typing import Any, List, Optional, Sequence, Tuple, Type, cast
 
 from langchain_core.runnables import RunnableConfig
+from redis import WatchError
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import FilterQuery
 from redisvl.query.filter import Num, Tag
@@ -418,38 +419,28 @@ async def aput_writes(
             }
             writes_objects.append(write_obj)
 
-            # For each write, check existence and then perform appropriate operation
-            async with self.checkpoints_index.client.pipeline(
-                transaction=False
-            ) as pipeline:
-                for write_obj in writes_objects:
-                    key = self._make_redis_checkpoint_writes_key(
-                        thread_id,
-                        checkpoint_ns,
-                        checkpoint_id,
-                        task_id,
-                        write_obj["idx"],
-                    )
-
-                    # First check if key exists
-                    key_exists = await self._redis.exists(key) == 1
-
-                    if all(w[0] in WRITES_IDX_MAP for w in writes):
-                        # UPSERT case - only update specific fields
-                        if key_exists:
-                            # Update only channel, type, and blob fields
-                            pipeline.json().set(key, "$.channel", write_obj["channel"])
-                            pipeline.json().set(key, "$.type", write_obj["type"])
-                            pipeline.json().set(key, "$.blob", write_obj["blob"])
+            upsert_case = all(w[0] in WRITES_IDX_MAP for w in writes)
+            for write_obj in writes_objects:
+                key = self._make_redis_checkpoint_writes_key(
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    task_id,
+                    write_obj["idx"],
+                )
+                async def tx(pipe, key=key, write_obj=write_obj, upsert_case=upsert_case):
+                    exists = await pipe.exists(key)
+                    if upsert_case:
+                        if exists:
+                            await pipe.json().set(key, "$.channel", write_obj["channel"])
+                            await pipe.json().set(key, "$.type", write_obj["type"])
+                            await pipe.json().set(key, "$.blob", write_obj["blob"])
                         else:
-                            # For new records, set the complete object
-                            pipeline.json().set(key, "$", write_obj)
+                            await pipe.json().set(key, "$", write_obj)
                     else:
-                        # INSERT case - only insert if doesn't exist
-                        if not key_exists:
-                            pipeline.json().set(key, "$", write_obj)
-
-                await pipeline.execute()
+                        if not exists:
+                            await pipe.json().set(key, "$", write_obj)
+                await self._redis.transaction(tx, key)
 
     def put_writes(
         self,

langgraph/checkpoint/redis/ashallow.py

@@ -8,6 +8,7 @@
 from typing import Any, AsyncIterator, Dict, List, Optional, Sequence, Tuple, cast
 
 from langchain_core.runnables import RunnableConfig
+from redis import WatchError
 from redisvl.index import AsyncSearchIndex
 from redisvl.query import FilterQuery
 from redisvl.query.filter import Num, Tag
@@ -317,9 +318,9 @@ async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
 
         # Ensure metadata matches CheckpointMetadata type
         sanitized_metadata = {
-            k.replace("\u0000", ""): v.replace("\u0000", "")
-            if isinstance(v, str)
-            else v
+            k.replace("\u0000", ""): (
+                v.replace("\u0000", "") if isinstance(v, str) else v
+            )
             for k, v in metadata_dict.items()
         }
         metadata = cast(CheckpointMetadata, sanitized_metadata)
@@ -386,37 +387,32 @@ async def aput_writes(
             }
             writes_objects.append(write_obj)
 
-            # For each write, check existence and then perform appropriate operation
-            async with self.checkpoints_index.client.pipeline(
-                transaction=False
-            ) as pipeline:
-                for write_obj in writes_objects:
-                    key = self._make_redis_checkpoint_writes_key(
-                        thread_id,
-                        checkpoint_ns,
-                        checkpoint_id,
-                        task_id,
-                        write_obj["idx"],
-                    )
-
-                    # First check if key exists
-                    key_exists = await self._redis.exists(key) == 1
-
-                    if all(w[0] in WRITES_IDX_MAP for w in writes):
-                        # UPSERT case - only update specific fields
-                        if key_exists:
-                            # Update only channel, type, and blob fields
-                            pipeline.json().set(key, "$.channel", write_obj["channel"])
-                            pipeline.json().set(key, "$.type", write_obj["type"])
-                            pipeline.json().set(key, "$.blob", write_obj["blob"])
+            upsert_case = all(w[0] in WRITES_IDX_MAP for w in writes)
+            for write_obj in writes_objects:
+                key = self._make_redis_checkpoint_writes_key(
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint_id,
+                    task_id,
+                    write_obj["idx"],
+                )
+                if upsert_case:
+                    async def tx(pipe, key=key, write_obj=write_obj):
+                        exists = await pipe.exists(key)
+                        if exists:
+                            await pipe.json().set(
+                                key, "$.channel", write_obj["channel"]
+                            )
+                            await pipe.json().set(key, "$.type", write_obj["type"])
+                            await pipe.json().set(key, "$.blob", write_obj["blob"])
                         else:
-                            # For new records, set the complete object
-                            pipeline.json().set(key, "$", write_obj)
-                    else:
-                        # INSERT case - only insert if doesn't exist
-                        pipeline.json().set(key, "$", write_obj)
+                            await pipe.json().set(key, "$", write_obj)
 
-                await pipeline.execute()
+                    await self._redis.transaction(tx, key)
+                else:
+                    # Unlike AsyncRedisSaver, the shallow implementation always overwrites
+                    # the full object, so we don't check for key existence here.
+                    await self._redis.json().set(key, "$", write_obj)
 
     async def aget_channel_values(
         self, thread_id: str, checkpoint_ns: str, checkpoint_id: str