Fix writes for task being saved against next checkpoint id

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,
             )

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()