Merge pull request #1395 from langchain-ai/nc/19aug/put-writes-idempotent

checkpoint-*: In put_writes clear any previously saved writes for this task_id

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -316,7 +316,15 @@ class PostgresSaver(BasePostgresSaver):
             writes (List[Tuple[str, Any]]): List of writes to store.
             task_id (str): Identifier for the task creating the writes.
         """
-        with self._cursor() as cur:
+        with self._cursor(pipeline=True) as cur:
+            cur.execute(
+                self.DELETE_WRITES_SQL,
+                config["configurable"]["thread_id"],
+                config["configurable"]["checkpoint_ns"],
+                config["configurable"]["checkpoint_id"],
+                task_id,
+                len(writes),
+            )
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 self._dump_writes(

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -272,7 +272,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
-        async with self._cursor() as cur:
+        async with self._cursor(pipeline=True) as cur:
+            await cur.execute(
+                self.DELETE_WRITES_SQL,
+                config["configurable"]["thread_id"],
+                config["configurable"]["checkpoint_ns"],
+                config["configurable"]["checkpoint_id"],
+                task_id,
+                len(writes),
+            )
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 await asyncio.to_thread(

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -105,6 +105,15 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
 """
 
+DELETE_WRITES_SQL = """
+    DELETE FROM checkpoint_writes
+    WHERE thread_id = %s
+    AND checkpoint_ns = %s
+    AND checkpoint_id = %s
+    AND task_id = %s
+    AND idx >= %s
+"""
+
 
 class BasePostgresSaver(BaseCheckpointSaver):
     SELECT_SQL = SELECT_SQL
@@ -112,6 +121,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+    DELETE_WRITES_SQL = DELETE_WRITES_SQL
+
     jsonplus_serde = JsonPlusSerializer()
 
     def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint:

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py

@@ -424,6 +424,16 @@ class SqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self.lock, self.cursor() as cur:
+            cur.execute(
+                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
+                (
+                    str(config["configurable"]["thread_id"]),
+                    str(config["configurable"]["checkpoint_ns"]),
+                    str(config["configurable"]["checkpoint_id"]),
+                    task_id,
+                    len(writes),
+                ),
+            )
             cur.executemany(
                 "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py

@@ -432,19 +432,29 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         await self.setup()
-        async with self.conn.executemany(
-            "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
-            [
+        async with self.conn.cursor() as cur:
+            await cur.execute(
+                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
                 (
                     str(config["configurable"]["thread_id"]),
                     str(config["configurable"]["checkpoint_ns"]),
                     str(config["configurable"]["checkpoint_id"]),
                     task_id,
-                    idx,
-                    channel,
-                    *self.serde.dumps_typed(value),
-                )
-                for idx, (channel, value) in enumerate(writes)
-            ],
-        ):
-            await self.conn.commit()
+                    len(writes),
+                ),
+            )
+            await cur.executemany(
+                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                [
+                    (
+                        str(config["configurable"]["thread_id"]),
+                        str(config["configurable"]["checkpoint_ns"]),
+                        str(config["configurable"]["checkpoint_id"]),
+                        task_id,
+                        idx,
+                        channel,
+                        *self.serde.dumps_typed(value),
+                    )
+                    for idx, (channel, value) in enumerate(writes)
+                ],
+            )

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -288,6 +288,7 @@ class MemorySaver(
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
         key = (thread_id, checkpoint_ns, checkpoint_id)
+        self.writes[key] = [w for w in self.writes[key] if w[0] != task_id]
         self.writes[key].extend(
             [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
         )