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