From eb593d47ddb4226d40fbd1ff3630a85511264ff7 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 3 Dec 2024 16:37:02 -0800
Subject: [PATCH] Fix writes for task being saved against next checkpoint id

---
 libs/langgraph/langgraph/pregel/loop.py   | 9 ++++-----
 libs/langgraph/langgraph/pregel/runner.py | 4 ++++
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index cb141c4a1..678e355ab 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -290,16 +290,15 @@ class PregelLoop(LoopProtocol):
         if self.checkpointer_put_writes is not None:
             self.submit(
                 self.checkpointer_put_writes,
-                {
-                    **self.checkpoint_config,
-                    CONF: {
-                        **self.checkpoint_config[CONF],
+                patch_configurable(
+                    self.checkpoint_config,
+                    {
                         CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
                             CONFIG_KEY_CHECKPOINT_NS, ""
                         ),
                         CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
                     },
-                },
+                ),
                 writes,
                 task_id,
             )
diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py
index 41c9e7e69..5ba0209c8 100644
--- a/libs/langgraph/langgraph/pregel/runner.py
+++ b/libs/langgraph/langgraph/pregel/runner.py
@@ -423,6 +423,10 @@ class PregelRunner:
                 break
             # give control back to the caller
             yield
+        # wait for pending done callbacks
+        # if a 2nd future finishes while `wait` is returning, it's possible
+        # that done callbacks for the 2nd future aren't called until next tick
+        await asyncio.sleep(0)
         # cancel waiter task
         for fut in futures:
             fut.cancel()