codekingpro/portable-devtools

1from __future__ import annotations2 3import asyncio4import concurrent.futures5import inspect6import threading7import time8import weakref9from collections.abc import (10    AsyncIterator,11    Awaitable,12    Callable,13    Collection,14    Iterable,15    Iterator,16    Mapping,17    Sequence,18)19from functools import partial20from typing import (21    Any,22    Generic,23    TypeVar,24    cast,25)26 27from langchain_core.callbacks import Callbacks28 29from langgraph._internal._constants import (30    CONF,31    CONFIG_KEY_CALL,32    CONFIG_KEY_SCRATCHPAD,33    ERROR,34    ERROR_SOURCE_NODE,35    INTERRUPT,36    NO_WRITES,37    RESUME,38    RETURN,39)40from langgraph._internal._future import chain_future, run_coroutine_threadsafe41from langgraph._internal._scratchpad import PregelScratchpad42from langgraph._internal._typing import MISSING43from langgraph.constants import TAG_HIDDEN44from langgraph.errors import GraphBubbleUp, GraphInterrupt45from langgraph.pregel._algo import Call46from langgraph.pregel._executor import Submit47from langgraph.pregel._retry import arun_with_retry, run_with_retry48from langgraph.types import (49    CachePolicy,50    PregelExecutableTask,51    RetryPolicy,52    TimeoutPolicy,53)54 55F = TypeVar("F", concurrent.futures.Future, asyncio.Future)56E = TypeVar("E", threading.Event, asyncio.Event)57 58# List of filenames to exclude from exception traceback59# Note: Frames will be removed if they are the last frame in traceback, recursively60EXCLUDED_FRAME_FNAMES = (61    "langgraph/pregel/retry.py",62    "langgraph/pregel/runner.py",63    "langgraph/pregel/executor.py",64    "langgraph/utils/runnable.py",65    "langchain_core/runnables/config.py",66    "concurrent/futures/thread.py",67    "concurrent/futures/_base.py",68)69 70SKIP_RERAISE_SET: weakref.WeakSet[concurrent.futures.Future | asyncio.Future] = (71    weakref.WeakSet()72)73 74 75class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):76    event: E77    callback: weakref.ref[Callable[[PregelExecutableTask, BaseException | None], None]]78    # Stop condition is injected by PregelRunner instead of hard-coded here.79    # This lets the runner treat graph-error-handled exceptions as non-fatal80    # so `on_done` does not trigger an early stop for those futures.81    should_stop: Callable[[set[F]], bool]82    counter: int83    done: set[F]84    lock: threading.Lock85 86    def __init__(87        self,88        event: E,89        callback: weakref.ref[90            Callable[[PregelExecutableTask, BaseException | None], None]91        ],92        should_stop: Callable[[set[F]], bool],93        future_type: type[F],94        # used for generic typing, newer py supports FutureDict[...](...)95    ) -> None:96        super().__init__()97        self.lock = threading.Lock()98        self.event = event99        self.callback = callback100        self.should_stop = should_stop101        self.counter = 0102        self.done: set[F] = set()103 104    def __setitem__(105        self,106        key: F,107        value: PregelExecutableTask | None,108    ) -> None:109        super().__setitem__(key, value)  # type: ignore[index]110        if value is not None:111            with self.lock:112                self.event.clear()113                self.counter += 1114            key.add_done_callback(partial(self.on_done, value))115 116    def on_done(117        self,118        task: PregelExecutableTask,119        fut: F,120    ) -> None:121        # Called automatically by future.add_done_callback registered in __setitem__.122        try:123            if cb := self.callback():124                cb(task, _exception(fut))125        finally:126            with self.lock:127                self.done.add(fut)128                self.counter -= 1129                # Wake waiter when all tracked futures are done, or when runner-level130                # stop condition is met (for example, a non-handled fatal exception).131                if self.counter == 0 or self.should_stop(self.done):132                    self.event.set()133 134 135class PregelRunner:136    """Responsible for executing a set of Pregel tasks concurrently, committing137    their writes, yielding control to caller when there is output to emit, and138    interrupting other tasks if appropriate."""139 140    def __init__(141        self,142        *,143        submit: weakref.ref[Submit],144        put_writes: weakref.ref[Callable[[str, Sequence[tuple[str, Any]]], None]],145        use_astream: bool = False,146        node_finished: Callable[[str], None] | None = None,147        node_error_handler_map: Mapping[str, str] | None = None,148        schedule_error_handler: Callable[149            [PregelExecutableTask, BaseException], PregelExecutableTask | None150        ]151        | None = None,152        aschedule_error_handler: Callable[153            [PregelExecutableTask, BaseException],154            Awaitable[PregelExecutableTask | None],155        ]156        | None = None,157    ) -> None:158        self.submit = submit159        self.put_writes = put_writes160        self.use_astream = use_astream161        self.node_finished = node_finished162        self.node_error_handler_map = dict(node_error_handler_map or {})163        self.error_handler_nodes = set(self.node_error_handler_map.values())164        self.schedule_error_handler = schedule_error_handler165        self.aschedule_error_handler = aschedule_error_handler166        # Exception object ids that are already routed to graph-level error handler.167        # These ids are consulted by stop/panic checks to avoid re-raising handled168        # exceptions via the normal fatal path in the same run.169        self._handled_exception_ids: set[int] = set()170 171    def _should_route_to_error_handler(self, task: PregelExecutableTask) -> bool:172        if task.name in self.error_handler_nodes:173            return False174        return task.name in self.node_error_handler_map175 176    def tick(177        self,178        tasks: Iterable[PregelExecutableTask],179        *,180        reraise: bool = True,181        timeout: float | None = None,182        retry_policy: Sequence[RetryPolicy] | None = None,183        get_waiter: Callable[[], concurrent.futures.Future[None]] | None = None,184        schedule_task: Callable[185            [PregelExecutableTask, int, Call | None],186            PregelExecutableTask | None,187        ],188    ) -> Iterator[None]:189        tasks = tuple(tasks)190        futures = FuturesDict(191            callback=weakref.WeakMethod(self.commit),192            event=threading.Event(),193            should_stop=partial(194                _should_stop_others, handled_exception_ids=self._handled_exception_ids195            ),196            future_type=concurrent.futures.Future,197        )198        # give control back to the caller199        yield200        # fast path if single task with no timeout and no waiter201        if len(tasks) == 0:202            return203        elif len(tasks) == 1 and timeout is None and get_waiter is None:204            t = tasks[0]205            scheduled_error_handler = False206            try:207                run_with_retry(208                    t,209                    retry_policy,210                    configurable={211                        CONFIG_KEY_CALL: partial(212                            _call,213                            weakref.ref(t),214                            retry_policy=retry_policy,215                            futures=weakref.ref(futures),216                            schedule_task=schedule_task,217                            submit=self.submit,218                        ),219                    },220                )221                self.commit(t, None)222            except Exception as exc:223                self.commit(t, exc)224                if (225                    not isinstance(exc, GraphBubbleUp)226                    and self._should_route_to_error_handler(t)227                    and self.schedule_error_handler is not None228                ):229                    self._handled_exception_ids.add(id(exc))230                    if handler_task := self.schedule_error_handler(t, exc):231                        tasks = (handler_task,)232                        scheduled_error_handler = True233                        # Continue to the regular scheduling path for handler execution.234                if reraise and futures:235                    if id(exc) not in self._handled_exception_ids:236                        # will be re-raised after futures are done237                        fut: concurrent.futures.Future = concurrent.futures.Future()238                        fut.set_exception(exc)239                        futures.done.add(fut)240                elif reraise and id(exc) not in self._handled_exception_ids:241                    if tb := exc.__traceback__:242                        while tb.tb_next is not None and any(243                            tb.tb_frame.f_code.co_filename.endswith(name)244                            for name in EXCLUDED_FRAME_FNAMES245                        ):246                            tb = tb.tb_next247                        exc.__traceback__ = tb248                    raise249            if not futures and not scheduled_error_handler:250                # maybe `t` scheduled another task251                return252            else:253                if not scheduled_error_handler:254                    tasks = ()  # don't reschedule this task255        # add waiter task if requested256        if get_waiter is not None:257            futures[get_waiter()] = None258        # schedule tasks259        for t in tasks:260            fut = self.submit()(  # type: ignore[misc]261                run_with_retry,262                t,263                retry_policy,264                configurable={265                    CONFIG_KEY_CALL: partial(266                        _call,267                        weakref.ref(t),268                        retry_policy=retry_policy,269                        futures=weakref.ref(futures),270                        schedule_task=schedule_task,271                        submit=self.submit,272                    ),273                },274                __reraise_on_exit__=reraise,275            )276            futures[fut] = t277        # execute tasks, and wait for one to fail or all to finish.278        # each task is independent from all other concurrent tasks279        # yield updates/debug output as each task finishes280        end_time = timeout + time.monotonic() if timeout else None281        handled_futures: set[concurrent.futures.Future[Any]] = set()282        while len(futures) > (1 if get_waiter is not None else 0):283            done, inflight = concurrent.futures.wait(284                futures,285                return_when=concurrent.futures.FIRST_COMPLETED,286                timeout=(max(0, end_time - time.monotonic()) if end_time else None),287            )288            if not done:289                break  # timed out290            done_for_stop: set[concurrent.futures.Future[Any]] = set()291            for fut in done:292                task = futures.pop(fut)293                if task is None:294                    # waiter task finished, schedule another295                    if inflight and get_waiter is not None:296                        futures[get_waiter()] = None297                elif (298                    (task_exc := _exception(fut))299                    and self._should_route_to_error_handler(task)300                    and not isinstance(task_exc, GraphBubbleUp)301                ):302                    self._handled_exception_ids.add(id(task_exc))303                    SKIP_RERAISE_SET.add(fut)304                    handled_futures.add(fut)305                    if self.schedule_error_handler is not None:306                        if handler_task := self.schedule_error_handler(task, task_exc):307                            handler_fut = self.submit()(  # type: ignore[misc]308                                run_with_retry,309                                handler_task,310                                retry_policy,311                                configurable={312                                    CONFIG_KEY_CALL: partial(313                                        _call,314                                        weakref.ref(handler_task),315                                        retry_policy=retry_policy,316                                        futures=weakref.ref(futures),317                                        schedule_task=schedule_task,318                                        submit=self.submit,319                                    ),320                                },321                                __reraise_on_exit__=reraise,322                            )323                            futures[handler_fut] = handler_task324                else:325                    done_for_stop.add(fut)326            else:327                # remove references to loop vars328                del fut, task329            # maybe stop other tasks330            if _should_stop_others(331                done_for_stop, handled_exception_ids=self._handled_exception_ids332            ):333                break334            # give control back to the caller335            yield336        # wait for done callbacks337        futures.event.wait(338            timeout=(max(0, end_time - time.monotonic()) if end_time else None)339        )340        # give control back to the caller341        yield342        # panic on failure or timeout343        try:344            _panic_or_proceed(345                futures.done.union(f for f, t in futures.items() if t is not None),346                panic=reraise,347                handled_exception_ids=self._handled_exception_ids,348                handled_futures=handled_futures,349            )350        except Exception as exc:351            if tb := exc.__traceback__:352                while tb.tb_next is not None and any(353                    tb.tb_frame.f_code.co_filename.endswith(name)354                    for name in EXCLUDED_FRAME_FNAMES355                ):356                    tb = tb.tb_next357                exc.__traceback__ = tb358            raise359 360    async def atick(361        self,362        tasks: Iterable[PregelExecutableTask],363        *,364        reraise: bool = True,365        timeout: float | None = None,366        retry_policy: Sequence[RetryPolicy] | None = None,367        get_waiter: Callable[[], asyncio.Future[None]] | None = None,368        schedule_task: Callable[369            [PregelExecutableTask, int, Call | None],370            Awaitable[PregelExecutableTask | None],371        ],372    ) -> AsyncIterator[None]:373        try:374            loop = asyncio.get_event_loop()375        except RuntimeError:376            loop = asyncio.new_event_loop()377        asyncio.set_event_loop(loop)378        tasks = tuple(tasks)379        futures = FuturesDict(380            callback=weakref.WeakMethod(self.commit),381            event=asyncio.Event(),382            should_stop=partial(383                _should_stop_others, handled_exception_ids=self._handled_exception_ids384            ),385            future_type=asyncio.Future,386        )387        # give control back to the caller388        yield389        # fast path if single task with no waiter and no timeout390        if len(tasks) == 0:391            return392        elif len(tasks) == 1 and get_waiter is None and timeout is None:393            t = tasks[0]394            scheduled_error_handler = False395            try:396                await arun_with_retry(397                    t,398                    retry_policy,399                    stream=self.use_astream,400                    configurable={401                        CONFIG_KEY_CALL: partial(402                            _acall,403                            weakref.ref(t),404                            stream=self.use_astream,405                            retry_policy=retry_policy,406                            futures=weakref.ref(futures),407                            schedule_task=schedule_task,408                            submit=self.submit,409                            loop=loop,410                        ),411                    },412                )413                self.commit(t, None)414            except Exception as exc:415                self.commit(t, exc)416                if (417                    not isinstance(exc, GraphBubbleUp)418                    and self._should_route_to_error_handler(t)419                    and self.aschedule_error_handler is not None420                ):421                    self._handled_exception_ids.add(id(exc))422                    if handler_task := await self.aschedule_error_handler(t, exc):423                        tasks = (handler_task,)424                        scheduled_error_handler = True425                if reraise and futures:426                    if id(exc) not in self._handled_exception_ids:427                        # will be re-raised after futures are done428                        fut: asyncio.Future = loop.create_future()429                        fut.set_exception(exc)430                        futures.done.add(fut)431                elif reraise and id(exc) not in self._handled_exception_ids:432                    if tb := exc.__traceback__:433                        while tb.tb_next is not None and any(434                            tb.tb_frame.f_code.co_filename.endswith(name)435                            for name in EXCLUDED_FRAME_FNAMES436                        ):437                            tb = tb.tb_next438                        exc.__traceback__ = tb439                    raise440            if not futures and not scheduled_error_handler:441                # maybe `t` scheduled another task442                return443            else:444                if not scheduled_error_handler:445                    tasks = ()  # don't reschedule this task446        # add waiter task if requested447        if get_waiter is not None:448            futures[get_waiter()] = None449        # schedule tasks450        for t in tasks:451            fut = cast(452                asyncio.Future,453                self.submit()(  # type: ignore[misc]454                    arun_with_retry,455                    t,456                    retry_policy,457                    stream=self.use_astream,458                    configurable={459                        CONFIG_KEY_CALL: partial(460                            _acall,461                            weakref.ref(t),462                            retry_policy=retry_policy,463                            stream=self.use_astream,464                            futures=weakref.ref(futures),465                            schedule_task=schedule_task,466                            submit=self.submit,467                            loop=loop,468                        ),469                    },470                    __name__=t.name,471                    __cancel_on_exit__=True,472                    __reraise_on_exit__=reraise,473                ),474            )475            futures[fut] = t476        # execute tasks, and wait for one to fail or all to finish.477        # each task is independent from all other concurrent tasks478        # yield updates/debug output as each task finishes479        end_time = timeout + loop.time() if timeout else None480        handled_futures: set[asyncio.Future[Any]] = set()481        while len(futures) > (1 if get_waiter is not None else 0):482            done, inflight = await asyncio.wait(483                futures,484                return_when=asyncio.FIRST_COMPLETED,485                timeout=(max(0, end_time - loop.time()) if end_time else None),486            )487            if not done:488                break  # timed out489            done_for_stop: set[asyncio.Future[Any]] = set()490            for fut in done:491                task = futures.pop(fut)492                if task is None:493                    # waiter task finished, schedule another494                    if inflight and get_waiter is not None:495                        futures[get_waiter()] = None496                elif (497                    (task_exc := _exception(fut))498                    and self._should_route_to_error_handler(task)499                    and not isinstance(task_exc, GraphBubbleUp)500                ):501                    self._handled_exception_ids.add(id(task_exc))502                    SKIP_RERAISE_SET.add(fut)503                    handled_futures.add(fut)504                    if self.aschedule_error_handler is not None:505                        if handler_task := await self.aschedule_error_handler(506                            task, task_exc507                        ):508                            handler_fut = cast(509                                asyncio.Future,510                                self.submit()(  # type: ignore[misc]511                                    arun_with_retry,512                                    handler_task,513                                    retry_policy,514                                    stream=self.use_astream,515                                    configurable={516                                        CONFIG_KEY_CALL: partial(517                                            _acall,518                                            weakref.ref(handler_task),519                                            retry_policy=retry_policy,520                                            stream=self.use_astream,521                                            futures=weakref.ref(futures),522                                            schedule_task=schedule_task,523                                            submit=self.submit,524                                            loop=loop,525                                        ),526                                    },527                                    __name__=handler_task.name,528                                    __cancel_on_exit__=True,529                                    __reraise_on_exit__=reraise,530                                ),531                            )532                            futures[handler_fut] = handler_task533                else:534                    done_for_stop.add(fut)535            else:536                # remove references to loop vars537                del fut, task538            # maybe stop other tasks539            if _should_stop_others(540                done_for_stop, handled_exception_ids=self._handled_exception_ids541            ):542                break543            # give control back to the caller544            yield545        # wait for done callbacks546        await asyncio.wait_for(547            futures.event.wait(),548            timeout=(max(0, end_time - loop.time()) if end_time else None),549        )550        # give control back to the caller551        yield552        # cancel waiter task553        for fut in futures:554            fut.cancel()555        # panic on failure or timeout556        try:557            _panic_or_proceed(558                futures.done.union(f for f, t in futures.items() if t is not None),559                timeout_exc_cls=asyncio.TimeoutError,560                panic=reraise,561                handled_exception_ids=self._handled_exception_ids,562                handled_futures=handled_futures,563            )564        except Exception as exc:565            if tb := exc.__traceback__:566                while tb.tb_next is not None and any(567                    tb.tb_frame.f_code.co_filename.endswith(name)568                    for name in EXCLUDED_FRAME_FNAMES569                ):570                    tb = tb.tb_next571                exc.__traceback__ = tb572            raise573 574    def commit(575        self,576        task: PregelExecutableTask,577        exception: BaseException | None,578    ) -> None:579        if isinstance(exception, asyncio.CancelledError):580            # for cancelled tasks, also save error in task,581            # so loop can finish super-step582            task.writes.append((ERROR, exception))583            self.put_writes()(task.id, task.writes)  # type: ignore[misc]584        elif exception:585            if isinstance(exception, GraphInterrupt):586                # save interrupt to checkpointer587                if exception.args[0]:588                    writes = [(INTERRUPT, exception.args[0])]589                    if resumes := [w for w in task.writes if w[0] == RESUME]:590                        writes.extend(resumes)591                    self.put_writes()(task.id, writes)  # type: ignore[misc]592            elif isinstance(exception, GraphBubbleUp):593                # exception will be raised in _panic_or_proceed594                pass595            else:596                # save error to checkpointer597                task.writes.append((ERROR, exception))598                if self._should_route_to_error_handler(task) and not isinstance(599                    exception, GraphBubbleUp600                ):601                    task.writes.append((ERROR_SOURCE_NODE, task.name))602                    self._handled_exception_ids.add(id(exception))603                self.put_writes()(task.id, task.writes)  # type: ignore[misc]604        else:605            if self.node_finished and (606                task.config is None or TAG_HIDDEN not in task.config.get("tags", [])607            ):608                self.node_finished(task.name)609            if not task.writes:610                # add no writes marker611                task.writes.append((NO_WRITES, None))612            # save task writes to checkpointer613            self.put_writes()(task.id, task.writes)  # type: ignore[misc]614 615 616def _should_stop_others(617    done: set[F],618    *,619    handled_exception_ids: set[int] | None = None,620) -> bool:621    """Check if any task failed, if so, cancel all other tasks.622    GraphInterrupts are not considered failures."""623    for fut in done:624        if fut.cancelled():625            continue626        elif exc := fut.exception():627            if (628                id(exc) not in (handled_exception_ids or set())629                and not isinstance(exc, GraphBubbleUp)630                and fut not in SKIP_RERAISE_SET631            ):632                return True633 634    return False635 636 637def _exception(638    fut: concurrent.futures.Future[Any] | asyncio.Future[Any],639) -> BaseException | None:640    """Return the exception from a future, without raising CancelledError."""641    if fut.cancelled():642        if isinstance(fut, asyncio.Future):643            return asyncio.CancelledError()644        else:645            return concurrent.futures.CancelledError()646    else:647        return fut.exception()648 649 650def _panic_or_proceed(651    futs: set[concurrent.futures.Future] | set[asyncio.Future],652    *,653    timeout_exc_cls: type[Exception] = TimeoutError,654    panic: bool = True,655    handled_exception_ids: set[int] | None = None,656    handled_futures: Collection[concurrent.futures.Future[Any] | asyncio.Future[Any]]657    | None = None,658) -> None:659    """Cancel remaining tasks if any failed, re-raise exception if panic is True."""660    done: set[concurrent.futures.Future[Any] | asyncio.Future[Any]] = set()661    inflight: set[concurrent.futures.Future[Any] | asyncio.Future[Any]] = set()662    for fut in futs:663        if fut.cancelled():664            continue665        elif fut.done():666            done.add(fut)667        else:668            inflight.add(fut)669    interrupts: list[GraphInterrupt] = []670    while done:671        # if any task failed672        fut = done.pop()673        if exc := _exception(fut):674            if fut in (handled_futures or set()):675                continue676            if id(exc) in (handled_exception_ids or set()):677                continue678            # cancel all pending tasks679            while inflight:680                inflight.pop().cancel()681            # raise the exception682            if panic:683                if isinstance(exc, GraphInterrupt):684                    # collect interrupts685                    interrupts.append(exc)686                elif fut not in SKIP_RERAISE_SET:687                    raise exc688    # raise combined interrupts689    if interrupts:690        raise GraphInterrupt(tuple(i for exc in interrupts for i in exc.args[0]))691    if inflight:692        # if we got here means we timed out693        while inflight:694            # cancel all pending tasks695            inflight.pop().cancel()696        # raise timeout error697        raise timeout_exc_cls("Timed out")698 699 700def _call(701    task: weakref.ref[PregelExecutableTask],702    func: Callable[[Any], Awaitable[Any] | Any],703    input: Any,704    *,705    retry_policy: Sequence[RetryPolicy] | None = None,706    cache_policy: CachePolicy | None = None,707    timeout: TimeoutPolicy | None = None,708    callbacks: Callbacks = None,709    futures: weakref.ref[FuturesDict],710    schedule_task: Callable[711        [PregelExecutableTask, int, Call | None], PregelExecutableTask | None712    ],713    submit: weakref.ref[Submit],714) -> concurrent.futures.Future[Any]:715    if inspect.iscoroutinefunction(func):716        raise RuntimeError("In an sync context async tasks cannot be called")717 718    fut: concurrent.futures.Future | None = None719    # schedule PUSH tasks, collect futures720    scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]721    # schedule the next task, if the callback returns one722    if next_task := schedule_task(723        task(),  # type: ignore[arg-type]724        scratchpad.call_counter(),725        Call(726            func,727            input,728            retry_policy=retry_policy,729            cache_policy=cache_policy,730            callbacks=callbacks,731            timeout=timeout,732        ),733    ):734        if fut := next(735            (736                f737                for f, t in list(futures().items())  # type: ignore[union-attr]738                if t is not None and t == next_task.id739            ),740            None,741        ):742            # if the parent task was retried,743            # the next task might already be running744            pass745        elif next_task.writes:746            # if it already ran, return the result747            fut = concurrent.futures.Future()748            ret = next((v for c, v in next_task.writes if c == RETURN), MISSING)749            if ret is not MISSING:750                fut.set_result(ret)751            elif exc := next((v for c, v in next_task.writes if c == ERROR), None):752                fut.set_exception(753                    exc if isinstance(exc, BaseException) else Exception(exc)754                )755            else:756                fut.set_result(None)757        else:758            # schedule the next task759            fut = submit()(  # type: ignore[misc]760                run_with_retry,761                next_task,762                retry_policy,763                configurable={764                    CONFIG_KEY_CALL: partial(765                        _call,766                        weakref.ref(next_task),767                        futures=futures,768                        retry_policy=retry_policy,769                        callbacks=callbacks,770                        schedule_task=schedule_task,771                        submit=submit,772                    ),773                },774                __reraise_on_exit__=False,775                # starting a new task in the next tick ensures776                # updates from this tick are committed/streamed first777                __next_tick__=True,778            )779            # exceptions for call() tasks are raised into the parent task780            # so we should not re-raise at the end of the tick781            SKIP_RERAISE_SET.add(fut)782            futures()[fut] = next_task  # type: ignore[index]783    fut = cast(asyncio.Future | concurrent.futures.Future, fut)784    # return a chained future to ensure commit() callback is called785    # before the returned future is resolved, to ensure stream order etc786    return chain_future(fut, concurrent.futures.Future())787 788 789def _acall(790    task: weakref.ref[PregelExecutableTask],791    func: Callable[[Any], Awaitable[Any] | Any],792    input: Any,793    *,794    retry_policy: Sequence[RetryPolicy] | None = None,795    cache_policy: CachePolicy | None = None,796    timeout: TimeoutPolicy | None = None,797    callbacks: Callbacks = None,798    # injected dependencies799    futures: weakref.ref[FuturesDict],800    schedule_task: Callable[801        [PregelExecutableTask, int, Call | None],802        Awaitable[PregelExecutableTask | None],803    ],804    submit: weakref.ref[Submit],805    loop: asyncio.AbstractEventLoop,806    stream: bool = False,807) -> asyncio.Future[Any] | concurrent.futures.Future[Any]:808    # return a chained future to ensure commit() callback is called809    # before the returned future is resolved, to ensure stream order etc810    try:811        in_async = asyncio.current_task() is not None812    except RuntimeError:813        in_async = False814    # if in async context return an async future, otherwise return a sync future815    if in_async:816        fut: asyncio.Future[Any] | concurrent.futures.Future[Any] = asyncio.Future(817            loop=loop818        )819    else:820        fut = concurrent.futures.Future()821    # schedule the next task822    run_coroutine_threadsafe(823        _acall_impl(824            fut,825            task,826            func,827            input,828            retry_policy=retry_policy,829            cache_policy=cache_policy,830            timeout=timeout,831            callbacks=callbacks,832            futures=futures,833            schedule_task=schedule_task,834            submit=submit,835            loop=loop,836            stream=stream,837        ),838        loop,839        lazy=False,840    )841    return fut842 843 844async def _acall_impl(845    destination: asyncio.Future[Any] | concurrent.futures.Future[Any],846    task: weakref.ref[PregelExecutableTask],847    func: Callable[[Any], Awaitable[Any] | Any],848    input: Any,849    *,850    retry_policy: Sequence[RetryPolicy] | None = None,851    cache_policy: CachePolicy | None = None,852    timeout: TimeoutPolicy | None = None,853    callbacks: Callbacks = None,854    # injected dependencies855    futures: weakref.ref[FuturesDict[asyncio.Future, asyncio.Event]],856    schedule_task: Callable[857        [PregelExecutableTask, int, Call | None],858        Awaitable[PregelExecutableTask | None],859    ],860    submit: weakref.ref[Submit],861    loop: asyncio.AbstractEventLoop,862    stream: bool = False,863) -> None:864    try:865        fut: asyncio.Future | None = None866        # schedule PUSH tasks, collect futures867        scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]868        # schedule the next task, if the callback returns one869        if next_task := await schedule_task(870            task(),  # type: ignore[arg-type]871            scratchpad.call_counter(),872            Call(873                func,874                input,875                retry_policy=retry_policy,876                cache_policy=cache_policy,877                callbacks=callbacks,878                timeout=timeout,879            ),880        ):881            if fut := next(882                (883                    f884                    for f, t in list(futures().items())  # type: ignore[union-attr]885                    if t is not None and t == next_task.id886                ),887                None,888            ):889                # if the parent task was retried,890                # the next task might already be running891                pass892            elif next_task.writes:893                # if it already ran, return the result894                fut = asyncio.Future(loop=loop)895                ret = next((v for c, v in next_task.writes if c == RETURN), MISSING)896                if ret is not MISSING:897                    fut.set_result(ret)898                elif exc := next((v for c, v in next_task.writes if c == ERROR), None):899                    fut.set_exception(900                        exc if isinstance(exc, BaseException) else Exception(exc)901                    )902                else:903                    fut.set_result(None)904            else:905                # schedule the next task906                fut = cast(907                    asyncio.Future,908                    submit()(  # type: ignore[misc]909                        arun_with_retry,910                        next_task,911                        retry_policy,912                        stream=stream,913                        configurable={914                            CONFIG_KEY_CALL: partial(915                                _acall,916                                weakref.ref(next_task),917                                stream=stream,918                                futures=futures,919                                schedule_task=schedule_task,920                                submit=submit,921                                loop=loop,922                            ),923                        },924                        __name__=next_task.name,925                        __cancel_on_exit__=True,926                        __reraise_on_exit__=False,927                        # starting a new task in the next tick ensures928                        # updates from this tick are committed/streamed first929                        __next_tick__=True,930                    ),931                )932                # exceptions for call() tasks are raised into the parent task933                # so we should not re-raise at the end of the tick934                SKIP_RERAISE_SET.add(fut)935                futures()[fut] = next_task  # type: ignore[index]936        if fut is not None:937            chain_future(fut, destination)938        else:939            destination.set_exception(RuntimeError("Task not scheduled"))940    except Exception as exc:941        destination.set_exception(exc)942