Fix faulty test and address langgraph checkpoint 2.1.0 issues

Author: tylerhutcherson

langgraph/checkpoint/redis/aio.py

@@ -760,18 +760,30 @@ async def aput_writes(
                         exists = await self._redis.exists(key)
                         if exists:
                             # Update existing key
-                            await self._redis.json().set(key, "$.channel", write_obj["channel"])  # type: ignore[misc, arg-type]
-                            await self._redis.json().set(key, "$.type", write_obj["type"])  # type: ignore[misc, arg-type]
-                            await self._redis.json().set(key, "$.blob", write_obj["blob"])  # type: ignore[misc, arg-type]
+                            pipeline.json().set(
+                                key,
+                                "$.channel",
+                                write_obj["channel"],  # type: ignore[arg-type]
+                            )
+                            pipeline.json().set(
+                                key,
+                                "$.type",
+                                write_obj["type"],  # type: ignore[arg-type]
+                            )
+                            pipeline.json().set(
+                                key,
+                                "$.blob",
+                                write_obj["blob"],  # type: ignore[arg-type]
+                            )
                         else:
                             # Create new key
-                            await self._redis.json().set(key, "$", write_obj)  # type: ignore[misc]
+                            pipeline.json().set(key, "$", write_obj)
                             created_keys.append(key)
                     else:
                         # For non-upsert case, only set if key doesn't exist
                         exists = await self._redis.exists(key)
                         if not exists:
-                            await self._redis.json().set(key, "$", write_obj)  # type: ignore[misc]
+                            pipeline.json().set(key, "$", write_obj)
                             created_keys.append(key)
 
                 # Execute all operations atomically

langgraph/checkpoint/redis/base.py

@@ -383,7 +383,9 @@ def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
         # NOTE: we're using JSON serializer (not msgpack), so we need to remove null characters before writing
         return serialized_metadata.decode().replace("\\u0000", "")
 
-    def get_next_version(self, current: Optional[str], channel: None = None) -> str:
+    def get_next_version(  # type: ignore[override]
+        self, current: Optional[str], channel: ChannelProtocol[Any, Any, Any]
+    ) -> str:
         """Generate next version number."""
         if current is None:
             current_v = 0

poetry.lock

@@ -420,8 +420,9 @@ async def test_aput_writes_interruption(redis_url: str) -> None:
 
             # Either there are no pending writes or they are not the ones we tried to save
             if checkpoint_tuple and checkpoint_tuple.pending_writes:
+                print(checkpoint_tuple.pending_writes, flush=True)
                 for write in checkpoint_tuple.pending_writes:
-                    assert write.channel not in [
+                    assert write[1] not in [
                         "channel1",
                         "channel2",
                     ], "Transaction should have been rolled back"