Wait until next tick to start send task

libs/langgraph/langgraph/pregel/executor.py

@@ -1,6 +1,7 @@
 import asyncio
 import concurrent.futures
 import sys
+import time
 from contextlib import ExitStack
 from contextvars import copy_context
 from types import TracebackType
@@ -34,6 +35,7 @@ class Submit(Protocol[P, T]):
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]: ...
 
@@ -58,9 +60,13 @@ class BackgroundExecutor(ContextManager):
         __name__: Optional[str] = None,  # currently not used in sync version
         __cancel_on_exit__: bool = False,  # for sync, can cancel only if not started
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
-        task = self.executor.submit(fn, *args, **kwargs)
+        if __next_tick__:
+            task = self.executor.submit(next_tick, fn, *args, **kwargs)
+        else:
+            task = self.executor.submit(fn, *args, **kwargs)
         self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
@@ -137,11 +143,14 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
         if self.semaphore:
             coro = gated(self.semaphore, coro)
+        if __next_tick__:
+            coro = anext_tick(coro)
         if self.context_not_supported:
             task = self.loop.create_task(coro, name=__name__)
         else:
@@ -197,3 +206,15 @@ async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) ->
     """A coroutine that waits for a semaphore before running another coroutine."""
     async with semaphore:
         return await coro
+
+
+def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    """A function that yields control to other threads before running another function."""
+    time.sleep(0)
+    return fn(*args, **kwargs)
+
+
+async def anext_tick(coro: Coroutine[None, None, T]) -> T:
+    """A coroutine that yields control to event loop before running another coroutine."""
+    await asyncio.sleep(0)
+    return await coro

libs/langgraph/langgraph/pregel/runner.py

@@ -133,6 +133,7 @@ class PregelRunner:
                                 CONFIG_KEY_CALL: partial(call, next_task),
                             },
                             __reraise_on_exit__=reraise,
+                            __next_tick__=True,
                         )
                         fut.add_done_callback(partial(self.commit, next_task))
                         futures[fut] = next_task
@@ -228,6 +229,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
+        time.sleep(0)
         # panic on failure or timeout
         _panic_or_proceed(
             done_futures.union(f for f, t in futures.items() if t is not None),

libs/langgraph/tests/test_pregel.py

@@ -5837,6 +5837,7 @@ def test_state_graph_packets(
     @tool()
     def search_api(query: str) -> str:
         """Searches the API for the query."""
+        time.sleep(0.1)
         return f"result for {query}"
 
     tools = [search_api]