Fix semantics of put_writes/list (#1436)

* Fix semantics of put_writes/list

- put_writes(error) should not prevent saving future successful if task is retried successfully
- put_writes(writes) should be a no-op if non-error writes already exist for that task (this prevents tasks executed more than once from modifying writes previously saved / acted on)
- checkpoints should not include channel default values (ie those without a version)
- list() should fetch and return writes for each checkpoint

* Lint

* Rm print

* Fix import

* Lint

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

@@ -150,6 +150,7 @@ class PostgresSaver(BasePostgresSaver):
                 }
                 if value["parent_checkpoint_id"]
                 else None,
+                self._load_writes(value["pending_writes"]),
             )
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
@@ -317,16 +318,6 @@ class PostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         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

@@ -135,6 +135,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 }
                 if value["parent_checkpoint_id"]
                 else None,
+                await asyncio.to_thread(self._load_writes, value["pending_writes"]),
             )
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
@@ -273,16 +274,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         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

@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
@@ -105,15 +106,6 @@ 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
@@ -121,7 +113,6 @@ 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()
 
@@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )