From f51e7ea9a407d35e7853e1182878eeec1a997d5e Mon Sep 17 00:00:00 2001
From: vbarda <vadym@langchain.dev>
Date: Wed, 21 Aug 2024 15:38:30 -0400
Subject: [PATCH] pass pending writes in checkpointers

---
 .../langgraph/checkpoint/postgres/__init__.py    |  1 +
 .../langgraph/checkpoint/postgres/aio.py         |  1 +
 .../langgraph/checkpoint/sqlite/__init__.py      | 16 +++++++++++++++-
 .../langgraph/checkpoint/sqlite/aio.py           | 12 ++++++++++++
 .../langgraph/checkpoint/memory/__init__.py      |  4 ++++
 5 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
index 02d5880a0..1da25b35f 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+++ b/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]:
diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
index 7ddb81237..7be1f36dc 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+++ b/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]:
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
index 9bf4139ad..d67454034 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
@@ -318,7 +318,9 @@ class SqliteSaver(BaseCheckpointSaver):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        with self.cursor(transaction=False) as cur:
+        with self.cursor(transaction=False) as cur, self.cursor(
+            transaction=False
+        ) as writes_cur:
             cur.execute(query, param_values)
             for (
                 thread_id,
@@ -329,6 +331,14 @@ class SqliteSaver(BaseCheckpointSaver):
                 checkpoint,
                 metadata,
             ) in cur:
+                writes_cur.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (
+                        thread_id,
+                        checkpoint_ns,
+                        checkpoint_id,
+                    ),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -350,6 +360,10 @@ class SqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        for task_id, channel, type, value in writes_cur
+                    ],
                 )
 
     def put(
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
index 56d14f613..364adf063 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
@@ -346,6 +346,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                 checkpoint,
                 metadata,
             ) in cursor:
+                writes_cur = await self.conn.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (
+                        thread_id,
+                        checkpoint_ns,
+                        checkpoint_id,
+                    ),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -367,6 +375,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        async for task_id, channel, type, value in writes_cur
+                    ],
                 )
 
     async def aput(
diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py
index 6cc1a3b14..6b3714f05 100644
--- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py
@@ -206,6 +206,7 @@ class MemorySaver(
                     elif limit is not None:
                         limit -= 1
 
+                    writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
                     yield CheckpointTuple(
                         config={
                             "configurable": {
@@ -215,6 +216,9 @@ class MemorySaver(
                             }
                         },
                         checkpoint=self.serde.loads_typed(checkpoint),
+                        pending_writes=[
+                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                        ],
                         metadata=metadata,
                         parent_config={
                             "configurable": {