diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 64bb6c90e..0a117651d 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -156,10 +156,16 @@ class AsyncBackgroundExecutor(AsyncContextManager): if self.semaphore: coro = gated(self.semaphore, coro) if CONTEXT_NOT_SUPPORTED: - task = run_coroutine_threadsafe(coro, self.loop, name=__name__) + task = run_coroutine_threadsafe( + coro, self.loop, name=__name__, lazy=__next_tick__ + ) else: task = run_coroutine_threadsafe( - coro, self.loop, name=__name__, context=copy_context() + coro, + self.loop, + name=__name__, + context=copy_context(), + lazy=__next_tick__, ) self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__) task.add_done_callback(self.done) diff --git a/libs/langgraph/langgraph/utils/future.py b/libs/langgraph/langgraph/utils/future.py index e61afec91..a311133df 100644 --- a/libs/langgraph/langgraph/utils/future.py +++ b/libs/langgraph/langgraph/utils/future.py @@ -10,6 +10,7 @@ T = TypeVar("T") AnyFuture = Union[asyncio.Future, concurrent.futures.Future] CONTEXT_NOT_SUPPORTED = sys.version_info < (3, 11) +EAGER_NOT_SUPPORTED = sys.version_info < (3, 12) def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop: @@ -142,6 +143,7 @@ def _ensure_future( loop: asyncio.AbstractEventLoop, name: Optional[str] = None, context: Optional[contextvars.Context] = None, + lazy: bool = True, ) -> asyncio.Task[T]: called_wrap_awaitable = False if not asyncio.iscoroutine(coro_or_future): @@ -159,8 +161,12 @@ def _ensure_future( try: if CONTEXT_NOT_SUPPORTED: return loop.create_task(coro_or_future, name=name) - else: + elif EAGER_NOT_SUPPORTED or lazy: return loop.create_task(coro_or_future, name=name, context=context) + else: + return asyncio.eager_task_factory( + loop, coro_or_future, name=name, context=context + ) except RuntimeError: if not called_wrap_awaitable: coro_or_future.close() @@ -180,6 +186,8 @@ def _wrap_awaitable(awaitable: Awaitable[T]) -> Generator[None, None, T]: def run_coroutine_threadsafe( coro: Coroutine[None, None, T], loop: asyncio.AbstractEventLoop, + *, + lazy: bool, name: Optional[str] = None, context: Optional[contextvars.Context] = None, ) -> asyncio.Future[T]: @@ -187,18 +195,23 @@ def run_coroutine_threadsafe( Return a asyncio.Future to access the result. """ - future: asyncio.Future[T] = asyncio.Future(loop=loop) - def callback() -> None: - try: - chain_future( - _ensure_future(coro, loop=loop, name=name, context=context), future - ) - except (SystemExit, KeyboardInterrupt): - raise - except BaseException as exc: - future.set_exception(exc) - raise + if asyncio._get_running_loop() is loop: + return _ensure_future(coro, loop=loop, name=name, context=context, lazy=lazy) + else: + future: asyncio.Future[T] = asyncio.Future(loop=loop) - loop.call_soon_threadsafe(callback, context=context) - return future + def callback() -> None: + try: + chain_future( + _ensure_future(coro, loop=loop, name=name, context=context), + future, + ) + except (SystemExit, KeyboardInterrupt): + raise + except BaseException as exc: + future.set_exception(exc) + raise + + loop.call_soon_threadsafe(callback, context=context) + return future