coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/_loop.py

@@ -0,0 +1,1988 @@
+from __future__ import annotations
+
+import asyncio
+import binascii
+import concurrent.futures
+from collections import defaultdict, deque
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import (
+    AbstractAsyncContextManager,
+    AbstractContextManager,
+    AsyncExitStack,
+    ExitStack,
+)
+from datetime import datetime, timezone
+from inspect import signature
+from types import TracebackType
+from typing import (
+    Any,
+    Literal,
+    TypeVar,
+    cast,
+)
+
+from coco_core.callbacks import AsyncParentRunManager, ParentRunManager
+from coco_core.runnables import RunnableConfig
+from coco_graph.cache.base import BaseCache
+from coco_graph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    BaseCheckpointSaver,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    PendingWrite,
+)
+from coco_graph.store.base import BaseStore
+from typing_extensions import ParamSpec, Self
+
+from coco_graph._internal._config import patch_configurable
+from coco_graph._internal._constants import (
+    CONF,
+    CONFIG_KEY_CHECKPOINT_ID,
+    CONFIG_KEY_CHECKPOINT_MAP,
+    CONFIG_KEY_CHECKPOINT_NS,
+    CONFIG_KEY_REPLAY_STATE,
+    CONFIG_KEY_RESUME_MAP,
+    CONFIG_KEY_RESUMING,
+    CONFIG_KEY_RUNTIME,
+    CONFIG_KEY_SCRATCHPAD,
+    CONFIG_KEY_STREAM,
+    CONFIG_KEY_TASK_ID,
+    CONFIG_KEY_THREAD_ID,
+    ERROR,
+    ERROR_SOURCE_NODE,
+    INPUT,
+    INTERRUPT,
+    NS_END,
+    NS_SEP,
+    NULL_TASK_ID,
+    PUSH,
+    RESUME,
+    TASKS,
+)
+from coco_graph._internal._replay import ReplayState
+from coco_graph._internal._scratchpad import PregelScratchpad
+from coco_graph._internal._typing import EMPTY_SEQ, MISSING
+from coco_graph.callbacks import (
+    GraphInterruptEvent,
+    GraphLifecycleEvent,
+    GraphResumeEvent,
+)
+from coco_graph.channels.base import BaseChannel
+from coco_graph.channels.binop import _get_overwrite
+from coco_graph.channels.delta import DeltaChannel
+from coco_graph.channels.untracked_value import UntrackedValue
+from coco_graph.constants import TAG_HIDDEN
+from coco_graph.errors import (
+    EmptyInputError,
+    GraphInterrupt,
+)
+from coco_graph.managed.base import (
+    ManagedValueMapping,
+    ManagedValueSpec,
+)
+from coco_graph.pregel._algo import (
+    Call,
+    GetNextVersion,
+    PregelTaskWrites,
+    apply_writes,
+    checkpoint_null_version,
+    increment,
+    prepare_next_tasks,
+    prepare_node_error_handler_task,
+    prepare_single_task,
+    sanitize_untracked_values_in_send,
+    should_interrupt,
+    task_path_str,
+)
+from coco_graph.pregel._checkpoint import (
+    achannels_from_checkpoint,
+    channels_from_checkpoint,
+    copy_checkpoint,
+    create_checkpoint,
+    delta_channels_to_snapshot,
+    empty_checkpoint,
+    exit_delta_task_id,
+)
+from coco_graph.pregel._executor import (
+    AsyncBackgroundExecutor,
+    BackgroundExecutor,
+    Submit,
+)
+from coco_graph.pregel._io import (
+    map_command,
+    map_input,
+    map_output_updates,
+    map_output_values,
+    read_channels,
+)
+from coco_graph.pregel._messages import ensure_message_ids
+from coco_graph.pregel._read import PregelNode
+from coco_graph.pregel._utils import get_new_channel_versions, is_xxh3_128_hexdigest
+from coco_graph.pregel.debug import (
+    map_debug_checkpoint,
+    map_debug_task_results,
+    map_debug_tasks,
+)
+from coco_graph.pregel.protocol import StreamChunk, StreamProtocol
+from coco_graph.runtime import RunControl, Runtime
+from coco_graph.types import (
+    All,
+    CachePolicy,
+    Command,
+    Durability,
+    Interrupt,
+    PregelExecutableTask,
+    RetryPolicy,
+    Send,
+    StreamMode,
+)
+
+V = TypeVar("V")
+P = ParamSpec("P")
+
+
+WritesT = Sequence[tuple[str, Any]]
+
+
+def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
+    def __call__(value: StreamChunk) -> None:
+        for stream in streams:
+            if value[1] in stream.modes:
+                stream(value)
+
+    return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
+
+
+class PregelLoop:
+    config: RunnableConfig
+    store: BaseStore | None
+    stream: StreamProtocol | None
+    step: int
+    stop: int
+
+    input: Any | None
+    cache: BaseCache[WritesT] | None
+    checkpointer: BaseCheckpointSaver | None
+    nodes: Mapping[str, PregelNode]
+    specs: Mapping[str, BaseChannel | ManagedValueSpec]
+    input_keys: str | Sequence[str]
+    output_keys: str | Sequence[str]
+    stream_keys: str | Sequence[str]
+    is_replaying: bool
+    is_nested: bool
+    manager: None | AsyncParentRunManager | ParentRunManager
+    interrupt_after: All | Sequence[str]
+    interrupt_before: All | Sequence[str]
+    durability: Durability
+    retry_policy: Sequence[RetryPolicy]
+    cache_policy: CachePolicy | None
+
+    checkpointer_get_next_version: GetNextVersion
+    checkpointer_put_writes: Callable[[RunnableConfig, WritesT, str], Any] | None
+    checkpointer_put_writes_accepts_task_path: bool
+    _checkpointer_put_after_previous: (
+        Callable[
+            [
+                concurrent.futures.Future | None,
+                RunnableConfig,
+                Checkpoint,
+                str,
+                ChannelVersions,
+            ],
+            Any,
+        ]
+        | None
+    )
+    _migrate_checkpoint: Callable[[Checkpoint], None] | None
+    submit: Submit
+    channels: Mapping[str, BaseChannel]
+    # Futures from `checkpointer.put_writes` calls that produced delta-channel
+    # writes. `_checkpointer_put_after_previous` drains this list (swap to a
+    # local `futs` then reset to `[]` and wait/gather) before putting the
+    # next checkpoint, so a checkpoint never becomes durable before the
+    # writes that produced it. Initialised to `[]` in both sync and async
+    # `__enter__`; stays `None` only when no checkpointer.
+    _delta_write_futs: list[Any] | None = None
+
+    # Same pattern as `_delta_write_futs` but for error-handler writes.
+    # When `put_writes` persists an ERROR_SOURCE_NODE marker, the future is
+    # appended here.  `schedule_error_handler` / `aschedule_error_handler`
+    # drain this list so the write is durable before the handler starts.
+    _error_handler_write_futs: list[Any] | None = None
+
+    # Exit-mode accumulator: every delta-channel write produced during this
+    # run (input writes from `_first` + per-superstep writes captured in
+    # `after_tick`). At exit, `_put_exit_delta_writes` filters out channels
+    # that will snapshot, then persists the rest under an anchor parent.
+    # `None` when not in exit mode (so the capture sites are no-ops).
+    # Each tuple is `(step, task_id, channel, value)` — `step` drives the
+    # synthetic step-prefixed task_id used to preserve chronological order
+    # under the saver's `ORDER BY task_id, idx` sorting.
+    _exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
+
+    # Delta channels that saw an Overwrite since the last checkpoint. These
+    # channels must snapshot after live update applies overwrite semantics so
+    # sparse replay starts from the same post-overwrite value.
+    _delta_channels_with_overwrite: set[str]
+
+    # The checkpoint_config that points at the parent loaded at `__enter__`
+    # (or the synthetic-empty checkpoint, on first run). We capture it
+    # eagerly because every `_put_checkpoint` advances `self.checkpoint_config`
+    # to the newly-saved checkpoint's id — by exit time the original parent
+    # config would otherwise be lost. `_put_exit_delta_writes` uses this:
+    # on resumed runs as the anchor for exit delta writes; on first runs
+    # to derive the lazy stub's config (its `checkpoint_id` is the
+    # synthetic-empty id we want the stub persisted under).
+    _initial_checkpoint_config: RunnableConfig
+
+    # True iff the saver actually returned a tuple at `__enter__`. False
+    # on the first-ever run for a thread (no parent persisted yet).
+    # `_put_exit_delta_writes` uses this to decide between anchoring on
+    # the existing parent (True) or creating a lazy stub (False).
+    _has_persisted_parent: bool = False
+
+    managed: ManagedValueMapping
+    checkpoint: Checkpoint
+    checkpoint_id_saved: str
+    checkpoint_ns: tuple[str, ...]
+    checkpoint_config: RunnableConfig
+    checkpoint_metadata: CheckpointMetadata
+    checkpoint_pending_writes: list[PendingWrite]
+    checkpoint_previous_versions: dict[str, str | float | int]
+    prev_checkpoint_config: RunnableConfig | None
+
+    status: Literal[
+        "input",
+        "pending",
+        "done",
+        "draining",
+        "interrupt_before",
+        "interrupt_after",
+        "out_of_steps",
+    ]
+    control: RunControl | None
+    tasks: dict[str, PregelExecutableTask]
+    output: None | dict[str, Any] | Any = None
+    updated_channels: set[str] | None = None
+    _graph_lifecycle_events: deque[GraphLifecycleEvent]
+    _has_graph_lifecycle_callbacks: bool
+
+    # public
+
+    def __init__(
+        self,
+        input: Any | None,
+        *,
+        stream: StreamProtocol | None,
+        config: RunnableConfig,
+        store: BaseStore | None,
+        cache: BaseCache | None,
+        checkpointer: BaseCheckpointSaver | None,
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, BaseChannel | ManagedValueSpec],
+        input_keys: str | Sequence[str],
+        output_keys: str | Sequence[str],
+        stream_keys: str | Sequence[str],
+        trigger_to_nodes: Mapping[str, Sequence[str]],
+        durability: Durability,
+        interrupt_after: All | Sequence[str] = EMPTY_SEQ,
+        interrupt_before: All | Sequence[str] = EMPTY_SEQ,
+        manager: None | AsyncParentRunManager | ParentRunManager = None,
+        migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: CachePolicy | None = None,
+        has_graph_lifecycle_callbacks: bool = False,
+    ) -> None:
+        self.stream = stream
+        self.config = config
+        self.store = store
+        self.step = 0
+        self.stop = 0
+        self.input = input
+        self.checkpointer = checkpointer
+        self.cache = cache
+        self.nodes = nodes
+        self.specs = specs
+        self.input_keys = input_keys
+        self.output_keys = output_keys
+        self.stream_keys = stream_keys
+        self.interrupt_after = interrupt_after
+        self.interrupt_before = interrupt_before
+        self.manager = manager
+        self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
+        self.is_replaying = CONFIG_KEY_CHECKPOINT_ID in config[CONF]
+        self._migrate_checkpoint = migrate_checkpoint
+        self.trigger_to_nodes = trigger_to_nodes
+        self.retry_policy = retry_policy
+        self.cache_policy = cache_policy
+        self.durability = durability
+        self._has_graph_lifecycle_callbacks = has_graph_lifecycle_callbacks
+        self._graph_lifecycle_events = deque()
+        if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
+            self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
+        scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
+        if isinstance(scratchpad, PregelScratchpad):
+            # if count is > 0, append to checkpoint_ns
+            # if count is 0, leave as is
+            if cnt := scratchpad.subgraph_counter():
+                self.config = patch_configurable(
+                    self.config,
+                    {
+                        CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
+                            (
+                                config[CONF][CONFIG_KEY_CHECKPOINT_NS],
+                                str(cnt),
+                            )
+                        )
+                    },
+                )
+        if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
+            self.config = patch_configurable(
+                self.config,
+                {CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None},
+            )
+        if (
+            CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF]
+            and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
+            in self.config[CONF][CONFIG_KEY_CHECKPOINT_MAP]
+        ):
+            self.checkpoint_config = patch_configurable(
+                self.config,
+                {
+                    CONFIG_KEY_CHECKPOINT_ID: self.config[CONF][
+                        CONFIG_KEY_CHECKPOINT_MAP
+                    ][self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]]
+                },
+            )
+        else:
+            self.checkpoint_config = self.config
+        if thread_id := self.checkpoint_config[CONF].get(CONFIG_KEY_THREAD_ID):
+            if not isinstance(thread_id, str):
+                self.checkpoint_config = patch_configurable(
+                    self.checkpoint_config,
+                    {CONFIG_KEY_THREAD_ID: str(thread_id)},
+                )
+        self.checkpoint_ns = (
+            tuple(cast(str, self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP))
+            if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
+            else ()
+        )
+        self.prev_checkpoint_config = None
+        runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
+        self.control = runtime.control if isinstance(runtime, Runtime) else None
+
+    def _push_graph_lifecycle_event(
+        self,
+        kind: Literal["resume", "interrupt"],
+        *,
+        interrupts: tuple[Interrupt, ...] = (),
+    ) -> None:
+        # drain status never reaches lifecycle events: tick() returns False
+        # before pushing, and interrupts are raised through GraphInterrupt
+        if self.status == "draining":
+            raise RuntimeError("Draining status cannot emit lifecycle events")
+        status = self.status
+        if kind == "resume":
+            self._graph_lifecycle_events.append(
+                GraphResumeEvent(
+                    run_id=None,
+                    status=status,
+                    checkpoint_id=self.checkpoint["id"],
+                    checkpoint_ns=self.checkpoint_ns,
+                )
+            )
+        elif kind == "interrupt":
+            self._graph_lifecycle_events.append(
+                GraphInterruptEvent(
+                    run_id=None,
+                    status=status,
+                    checkpoint_id=self.checkpoint["id"],
+                    checkpoint_ns=self.checkpoint_ns,
+                    interrupts=interrupts,
+                )
+            )
+        else:
+            msg = f"Unknown graph lifecycle event type: {kind}"
+            raise AssertionError(msg)
+
+    def _pop_lifecycle_event(self) -> GraphLifecycleEvent | None:
+        if not self._graph_lifecycle_events:
+            return None
+        return self._graph_lifecycle_events.popleft()
+
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        if not writes:
+            return
+        # deduplicate writes to special channels, last write wins
+        if all(w[0] in WRITES_IDX_MAP for w in writes):
+            writes = list({w[0]: w for w in writes}.values())
+        if task_id == NULL_TASK_ID:
+            # writes for the null task are accumulated
+            self.checkpoint_pending_writes = [
+                w
+                for w in self.checkpoint_pending_writes
+                if w[0] != task_id or w[1] not in WRITES_IDX_MAP
+            ]
+            writes_to_save: WritesT = [
+                w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
+            ] + list(writes)
+        else:
+            # remove existing writes for this task
+            self.checkpoint_pending_writes = [
+                w for w in self.checkpoint_pending_writes if w[0] != task_id
+            ]
+            writes_to_save = writes
+
+        # check if any writes are to an UntrackedValue channel
+        if any(
+            isinstance(channel, UntrackedValue) for channel in self.channels.values()
+        ):
+            # we do not persist untracked values in checkpoints
+            writes_to_save = [
+                # sanitize UntrackedValues that are nested within Send packets
+                (
+                    (c, sanitize_untracked_values_in_send(v, self.channels))
+                    if c == TASKS and isinstance(v, Send)
+                    else (c, v)
+                )
+                for c, v in writes_to_save
+                # dont persist UntrackedValue channel writes
+                if not isinstance(self.specs.get(c), UntrackedValue)
+            ]
+
+        # save writes
+        self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
+        # Assign stable IDs to any id=None BaseMessages in DeltaChannel writes
+        # before the background thread serialises them. Without this, reducers
+        # that assign IDs inside apply_writes() race with serialisation and
+        # store id=None, causing get_state() replays to produce a different UUID
+        # on every call.
+        for c, v in writes_to_save:
+            if isinstance(self.specs.get(c), DeltaChannel):
+                ensure_message_ids(v)
+        if self.durability != "exit" and self.checkpointer_put_writes is not None:
+            config = patch_configurable(
+                self.checkpoint_config,
+                {
+                    CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
+                        CONFIG_KEY_CHECKPOINT_NS, ""
+                    ),
+                    CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
+                },
+            )
+            if self.checkpointer_put_writes_accepts_task_path:
+                if hasattr(self, "tasks"):
+                    task = self.tasks.get(task_id)
+                else:
+                    task = None
+                fut = self.submit(
+                    self.checkpointer_put_writes,
+                    config,
+                    writes_to_save,
+                    task_id,
+                    task_path_str(task.path) if task else "",
+                )
+            else:
+                fut = self.submit(
+                    self.checkpointer_put_writes,
+                    config,
+                    writes_to_save,
+                    task_id,
+                )
+            if self._delta_write_futs is not None and any(
+                isinstance(self.specs.get(c), DeltaChannel) for c, _ in writes_to_save
+            ):
+                self._delta_write_futs.append(fut)
+            # ERROR_SOURCE_NODE is only appended by commit() when the task
+            # has an error handler (_should_route_to_error_handler), so this
+            # check naturally limits future collection to those tasks.
+            if self._error_handler_write_futs is not None and any(
+                c == ERROR_SOURCE_NODE for c, _ in writes
+            ):
+                self._error_handler_write_futs.append(fut)
+        # output writes
+        if hasattr(self, "tasks"):
+            self.output_writes(task_id, writes)
+
+    def _put_pending_writes(self) -> None:
+        if self.checkpointer_put_writes is None:
+            return
+        if not self.checkpoint_pending_writes:
+            return
+        # patch config
+        config = patch_configurable(
+            self.checkpoint_config,
+            {
+                CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
+                    CONFIG_KEY_CHECKPOINT_NS, ""
+                ),
+                CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
+            },
+        )
+        # group by task id
+        by_task = defaultdict(list)
+        for task_id, channel, value in self.checkpoint_pending_writes:
+            by_task[task_id].append((channel, value))
+        # submit writes to checkpointer
+        for task_id, writes in by_task.items():
+            if self.checkpointer_put_writes_accepts_task_path and hasattr(
+                self, "tasks"
+            ):
+                task = self.tasks.get(task_id)
+                self.submit(
+                    self.checkpointer_put_writes,
+                    config,
+                    writes,
+                    task_id,
+                    task_path_str(task.path) if task else "",
+                )
+            else:
+                self.submit(
+                    self.checkpointer_put_writes,
+                    config,
+                    writes,
+                    task_id,
+                )
+
+    def accept_push(
+        self, task: PregelExecutableTask, write_idx: int, call: Call | None = None
+    ) -> PregelExecutableTask | None:
+        """Accept a PUSH from a task, potentially returning a new task to start."""
+        checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", ""))
+        null_version = checkpoint_null_version(self.checkpoint)
+        if pushed := cast(
+            PregelExecutableTask | None,
+            prepare_single_task(
+                (PUSH, task.path, write_idx, task.id, call),
+                None,
+                checkpoint=self.checkpoint,
+                checkpoint_id_bytes=checkpoint_id_bytes,
+                checkpoint_null_version=null_version,
+                pending_writes=self.checkpoint_pending_writes,
+                processes=self.nodes,
+                channels=self.channels,
+                managed=self.managed,
+                config=task.config,
+                step=self.step,
+                stop=self.stop,
+                for_execution=True,
+                store=self.store,
+                checkpointer=self.checkpointer,
+                manager=self.manager,
+                retry_policy=self.retry_policy,
+                cache_policy=self.cache_policy,
+            ),
+        ):
+            # produce debug output
+            self._emit("tasks", map_debug_tasks, [pushed])
+            # save the new task
+            self.tasks[pushed.id] = pushed
+            # match any pending writes to the new task
+            if not self.is_replaying:
+                self._reapply_writes_to_succeeded_nodes({pushed.id: pushed})
+            # return the new task, to be started if not run before
+            return pushed
+
+    def schedule_error_handler(
+        self, failed_task: PregelExecutableTask, error: BaseException
+    ) -> PregelExecutableTask | None:
+        raise NotImplementedError
+
+    async def aschedule_error_handler(
+        self, failed_task: PregelExecutableTask, error: BaseException
+    ) -> PregelExecutableTask | None:
+        raise NotImplementedError
+
+    def tick(self) -> bool:
+        """Execute a single iteration of the Pregel loop.
+
+        Returns:
+            True if more iterations are needed.
+        """
+
+        # check if iteration limit is reached
+        if self.step > self.stop:
+            self.status = "out_of_steps"
+            return False
+
+        # prepare next tasks
+        self.tasks = prepare_next_tasks(
+            self.checkpoint,
+            self.checkpoint_pending_writes,
+            self.nodes,
+            self.channels,
+            self.managed,
+            self.config,
+            self.step,
+            self.stop,
+            for_execution=True,
+            manager=self.manager,
+            store=self.store,
+            checkpointer=self.checkpointer,
+            trigger_to_nodes=self.trigger_to_nodes,
+            updated_channels=self.updated_channels,
+            retry_policy=self.retry_policy,
+            cache_policy=self.cache_policy,
+        )
+
+        # produce debug output
+        if self._checkpointer_put_after_previous is not None:
+            self._emit(
+                "checkpoints",
+                map_debug_checkpoint,
+                {
+                    **self.checkpoint_config,
+                    CONF: {
+                        **self.checkpoint_config[CONF],
+                        CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
+                    },
+                },
+                self.channels,
+                self.stream_keys,
+                self.checkpoint_metadata,
+                self.tasks.values(),
+                self.checkpoint_pending_writes,
+                self.prev_checkpoint_config,
+                self.output_keys,
+            )
+
+        # if no more tasks, we're done
+        if not self.tasks:
+            self.status = "done"
+            return False
+
+        if self.control is not None and self.control.drain_requested:
+            self.status = "draining"
+            return False
+
+        # if there are pending writes from a previous loop, apply them
+        if not self.is_replaying and self.checkpoint_pending_writes:
+            self._reapply_writes_to_succeeded_nodes(self.tasks)
+            self._resume_error_handlers_if_applicable()
+
+        # before execution, check if we should interrupt
+        if self.interrupt_before and should_interrupt(
+            self.checkpoint, self.interrupt_before, self.tasks.values()
+        ):
+            self.status = "interrupt_before"
+            raise GraphInterrupt()
+
+        # produce debug output
+        self._emit("tasks", map_debug_tasks, self.tasks.values())
+
+        # print output for any tasks we applied previous writes to
+        for task in self.tasks.values():
+            if task.writes:
+                self.output_writes(task.id, task.writes, cached=True)
+
+        return True
+
+    def after_tick(self) -> None:
+        # finish superstep
+        writes = [w for t in self.tasks.values() for w in t.writes]
+        self._delta_channels_with_overwrite.update(
+            ch
+            for ch, v in writes
+            if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
+        )
+        # all tasks have finished
+        self.updated_channels = apply_writes(
+            self.checkpoint,
+            self.channels,
+            self.tasks.values(),
+            self.checkpointer_get_next_version,
+            self.trigger_to_nodes,
+        )
+        # produce values output
+        if not self.updated_channels.isdisjoint(
+            (self.output_keys,)
+            if isinstance(self.output_keys, str)
+            else self.output_keys
+        ):
+            self._emit(
+                "values", map_output_values, self.output_keys, writes, self.channels
+            )
+        # capture delta-channel writes for exit-mode accumulator before clearing
+        if self._exit_delta_writes is not None:
+            for tid, ch, v in self.checkpoint_pending_writes:
+                if isinstance(self.specs.get(ch), DeltaChannel):
+                    self._exit_delta_writes.append((self.step, tid, ch, v))
+        # clear pending writes
+        self.checkpoint_pending_writes.clear()
+        # only replay (re-execute) done tasks on the first tick
+        self.is_replaying = False
+        # save checkpoint
+        self._put_checkpoint({"source": "loop"})
+        # after execution, check if we should interrupt
+        if self.interrupt_after and should_interrupt(
+            self.checkpoint, self.interrupt_after, self.tasks.values()
+        ):
+            self.status = "interrupt_after"
+            raise GraphInterrupt()
+        # unset resuming flag
+        self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
+
+    def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        raise NotImplementedError
+
+    async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        raise NotImplementedError
+
+    # private
+
+    def _reapply_writes_to_succeeded_nodes(
+        self, tasks: Mapping[str, PregelExecutableTask]
+    ) -> None:
+        """Restore successful channel writes from checkpoint to in-memory tasks.
+
+        Skips control signals (ERROR, ERROR_SOURCE_NODE, INTERRUPT, RESUME)
+        so that failed/interrupted tasks remain with empty writes and will be
+        re-executed (or routed to error handlers) by the runner.
+        """
+        for tid, k, v in self.checkpoint_pending_writes:
+            if k in (ERROR, ERROR_SOURCE_NODE, INTERRUPT, RESUME):
+                continue
+            if task := tasks.get(tid):
+                task.writes.append((k, v))
+
+    def _resume_error_handlers_if_applicable(self) -> None:
+        """On resume, schedule error handlers for tasks that failed in a prior run.
+
+        Called right after ``_reapply_writes_to_succeeded_nodes`` during ``tick()``.
+        At that point, ``_reapply_writes_to_succeeded_nodes`` has already skipped
+        ERROR / ERROR_SOURCE_NODE writes, so a previously-failed task still has
+        empty ``writes``.  Without intervention the runner (which executes only
+        tasks where ``not t.writes``) would re-run the original node.
+
+        This method prevents that re-execution for nodes that have an error
+        handler:
+
+        1. Scan ``checkpoint_pending_writes`` for ERROR_SOURCE_NODE markers
+           persisted by a prior ``commit()``.  Each marker means "this task
+           already failed and was routed to an error handler".
+        2. For each such task, write ``(ERROR, error)`` into ``task.writes``
+           so the task is no longer empty — the runner will skip it.
+        3. Prepare a fresh error-handler task and add it to ``self.tasks``.
+           Because the handler task starts with empty ``writes``, the runner
+           will pick it up and execute it.
+        """
+        # Phase 1: collect task-ids that have ERROR_SOURCE_NODE + ERROR pairs.
+        failed: dict[str, BaseException] = {}
+        for tid, chan, val in self.checkpoint_pending_writes:
+            if chan == ERROR_SOURCE_NODE:
+                error = next(
+                    (
+                        v
+                        for t, c, v in self.checkpoint_pending_writes
+                        if t == tid and c == ERROR
+                    ),
+                    None,
+                )
+                if error is not None:
+                    failed[tid] = error
+        # Phase 2: mark originals as done, schedule handler tasks.
+        for task_id, error in failed.items():
+            task = self.tasks.get(task_id)
+            if task is None:
+                continue
+            handler_node = self.nodes[task.name].error_handler_node
+            if not handler_node:
+                continue
+            # Non-empty writes → runner's `not t.writes` filter skips this task.
+            task.writes.append((ERROR, error))
+            # The handler task starts with empty writes → runner will execute it.
+            handler_task = prepare_node_error_handler_task(
+                task,
+                handler_node_name=handler_node,
+                failed_error=error,
+                checkpoint=self.checkpoint,
+                pending_writes=self.checkpoint_pending_writes,
+                processes=self.nodes,
+                channels=self.channels,
+                managed=self.managed,
+                config=task.config,
+                step=self.step,
+                stop=self.stop,
+                store=self.store,
+                checkpointer=self.checkpointer,
+                manager=self.manager,
+                retry_policy=self.retry_policy,
+                cache_policy=self.cache_policy,
+            )
+            if handler_task is not None:
+                self.tasks[handler_task.id] = handler_task
+
+    def _pending_interrupts(self) -> set[str]:
+        """Return the set of interrupt ids that are pending without corresponding resume values."""
+        # mapping of task ids to interrupt ids
+        pending_interrupts: dict[str, str] = {}
+
+        # set of resume task ids
+        pending_resumes: set[str] = set()
+
+        for task_id, write_type, value in self.checkpoint_pending_writes:
+            if write_type == INTERRUPT:
+                # interrupts is always a list, but there should only be one element
+                pending_interrupts[task_id] = value[0].id
+            elif write_type == RESUME:
+                pending_resumes.add(task_id)
+
+        resumed_interrupt_ids = {
+            pending_interrupts[task_id]
+            for task_id in pending_resumes
+            if task_id in pending_interrupts
+        }
+
+        # Keep only interrupts whose interrupt_id is not resumed
+        hanging_interrupts: set[str] = {
+            interrupt_id
+            for interrupt_id in pending_interrupts.values()
+            if interrupt_id not in resumed_interrupt_ids
+        }
+
+        return hanging_interrupts
+
+    def _first(
+        self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
+    ) -> set[str] | None:
+        # Resuming from a previous checkpoint requires two things:
+        # 1. A prior checkpoint exists (channel_versions is non-empty)
+        # 2. The input signals continuation (not a fresh run with new input)
+        # For subgraphs, the parent explicitly sets CONFIG_KEY_RESUMING.
+        # For the outer graph, we infer from the input:
+        #   - None input: resume after interrupt (invoke(None, config))
+        #   - Command input: any Command operates on existing state
+        #   - Same run_id: re-entry into an ongoing run (e.g. stream reconnect)
+        configurable = self.config.get(CONF, {})
+        input_is_command = isinstance(self.input, Command)
+        is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
+            configurable.get(
+                CONFIG_KEY_RESUMING,
+                self.input is None
+                or input_is_command
+                or (
+                    not self.is_nested
+                    and self.config.get("metadata", {}).get("run_id")
+                    == self.checkpoint_metadata.get("run_id", MISSING)
+                ),
+            )
+        )
+
+        # When replaying from a specific checkpoint, drop cached RESUME
+        # writes so that interrupt() calls re-fire instead of returning
+        # stale values. But if we're actively resuming, keep them —
+        # multi-interrupt scenarios need previously resolved values preserved.
+        is_time_traveling = self.is_replaying and (
+            # Time-travel to a subgraph checkpoint: the parent sets
+            # RESUMING=True (it can't distinguish time-travel from resume),
+            # so we check if this subgraph's own ns is in checkpoint_map.
+            # Normally the map only has ancestor entries (_algo.py); the
+            # subgraph's own entry only appears via get_state(subgraphs=True).
+            (
+                self.is_nested
+                and configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
+                in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
+            )
+            or not (
+                # Outer graph: resume arrives as Command(resume=...)
+                (input_is_command and cast(Command, self.input).resume is not None)
+                # Subgraphs: resume arrives via config flag from parent
+                # (subgraph input is a Send arg, not a Command)
+                or configurable.get(CONFIG_KEY_RESUMING, False)
+            )
+        )
+        if is_time_traveling:
+            self.checkpoint_pending_writes = [
+                w for w in self.checkpoint_pending_writes if w[1] != RESUME
+            ]
+
+        # map command to writes
+        if input_is_command:
+            if (resume := cast(Command, self.input).resume) is not None:
+                if not self.checkpointer:
+                    raise RuntimeError(
+                        "Cannot use Command(resume=...) without checkpointer"
+                    )
+
+                if resume_is_map := (
+                    isinstance(resume, dict)
+                    and all(is_xxh3_128_hexdigest(k) for k in resume)
+                ):
+                    self.config[CONF][CONFIG_KEY_RESUME_MAP] = resume
+                else:
+                    if len(self._pending_interrupts()) > 1:
+                        raise RuntimeError(
+                            "When there are multiple pending interrupts, you must specify the interrupt id when resuming. "
+                            "Docs: https://docs.langchain.com/oss/python/langgraph/add-human-in-the-loop#resume-multiple-interrupts-with-one-invocation."
+                        )
+
+            writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
+            # group writes by task ID
+            for tid, c, v in map_command(cmd=cast(Command, self.input)):
+                if not (c == RESUME and resume_is_map):
+                    writes[tid].append((c, v))
+            if not writes and not resume_is_map:
+                raise EmptyInputError("Received empty Command input")
+            # save writes
+            for tid, ws in writes.items():
+                self.put_writes(tid, ws)
+        # apply NULL writes
+        if null_writes := [
+            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
+        ]:
+            null_updated_channels = apply_writes(
+                self.checkpoint,
+                self.channels,
+                [PregelTaskWrites((), INPUT, null_writes, [])],
+                self.checkpointer_get_next_version,
+                self.trigger_to_nodes,
+            )
+            if updated_channels is not None:
+                updated_channels.update(null_updated_channels)
+        # proceed past previous checkpoint
+        if is_resuming:
+            self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
+            for k in self.channels:
+                if k in self.checkpoint["channel_versions"]:
+                    version = self.checkpoint["channel_versions"][k]
+                    self.checkpoint["versions_seen"][INTERRUPT][k] = version
+            # When time-traveling (replaying from a specific checkpoint),
+            # save a fork checkpoint so the replayed execution creates a
+            # new branch. Without this, if the execution hits an interrupt
+            # before after_tick() runs, no new checkpoint is created —
+            # the parent's latest checkpoint remains the old one and
+            # subsequent resumes load the wrong state.
+            # Skip for update_state forks (source=update/fork) since they
+            # already have their own fork checkpoint.
+            if is_time_traveling and self.checkpoint_metadata.get("source") not in (
+                "update",
+                "fork",
+            ):
+                # Clear old INTERRUPT writes from the loaded checkpoint.
+                # The fork will have a new checkpoint_id which changes
+                # task IDs — stale interrupt writes would accumulate and
+                # confuse the multiple-interrupt check in future resumes.
+                self.checkpoint_pending_writes = [
+                    w for w in self.checkpoint_pending_writes if w[1] != INTERRUPT
+                ]
+                self._put_checkpoint({"source": "fork"})
+            # produce values output
+            self._emit(
+                "values", map_output_values, self.output_keys, True, self.channels
+            )
+        # map inputs to channel updates
+        elif input_writes := deque(map_input(input_keys, self.input)):
+            # discard any unfinished tasks from previous checkpoint
+            discard_tasks = prepare_next_tasks(
+                self.checkpoint,
+                self.checkpoint_pending_writes,
+                self.nodes,
+                self.channels,
+                self.managed,
+                self.config,
+                self.step,
+                self.stop,
+                for_execution=True,
+                store=None,
+                checkpointer=None,
+                manager=None,
+                updated_channels=updated_channels,
+            )
+            self._delta_channels_with_overwrite.update(
+                c
+                for c, v in input_writes
+                if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0]
+            )
+            # apply input writes
+            updated_channels = apply_writes(
+                self.checkpoint,
+                self.channels,
+                [
+                    *discard_tasks.values(),
+                    PregelTaskWrites((), INPUT, input_writes, []),
+                ],
+                self.checkpointer_get_next_version,
+                self.trigger_to_nodes,
+            )
+            # Input writes go through `apply_writes` directly (above) — they
+            # never enter `checkpoint_pending_writes`, so the after_tick
+            # capture site does not see them. In exit mode, capture them
+            # here so `_exit_delta_writes` includes the input's delta writes
+            # alongside per-superstep writes; otherwise the input would be
+            # lost on read (it's not in final_checkpoint.channel_values for
+            # sub-freq channels, and walks ignore target.pending_writes).
+            if self._exit_delta_writes is not None:
+                for c, v in input_writes:
+                    if isinstance(self.specs.get(c), DeltaChannel):
+                        self._exit_delta_writes.append((self.step, NULL_TASK_ID, c, v))
+            # Persist delta-channel input writes so sub-freq inputs are
+            # recoverable via ancestor walk (mirrors the Command input path).
+            if self.durability != "exit":
+                delta_input = [
+                    (c, v)
+                    for c, v in input_writes
+                    if isinstance(self.specs.get(c), DeltaChannel)
+                ]
+                if delta_input:
+                    self.put_writes(NULL_TASK_ID, delta_input)
+            # save input checkpoint
+            self.updated_channels = updated_channels
+            self._put_checkpoint({"source": "input"})
+        elif CONFIG_KEY_RESUMING not in configurable:
+            raise EmptyInputError(f"Received no input for {input_keys}")
+        # Propagate resuming and replaying flags to subgraphs.
+        if not self.is_nested:
+            # Pass the resolved before-bound checkpoint ID so subgraphs can
+            # find their corresponding checkpoint without re-fetching the
+            # parent. For forks (source=update/fork), use the fork's parent
+            # checkpoint ID since the fork was created after the subgraph's
+            # checkpoints from the original execution.
+            #
+            # Only gate on is_time_traveling (not is_replaying). When the
+            # client resumes with an explicit checkpoint_id that happens to
+            # point at the current head (e.g. CocoGraph Studio sending
+            # `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
+            # is_replaying is True but is_time_traveling is False. In that
+            # case subgraphs should load their latest checkpoint normally,
+            # not go through ReplayState's before-bound lookup which would
+            # miss subgraph checkpoints created during processing of the
+            # current parent step.
+            replay_state: ReplayState | None = None
+            if is_time_traveling:
+                replay_checkpoint_id = self.checkpoint["id"]
+                if (
+                    self.checkpoint_metadata.get("source")
+                    in (
+                        "update",
+                        "fork",
+                    )
+                    and self.prev_checkpoint_config
+                ):
+                    replay_checkpoint_id = self.prev_checkpoint_config[CONF].get(
+                        CONFIG_KEY_CHECKPOINT_ID, replay_checkpoint_id
+                    )
+                replay_state = ReplayState(replay_checkpoint_id)
+            self.config = patch_configurable(
+                self.config,
+                {
+                    CONFIG_KEY_RESUMING: is_resuming,
+                    CONFIG_KEY_REPLAY_STATE: replay_state,
+                },
+            )
+        # set flag
+        self.status = "pending"
+        if is_resuming:
+            self._push_graph_lifecycle_event("resume")
+        return updated_channels
+
+    def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
+        # `is` (object identity) — not `==`. Three of four call sites pass a
+        # fresh dict ({"source":"input"|"loop"|"fork"}); only
+        # `_suppress_interrupt`(will rename to _on_loop_exit soon)
+        # at exit reuses the existing `self.checkpoint_metadata` instance. So
+        # `metadata is self.checkpoint_metadata` is True only on the exit call,
+        # which is what we use to gate exit-only behaviour (skip count-bump,
+        # don't replace metadata). Could be replaced by an explicit
+        # `exiting: bool = False` parameter; left as-is to match the existing
+        # idiom in this file.
+        # TODO: replace with an explicit `exiting: bool = False` parameter.
+        exiting = metadata is self.checkpoint_metadata
+        if exiting and self.checkpoint["id"] == self.checkpoint_id_saved:
+            # checkpoint already saved
+            return
+        # Per-delta-channel counter bookkeeping.
+        #
+        # Each delta channel tracks a (updates, supersteps) tuple:
+        # - `updates` increments only when the channel is written this step.
+        # - `supersteps` increments every superstep regardless.
+        #
+        # `_put_checkpoint` is called once per superstep with a fresh
+        # metadata dict (source="input"|"loop"|"fork") — those are the
+        # intermediate calls that bump counters. In exit mode,
+        # `_suppress_interrupt`(will rename to _on_loop_exit soon)
+        # additionally calls `_put_checkpoint(self.checkpoint_metadata)` AT
+        # EXIT to commit the final checkpoint — this runs *after* the last
+        # intermediate call already counted the last superstep. So the
+        # exit call must NOT bump again or it would double-count the last
+        # superstep.
+        if not exiting:
+            prev_counters = dict(
+                self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
+            )
+            new_counters: dict[str, tuple[int, int]] = {}
+            updated = self.updated_channels or set()
+            for ch_name, ch in self.channels.items():
+                if not isinstance(ch, DeltaChannel):
+                    continue
+                u, s = prev_counters.get(ch_name, (0, 0))
+                s += 1
+                if ch_name in updated:
+                    u += 1
+                new_counters[ch_name] = (u, s)
+            metadata["step"] = self.step
+            metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})
+            self.checkpoint_metadata = metadata
+        else:
+            new_counters = dict(
+                self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
+            )
+        # do checkpoint?
+        do_checkpoint = self._checkpointer_put_after_previous is not None and (
+            exiting or self.durability != "exit"
+        )
+        # create new checkpoint
+        channels_to_snapshot = (
+            delta_channels_to_snapshot(self.channels, new_counters)
+            | self._delta_channels_with_overwrite
+            if do_checkpoint
+            else set()
+        )
+        self.checkpoint = create_checkpoint(
+            self.checkpoint,
+            self.channels if do_checkpoint else None,
+            self.step,
+            id=self.checkpoint["id"] if exiting else None,
+            updated_channels=self.updated_channels,
+            get_next_version=self.checkpointer_get_next_version
+            if do_checkpoint
+            else None,
+            channels_to_snapshot=channels_to_snapshot,
+        )
+        for k in channels_to_snapshot:
+            new_counters[k] = (0, 0)
+        if do_checkpoint:
+            self._delta_channels_with_overwrite.difference_update(channels_to_snapshot)
+        non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)}
+        if non_zero:
+            self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero
+        elif "counters_since_delta_snapshot" in self.checkpoint_metadata:
+            del self.checkpoint_metadata["counters_since_delta_snapshot"]
+        # sanitize TASK channel in the checkpoint before saving (durability=="exit")
+        if TASKS in self.checkpoint["channel_values"] and any(
+            isinstance(channel, UntrackedValue) for channel in self.channels.values()
+        ):
+            sanitized_tasks = [
+                sanitize_untracked_values_in_send(value, self.channels)
+                if isinstance(value, Send)
+                else value
+                for value in self.checkpoint["channel_values"][TASKS]
+            ]
+            self.checkpoint["channel_values"][TASKS] = sanitized_tasks
+        # bail if no checkpointer
+
+        if do_checkpoint and self._checkpointer_put_after_previous is not None:
+            self.prev_checkpoint_config = (
+                self.checkpoint_config
+                if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
+                and self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
+                else None
+            )
+            self.checkpoint_config = {
+                **self.checkpoint_config,
+                CONF: {
+                    **self.checkpoint_config[CONF],
+                    CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
+                        CONFIG_KEY_CHECKPOINT_NS, ""
+                    ),
+                },
+            }
+
+            channel_versions = self.checkpoint["channel_versions"].copy()
+            new_versions = get_new_channel_versions(
+                self.checkpoint_previous_versions, channel_versions
+            )
+            self.checkpoint_previous_versions = channel_versions
+
+            # save it, without blocking
+            # if there's a previous checkpoint save in progress, wait for it
+            # ensuring checkpointers receive checkpoints in order
+            self._put_checkpoint_fut = self.submit(
+                self._checkpointer_put_after_previous,
+                getattr(self, "_put_checkpoint_fut", None),
+                self.checkpoint_config,
+                copy_checkpoint(self.checkpoint),
+                self.checkpoint_metadata,
+                new_versions,
+            )
+            self.checkpoint_config = {
+                **self.checkpoint_config,
+                CONF: {
+                    **self.checkpoint_config[CONF],
+                    CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
+                },
+            }
+        if not exiting:
+            # increment step
+            self.step += 1
+
+    def _put_exit_delta_writes(self) -> None:
+        """Stage stub + accumulated delta writes so final_checkpoint's put
+        waits on them (visibility invariant: both must be durable before
+        final_checkpoint becomes visible to readers).
+
+        Stub is created lazily — only when no persisted parent exists AND at
+        least one delta channel has writes that won't be snapshotted.
+        """
+        if (
+            not self._exit_delta_writes
+            or self.checkpointer is None
+            or self._checkpointer_put_after_previous is None
+            or self.checkpointer_put_writes is None
+        ):
+            return
+
+        counters = dict(
+            self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
+        )
+        channels_to_snapshot = (
+            delta_channels_to_snapshot(self.channels, counters)
+            | self._delta_channels_with_overwrite
+        )
+
+        pending = [
+            (step, tid, ch, v)
+            for (step, tid, ch, v) in self._exit_delta_writes
+            if ch not in channels_to_snapshot
+        ]
+        if not pending:
+            return
+
+        if self._has_persisted_parent:
+            # _initial_checkpoint_config's checkpoint_id is the saved parent's
+            # id (saver returned a real tuple at __enter__).
+            anchor_config = self._initial_checkpoint_config
+        else:
+            stub_cp = empty_checkpoint()
+            stub_cp["id"] = self.checkpoint_id_saved
+            stub_cp["ts"] = datetime.now(timezone.utc).isoformat()
+            # Stub has no parent (checkpoint_id=None in config).
+            stub_put_config = patch_configurable(
+                self._initial_checkpoint_config,
+                {CONFIG_KEY_CHECKPOINT_ID: None},
+            )
+            # Anchor config for put_writes: checkpoint_id = stub's id.
+            anchor_config = patch_configurable(
+                self._initial_checkpoint_config,
+                {CONFIG_KEY_CHECKPOINT_ID: stub_cp["id"]},
+            )
+            self._put_checkpoint_fut = self.submit(
+                self._checkpointer_put_after_previous,
+                getattr(self, "_put_checkpoint_fut", None),
+                stub_put_config,
+                stub_cp,
+                {"step": -2},
+                {},
+            )
+            # Set checkpoint_config so final_checkpoint's _put_checkpoint
+            # sees the stub as its parent.
+            self.checkpoint_config = anchor_config
+
+        # Step-prefixed synthetic task_id preserves chronological superstep
+        # order under the saver's ORDER BY task_id, idx sorting.
+        grouped: dict[tuple[int, str], list[tuple[str, Any]]] = {}
+        for step, tid, ch, v in pending:
+            grouped.setdefault((step, tid), []).append((ch, v))
+        anchor_write_config = patch_configurable(
+            anchor_config,
+            {
+                CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
+                    CONFIG_KEY_CHECKPOINT_NS, ""
+                ),
+                CONFIG_KEY_CHECKPOINT_ID: anchor_config[CONF][CONFIG_KEY_CHECKPOINT_ID],
+            },
+        )
+        for (step, tid), entries in grouped.items():
+            synth_tid = exit_delta_task_id(step, tid)
+            if self.checkpointer_put_writes_accepts_task_path:
+                fut = self.submit(
+                    self.checkpointer_put_writes,
+                    anchor_write_config,
+                    entries,
+                    synth_tid,
+                    "",
+                )
+            else:
+                fut = self.submit(
+                    self.checkpointer_put_writes,
+                    anchor_write_config,
+                    entries,
+                    synth_tid,
+                )
+            if self._delta_write_futs is not None:
+                self._delta_write_futs.append(fut)
+
+    def _suppress_interrupt(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        # persist current checkpoint and writes
+        if self.durability == "exit" and (
+            # if it's a top graph
+            not self.is_nested
+            # or a nested graph with error or interrupt
+            or exc_value is not None
+            # or a nested graph with checkpointer=True
+            or all(NS_END not in part for part in self.checkpoint_ns)
+        ):
+            self._put_exit_delta_writes()
+            self._put_checkpoint(self.checkpoint_metadata)
+            self._put_pending_writes()
+        # suppress interrupt
+        if isinstance(exc_value, GraphInterrupt) and not self.is_nested:
+            interrupt = exc_value
+            interrupts = tuple(interrupt.args[0]) if interrupt.args else ()
+            self._push_graph_lifecycle_event("interrupt", interrupts=interrupts)
+            # emit one last "values" event, with pending writes applied
+            if (
+                hasattr(self, "tasks")
+                and self.checkpoint_pending_writes
+                and any(task.writes for task in self.tasks.values())
+            ):
+                updated_channels = apply_writes(
+                    self.checkpoint,
+                    self.channels,
+                    self.tasks.values(),
+                    self.checkpointer_get_next_version,
+                    self.trigger_to_nodes,
+                )
+                if not updated_channels.isdisjoint(
+                    (self.output_keys,)
+                    if isinstance(self.output_keys, str)
+                    else self.output_keys
+                ):
+                    self._emit(
+                        "values",
+                        map_output_values,
+                        self.output_keys,
+                        [w for t in self.tasks.values() for w in t.writes],
+                        self.channels,
+                    )
+            # emit INTERRUPT if exception is empty (otherwise emitted by put_writes)
+            if not interrupt.args or not interrupt.args[0]:
+                interrupt_payload = interrupt.args[0] if interrupt.args else ()
+                self._emit(
+                    "updates",
+                    lambda: iter([{INTERRUPT: interrupt_payload}]),
+                )
+            # save final output
+            self.output = read_channels(self.channels, self.output_keys)
+            # suppress interrupt
+            return True
+        elif exc_type is None:
+            # save final output
+            self.output = read_channels(self.channels, self.output_keys)
+
+    def _emit(
+        self,
+        mode: StreamMode,
+        values: Callable[P, Iterator[Any]],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> None:
+        if self.stream is None:
+            return
+        debug_remap = mode in ("checkpoints", "tasks") and "debug" in self.stream.modes
+        if mode not in self.stream.modes and not debug_remap:
+            return
+        for v in values(*args, **kwargs):
+            if mode in self.stream.modes:
+                self.stream((self.checkpoint_ns, mode, v))
+            # "debug" mode is "checkpoints" or "tasks" with a wrapper dict
+            if debug_remap:
+                self.stream(
+                    (
+                        self.checkpoint_ns,
+                        "debug",
+                        {
+                            "step": self.step - 1
+                            if mode == "checkpoints"
+                            else self.step,
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                            "type": "checkpoint"
+                            if mode == "checkpoints"
+                            else "task_result"
+                            if "result" in v
+                            else "task",
+                            "payload": v,
+                        },
+                    )
+                )
+
+    def output_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
+        if task := self.tasks.get(task_id):
+            if task.config is not None and TAG_HIDDEN in task.config.get(
+                "tags", EMPTY_SEQ
+            ):
+                return
+            if writes[0][0] == INTERRUPT:
+                # in loop.py we append a bool to the PUSH task paths to indicate
+                # whether or not a call was present. If so,
+                # we don't emit the interrupt as it'll be emitted by the parent
+                if task.path[0] == PUSH and task.path[-1] is True:
+                    return
+                interrupts = [
+                    {
+                        INTERRUPT: tuple(
+                            v
+                            for w in writes
+                            if w[0] == INTERRUPT
+                            for v in (w[1] if isinstance(w[1], Sequence) else (w[1],))
+                        )
+                    }
+                ]
+                stream_modes = self.stream.modes if self.stream else []
+                if "updates" in stream_modes:
+                    self._emit("updates", lambda: iter(interrupts))
+                if "values" in stream_modes:
+                    current_values = read_channels(self.channels, self.output_keys)
+                    # self.output_keys is a sequence, stream chunk contains entire state and interrupts
+                    if isinstance(current_values, dict):
+                        current_values[INTERRUPT] = interrupts[0][INTERRUPT]
+                        self._emit("values", lambda: iter([current_values]))
+                    # self.output_keys is a string, stream chunk contains only interrupts
+                    else:
+                        self._emit("values", lambda: iter(interrupts))
+            elif writes[0][0] != ERROR:
+                self._emit(
+                    "updates",
+                    map_output_updates,
+                    self.output_keys,
+                    [(task, writes)],
+                    cached,
+                )
+            if not cached:
+                self._emit(
+                    "tasks",
+                    map_debug_task_results,
+                    (task, writes),
+                    self.stream_keys,
+                )
+
+
+class SyncPregelLoop(PregelLoop, AbstractContextManager):
+    def __init__(
+        self,
+        input: Any | None,
+        *,
+        stream: StreamProtocol | None,
+        config: RunnableConfig,
+        store: BaseStore | None,
+        cache: BaseCache | None,
+        checkpointer: BaseCheckpointSaver | None,
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, BaseChannel | ManagedValueSpec],
+        trigger_to_nodes: Mapping[str, Sequence[str]],
+        durability: Durability,
+        manager: None | AsyncParentRunManager | ParentRunManager = None,
+        interrupt_after: All | Sequence[str] = EMPTY_SEQ,
+        interrupt_before: All | Sequence[str] = EMPTY_SEQ,
+        input_keys: str | Sequence[str] = EMPTY_SEQ,
+        output_keys: str | Sequence[str] = EMPTY_SEQ,
+        stream_keys: str | Sequence[str] = EMPTY_SEQ,
+        migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: CachePolicy | None = None,
+        has_graph_lifecycle_callbacks: bool = False,
+    ) -> None:
+        super().__init__(
+            input,
+            stream=stream,
+            config=config,
+            checkpointer=checkpointer,
+            cache=cache,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+            input_keys=input_keys,
+            output_keys=output_keys,
+            stream_keys=stream_keys,
+            interrupt_after=interrupt_after,
+            interrupt_before=interrupt_before,
+            manager=manager,
+            migrate_checkpoint=migrate_checkpoint,
+            trigger_to_nodes=trigger_to_nodes,
+            retry_policy=retry_policy,
+            cache_policy=cache_policy,
+            durability=durability,
+            has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks,
+        )
+        self.stack = ExitStack()
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.put_writes
+            self.checkpointer_put_writes_accepts_task_path = (
+                signature(checkpointer.put_writes).parameters.get("task_path")
+                is not None
+            )
+        else:
+            self.checkpointer_get_next_version = increment
+            self._checkpointer_put_after_previous = None  # type: ignore[assignment]
+            self.checkpointer_put_writes = None
+            self.checkpointer_put_writes_accepts_task_path = False
+
+    def _checkpointer_put_after_previous(
+        self,
+        prev: concurrent.futures.Future | None,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        if self._delta_write_futs:
+            futs, self._delta_write_futs = self._delta_write_futs, []
+            concurrent.futures.wait(futs)
+        try:
+            if prev is not None:
+                prev.result()
+        finally:
+            cast(BaseCheckpointSaver, self.checkpointer).put(
+                config, checkpoint, metadata, new_versions
+            )
+
+    def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        if self.cache is None:
+            return ()
+        matched: list[PregelExecutableTask] = []
+        if cached := {
+            (t.cache_key.ns, t.cache_key.key): t
+            for t in self.tasks.values()
+            if t.cache_key and not t.writes
+        }:
+            for key, values in self.cache.get(tuple(cached)).items():
+                task = cached[key]
+                task.writes.extend(values)
+                matched.append(task)
+        return matched
+
+    def accept_push(
+        self, task: PregelExecutableTask, write_idx: int, call: Call | None = None
+    ) -> PregelExecutableTask | None:
+        if pushed := super().accept_push(task, write_idx, call):
+            for task in self.match_cached_writes():
+                self.output_writes(task.id, task.writes, cached=True)
+        return pushed
+
+    def schedule_error_handler(
+        self, failed_task: PregelExecutableTask, error: BaseException
+    ) -> PregelExecutableTask | None:
+        handler_node = self.nodes[failed_task.name].error_handler_node
+        if not handler_node:
+            return None
+        # ensure error + ERROR_SOURCE_NODE writes are durable before handler runs
+        if self._error_handler_write_futs:
+            futs, self._error_handler_write_futs = self._error_handler_write_futs, []
+            concurrent.futures.wait(futs)
+        handler_task = prepare_node_error_handler_task(
+            failed_task,
+            handler_node_name=handler_node,
+            failed_error=error,
+            checkpoint=self.checkpoint,
+            pending_writes=self.checkpoint_pending_writes,
+            processes=self.nodes,
+            channels=self.channels,
+            managed=self.managed,
+            config=failed_task.config,
+            step=self.step,
+            stop=self.stop,
+            store=self.store,
+            checkpointer=self.checkpointer,
+            manager=self.manager,
+            retry_policy=self.retry_policy,
+            cache_policy=self.cache_policy,
+        )
+        if handler_task is None:
+            return None
+        self.tasks[handler_task.id] = handler_task
+        if not self.is_replaying:
+            self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
+        for task in self.match_cached_writes():
+            self.output_writes(task.id, task.writes, cached=True)
+        return handler_task
+
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        super().put_writes(task_id, writes)
+        if not writes or self.cache is None or not hasattr(self, "tasks"):
+            return
+        task = self.tasks.get(task_id)
+        if task is None or task.cache_key is None:
+            return
+        self.submit(
+            self.cache.set,
+            {
+                (task.cache_key.ns, task.cache_key.key): (
+                    task.writes,
+                    task.cache_key.ttl,
+                )
+            },
+        )
+
+    # context manager
+
+    def __enter__(self) -> Self:
+        self._graph_lifecycle_events = deque()
+        if not self.checkpointer:
+            saved = None
+        elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
+            # Explicit checkpoint_id requested — fetch that exact checkpoint.
+            # This covers both normal replay and subgraphs resolved via
+            # checkpoint_map during time-travel.
+            saved = self.checkpointer.get_tuple(self.checkpoint_config)
+        elif replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE):
+            # Subgraph replay: the parent is replaying and passed us a
+            # replay_state with its checkpoint_id. Look up our checkpoint
+            # from the parent's checkpoint_map instead of fetching latest.
+            saved = replay_state.get_checkpoint(
+                self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, ""),
+                self.checkpointer,
+                self.checkpoint_config,
+            )
+            # Clear RESUMING so _first re-applies input instead of resuming.
+            # This recreates ephemeral routing channels so nodes trigger
+            # naturally via version comparison.
+            self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
+        else:
+            # Normal case: fetch the most recent checkpoint for this
+            # graph/thread. Returns None on first invocation.
+            saved = self.checkpointer.get_tuple(self.checkpoint_config)
+
+        # Capture before the synthetic-empty fallback below overwrites `saved`.
+        # `_put_exit_delta_writes` uses this on first run (no persisted parent)
+        # to lazy-create a stub instead of anchoring delta writes on a parent.
+        self._has_persisted_parent = saved is not None
+        if saved is None:
+            saved = CheckpointTuple(
+                self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
+            )
+        elif self._migrate_checkpoint is not None:
+            self._migrate_checkpoint(saved.checkpoint)
+        self.checkpoint_config = {
+            **self.checkpoint_config,
+            **saved.config,
+            CONF: {
+                CONFIG_KEY_CHECKPOINT_NS: "",
+                **self.checkpoint_config.get(CONF, {}),
+                **saved.config.get(CONF, {}),
+            },
+        }
+        self._initial_checkpoint_config = self.checkpoint_config
+        self.prev_checkpoint_config = saved.parent_config
+        self.checkpoint_id_saved = saved.checkpoint["id"]
+        self.checkpoint = saved.checkpoint
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = (
+            [(str(tid), k, v) for tid, k, v in saved.pending_writes]
+            if saved.pending_writes is not None
+            else []
+        )
+        self._delta_write_futs = []
+        self._error_handler_write_futs = []
+        self._delta_channels_with_overwrite = set()
+        self._exit_delta_writes = (
+            [] if self.durability == "exit" and self.checkpointer is not None else None
+        )
+        self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
+        self.channels, self.managed = channels_from_checkpoint(
+            self.specs,
+            self.checkpoint,
+            saver=self.checkpointer,
+            config=self.checkpoint_config,
+        )
+        self.stack.push(self._suppress_interrupt)
+        self.status = "input"
+        self.step = self.checkpoint_metadata["step"] + 1
+        self.stop = self.step + self.config["recursion_limit"] + 1
+        self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
+        self.updated_channels = self._first(
+            input_keys=self.input_keys,
+            updated_channels=set(self.checkpoint.get("updated_channels"))  # type: ignore[arg-type]
+            if self.checkpoint.get("updated_channels")
+            else None,
+        )
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        # unwind stack
+        return self.stack.__exit__(exc_type, exc_value, traceback)
+
+
+class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
+    def __init__(
+        self,
+        input: Any | None,
+        *,
+        stream: StreamProtocol | None,
+        config: RunnableConfig,
+        store: BaseStore | None,
+        cache: BaseCache | None,
+        checkpointer: BaseCheckpointSaver | None,
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, BaseChannel | ManagedValueSpec],
+        trigger_to_nodes: Mapping[str, Sequence[str]],
+        durability: Durability,
+        interrupt_after: All | Sequence[str] = EMPTY_SEQ,
+        interrupt_before: All | Sequence[str] = EMPTY_SEQ,
+        manager: None | AsyncParentRunManager | ParentRunManager = None,
+        input_keys: str | Sequence[str] = EMPTY_SEQ,
+        output_keys: str | Sequence[str] = EMPTY_SEQ,
+        stream_keys: str | Sequence[str] = EMPTY_SEQ,
+        migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
+        retry_policy: Sequence[RetryPolicy] = (),
+        cache_policy: CachePolicy | None = None,
+        has_graph_lifecycle_callbacks: bool = False,
+    ) -> None:
+        super().__init__(
+            input,
+            stream=stream,
+            config=config,
+            checkpointer=checkpointer,
+            cache=cache,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+            input_keys=input_keys,
+            output_keys=output_keys,
+            stream_keys=stream_keys,
+            interrupt_after=interrupt_after,
+            interrupt_before=interrupt_before,
+            manager=manager,
+            migrate_checkpoint=migrate_checkpoint,
+            trigger_to_nodes=trigger_to_nodes,
+            retry_policy=retry_policy,
+            cache_policy=cache_policy,
+            durability=durability,
+            has_graph_lifecycle_callbacks=has_graph_lifecycle_callbacks,
+        )
+        self.stack = AsyncExitStack()
+        if checkpointer:
+            self.checkpointer_get_next_version = checkpointer.get_next_version
+            self.checkpointer_put_writes = checkpointer.aput_writes
+            self.checkpointer_put_writes_accepts_task_path = (
+                signature(checkpointer.aput_writes).parameters.get("task_path")
+                is not None
+            )
+        else:
+            self.checkpointer_get_next_version = increment
+            self._checkpointer_put_after_previous = None  # type: ignore[assignment]
+            self.checkpointer_put_writes = None
+            self.checkpointer_put_writes_accepts_task_path = False
+
+    async def _checkpointer_put_after_previous(
+        self,
+        prev: asyncio.Task | None,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        # Drain DeltaChannel write futures before committing the checkpoint so
+        # ancestor walks never see a checkpoint without its backing writes.
+        if self._delta_write_futs:
+            futs, self._delta_write_futs = self._delta_write_futs, []
+            await asyncio.gather(*futs)
+        try:
+            if prev is not None:
+                await prev
+        finally:
+            await cast(BaseCheckpointSaver, self.checkpointer).aput(
+                config, checkpoint, metadata, new_versions
+            )
+
+    async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
+        if self.cache is None:
+            return []
+        matched: list[PregelExecutableTask] = []
+        if cached := {
+            (t.cache_key.ns, t.cache_key.key): t
+            for t in self.tasks.values()
+            if t.cache_key and not t.writes
+        }:
+            for key, values in (await self.cache.aget(tuple(cached))).items():
+                task = cached[key]
+                task.writes.extend(values)
+                matched.append(task)
+        return matched
+
+    async def aaccept_push(
+        self, task: PregelExecutableTask, write_idx: int, call: Call | None = None
+    ) -> PregelExecutableTask | None:
+        if pushed := super().accept_push(task, write_idx, call):
+            for task in await self.amatch_cached_writes():
+                self.output_writes(task.id, task.writes, cached=True)
+        return pushed
+
+    async def aschedule_error_handler(
+        self, failed_task: PregelExecutableTask, error: BaseException
+    ) -> PregelExecutableTask | None:
+        handler_node = self.nodes[failed_task.name].error_handler_node
+        if not handler_node:
+            return None
+        # ensure error + ERROR_SOURCE_NODE writes are durable before handler runs
+        if self._error_handler_write_futs:
+            futs, self._error_handler_write_futs = self._error_handler_write_futs, []
+            await asyncio.gather(*futs)
+        handler_task = prepare_node_error_handler_task(
+            failed_task,
+            handler_node_name=handler_node,
+            failed_error=error,
+            checkpoint=self.checkpoint,
+            pending_writes=self.checkpoint_pending_writes,
+            processes=self.nodes,
+            channels=self.channels,
+            managed=self.managed,
+            config=failed_task.config,
+            step=self.step,
+            stop=self.stop,
+            store=self.store,
+            checkpointer=self.checkpointer,
+            manager=self.manager,
+            retry_policy=self.retry_policy,
+            cache_policy=self.cache_policy,
+        )
+        if handler_task is None:
+            return None
+        self.tasks[handler_task.id] = handler_task
+        if not self.is_replaying:
+            self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
+        for task in await self.amatch_cached_writes():
+            self.output_writes(task.id, task.writes, cached=True)
+        return handler_task
+
+    def put_writes(self, task_id: str, writes: WritesT) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        super().put_writes(task_id, writes)
+        if not writes or self.cache is None or not hasattr(self, "tasks"):
+            return
+        task = self.tasks.get(task_id)
+        if task is None or task.cache_key is None:
+            return
+        if writes[0][0] in (INTERRUPT, ERROR):
+            # only cache successful tasks
+            return
+        self.submit(
+            self.cache.aset,
+            {
+                (task.cache_key.ns, task.cache_key.key): (
+                    task.writes,
+                    task.cache_key.ttl,
+                )
+            },
+        )
+
+    # context manager
+
+    async def __aenter__(self) -> Self:
+        self._graph_lifecycle_events = deque()
+        if not self.checkpointer:
+            saved = None
+        elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
+            # Explicit checkpoint_id requested — fetch that exact checkpoint.
+            # This covers both normal replay and subgraphs resolved via
+            # checkpoint_map during time-travel.
+            saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
+        elif replay_state := self.config[CONF].get(CONFIG_KEY_REPLAY_STATE):
+            # Subgraph replay: the parent is replaying and passed us a
+            # replay_state with its checkpoint_id. Look up our checkpoint
+            # from the parent's checkpoint_map instead of fetching latest.
+            saved = await replay_state.aget_checkpoint(
+                self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, ""),
+                self.checkpointer,
+                self.checkpoint_config,
+            )
+            # Clear RESUMING so _first re-applies input instead of resuming.
+            # This recreates ephemeral routing channels so nodes trigger
+            # naturally via version comparison.
+            self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
+        else:
+            # Normal case: fetch the most recent checkpoint for this
+            # graph/thread. Returns None on first invocation.
+            saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
+
+        # Capture before the synthetic-empty fallback below overwrites `saved`.
+        # `_put_exit_delta_writes` uses this on first run (no persisted parent)
+        # to lazy-create a stub instead of anchoring delta writes on a parent.
+        self._has_persisted_parent = saved is not None
+        if saved is None:
+            saved = CheckpointTuple(
+                self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
+            )
+        elif self._migrate_checkpoint is not None:
+            self._migrate_checkpoint(saved.checkpoint)
+        self.checkpoint_config = {
+            **self.checkpoint_config,
+            **saved.config,
+            CONF: {
+                CONFIG_KEY_CHECKPOINT_NS: "",
+                **self.checkpoint_config.get(CONF, {}),
+                **saved.config.get(CONF, {}),
+            },
+        }
+        self._initial_checkpoint_config = self.checkpoint_config
+        self.prev_checkpoint_config = saved.parent_config
+        self.checkpoint_id_saved = saved.checkpoint["id"]
+        self.checkpoint = saved.checkpoint
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = (
+            [(str(tid), k, v) for tid, k, v in saved.pending_writes]
+            if saved.pending_writes is not None
+            else []
+        )
+        self._delta_write_futs = []
+        self._error_handler_write_futs = []
+        self._delta_channels_with_overwrite = set()
+        self._exit_delta_writes = (
+            [] if self.durability == "exit" and self.checkpointer is not None else None
+        )
+        self.submit = await self.stack.enter_async_context(
+            AsyncBackgroundExecutor(self.config)
+        )
+        self.channels, self.managed = await achannels_from_checkpoint(
+            self.specs,
+            self.checkpoint,
+            saver=self.checkpointer,
+            config=self.checkpoint_config,
+        )
+        self.stack.push(self._suppress_interrupt)
+        self.status = "input"
+        self.step = self.checkpoint_metadata["step"] + 1
+        self.stop = self.step + self.config["recursion_limit"] + 1
+        self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
+        self.updated_channels = self._first(
+            input_keys=self.input_keys,
+            updated_channels=set(self.checkpoint.get("updated_channels"))  # type: ignore[arg-type]
+            if self.checkpoint.get("updated_channels")
+            else None,
+        )
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        # unwind stack
+        exit_task = asyncio.create_task(
+            self.stack.__aexit__(exc_type, exc_value, traceback)
+        )
+        try:
+            return await exit_task
+        except asyncio.CancelledError as e:
+            # Bubble up the exit task upon cancellation to permit the API
+            # consumer to await it before e.g., reusing the DB connection.
+            e.args = (*e.args, exit_task)
+            raise