Merge pull request #1059 from langchain-ai/nc/18jul/loop

Make Pregel loop runnable step-by-step

libs/langgraph/langgraph/channels/manager.py

@@ -61,4 +61,6 @@ def create_checkpoint(
         channel_versions=checkpoint["channel_versions"],
         versions_seen=checkpoint["versions_seen"],
         pending_sends=checkpoint.get("pending_sends", []),
+        # checkpoints are saved only at the end of a step, ie. when current tasks should be cleared
+        current_tasks={},
     )

libs/langgraph/langgraph/checkpoint/base.py

@@ -1,5 +1,4 @@
 from abc import ABC
-from collections import defaultdict
 from datetime import datetime, timezone
 from typing import (
     Any,
@@ -25,6 +24,7 @@ from langgraph.serde.base import SerializerProtocol
 from langgraph.serde.jsonplus import JsonPlusSerializer
 
 V = TypeVar("V", int, float, str)
+PendingWrite = Tuple[str, str, Any]
 
 
 # Marked as total=False to allow for future expansion.
@@ -53,6 +53,10 @@ class CheckpointMetadata(TypedDict, total=False):
     """
 
 
+class TaskInfo(TypedDict):
+    status: Literal["scheduled", "success", "error"]
+
+
 class Checkpoint(TypedDict):
     """State snapshot at a given point in time."""
 
@@ -74,7 +78,7 @@ class Checkpoint(TypedDict):
     The keys are channel names and the values are the logical time step
     at which the channel was last updated.
     """
-    versions_seen: defaultdict[str, dict[str, Union[str, int, float]]]
+    versions_seen: dict[str, dict[str, Union[str, int, float]]]
     """Map from node ID to map from channel name to version seen.
     
     This keeps track of the versions of the channels that each node has seen.
@@ -84,6 +88,8 @@ class Checkpoint(TypedDict):
     pending_sends: List[Send]
     """List of packets sent to nodes but not yet processed.
     Cleared by the next checkpoint."""
+    current_tasks: Dict[str, TaskInfo]
+    """Map from task ID to task info."""
 
 
 def empty_checkpoint() -> Checkpoint:
@@ -93,8 +99,9 @@ def empty_checkpoint() -> Checkpoint:
         ts=datetime.now(timezone.utc).isoformat(),
         channel_values={},
         channel_versions={},
-        versions_seen=defaultdict(dict),
+        versions_seen={},
         pending_sends=[],
+        current_tasks={},
     )
 
 
@@ -105,11 +112,9 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
         id=checkpoint["id"],
         channel_values=checkpoint["channel_values"].copy(),
         channel_versions=checkpoint["channel_versions"].copy(),
-        versions_seen=defaultdict(
-            dict,
-            {k: v.copy() for k, v in checkpoint["versions_seen"].items()},
-        ),
+        versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
         pending_sends=checkpoint.get("pending_sends", []).copy(),
+        current_tasks=checkpoint.get("current_tasks", {}).copy(),
     )
 
 
@@ -118,7 +123,7 @@ class CheckpointTuple(NamedTuple):
     checkpoint: Checkpoint
     metadata: CheckpointMetadata
     parent_config: Optional[RunnableConfig] = None
-    pending_writes: Optional[List[Tuple[str, str, Any]]] = None
+    pending_writes: Optional[List[PendingWrite]] = None
 
 
 CheckpointThreadId = ConfigurableFieldSpec(

libs/langgraph/langgraph/constants.py

@@ -1,10 +1,11 @@
 from typing import Any
 
+INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 INTERRUPT = "__interrupt__"
 TASKS = "__pregel_tasks"
-RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ}
+RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT}
 TAG_HIDDEN = "langsmith:hidden"
 
 START = "__start__"

libs/langgraph/langgraph/pregel/algo.py

@@ -0,0 +1,427 @@
+import json
+from collections import defaultdict, deque
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    Literal,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Protocol,
+    Sequence,
+    Union,
+    overload,
+)
+from uuid import UUID, uuid5
+
+from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
+from langchain_core.runnables.config import (
+    RunnableConfig,
+    merge_configs,
+    patch_config,
+)
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.context import Context
+from langgraph.channels.manager import ChannelsManager, create_checkpoint
+from langgraph.checkpoint.base import Checkpoint, copy_checkpoint
+from langgraph.constants import (
+    CONFIG_KEY_READ,
+    CONFIG_KEY_SEND,
+    INTERRUPT,
+    RESERVED,
+    TAG_HIDDEN,
+    TASKS,
+    Send,
+)
+from langgraph.errors import EmptyChannelError, InvalidUpdateError
+from langgraph.managed.base import ManagedValueMapping, is_managed_value
+from langgraph.pregel.io import read_channel, read_channels
+from langgraph.pregel.log import logger
+from langgraph.pregel.read import PregelNode
+from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription
+
+
+class WritesProtocol(Protocol):
+    name: str
+    writes: Sequence[tuple[str, Any]]
+    triggers: Sequence[str]
+
+
+class PregelTaskWrites(NamedTuple):
+    name: str
+    writes: Sequence[tuple[str, Any]]
+    triggers: Sequence[str]
+
+
+def should_interrupt(
+    checkpoint: Checkpoint,
+    interrupt_nodes: Union[All, Sequence[str]],
+    tasks: list[PregelExecutableTask],
+) -> bool:
+    version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
+    null_version = version_type()
+    seen = checkpoint["versions_seen"].get(INTERRUPT, {})
+    return (
+        # interrupt if any channel has been updated since last interrupt
+        any(
+            version > seen.get(chan, null_version)
+            for chan, version in checkpoint["channel_versions"].items()
+        )
+        # and any triggered node is in interrupt_nodes list
+        and any(
+            task.name
+            for task in tasks
+            if (
+                (not task.config or TAG_HIDDEN not in task.config.get("tags"))
+                if interrupt_nodes == "*"
+                else task.name in interrupt_nodes
+            )
+        )
+    )
+
+
+def local_read(
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    task: WritesProtocol,
+    config: RunnableConfig,
+    select: Union[list[str], str],
+    fresh: bool = False,
+) -> Union[dict[str, Any], Any]:
+    if fresh:
+        new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
+        context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
+        with ChannelsManager(
+            {k: v for k, v in channels.items() if k not in context_channels},
+            new_checkpoint,
+            config,
+        ) as channels:
+            all_channels = {**channels, **context_channels}
+            apply_writes(new_checkpoint, all_channels, [task], None)
+            return read_channels(all_channels, select)
+    else:
+        return read_channels(channels, select)
+
+
+def local_write(
+    commit: Callable[[Sequence[tuple[str, Any]]], None],
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    writes: Sequence[tuple[str, Any]],
+) -> None:
+    for chan, value in writes:
+        if chan == TASKS:
+            if not isinstance(value, Send):
+                raise InvalidUpdateError(
+                    f"Invalid packet type, expected Packet, got {value}"
+                )
+            if value.node not in processes:
+                raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
+        elif chan not in channels:
+            logger.warning(f"Skipping write for channel '{chan}' which has no readers")
+    commit(writes)
+
+
+def increment(current: Optional[int], channel: BaseChannel) -> int:
+    return current + 1 if current is not None else 1
+
+
+def apply_writes(
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    tasks: Sequence[WritesProtocol],
+    get_next_version: Optional[Callable[[int, BaseChannel], int]],
+) -> None:
+    # update seen versions
+    for task in tasks:
+        checkpoint["versions_seen"].setdefault(task.name, {}).update(
+            {
+                chan: checkpoint["channel_versions"][chan]
+                for chan in task.triggers
+                if chan in checkpoint["channel_versions"]
+            }
+        )
+
+    # Find the highest version of all channels
+    if checkpoint["channel_versions"]:
+        max_version = max(checkpoint["channel_versions"].values())
+    else:
+        max_version = None
+    # Consume all channels that were read
+    for chan in {
+        chan for task in tasks for chan in task.triggers if chan not in RESERVED
+    }:
+        if channels[chan].consume():
+            if get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+
+    # clear pending sends
+    if checkpoint["pending_sends"]:
+        checkpoint["pending_sends"].clear()
+
+    # Group writes by channel
+    pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    for task in tasks:
+        for chan, val in task.writes:
+            if chan == TASKS:
+                checkpoint["pending_sends"].append(val)
+            else:
+                pending_writes_by_channel[chan].append(val)
+
+    # Find the highest version of all channels
+    if checkpoint["channel_versions"]:
+        max_version = max(checkpoint["channel_versions"].values())
+    else:
+        max_version = None
+
+    # Apply writes to channels
+    updated_channels: set[str] = set()
+    for chan, vals in pending_writes_by_channel.items():
+        if chan in channels:
+            try:
+                updated = channels[chan].update(vals)
+            except InvalidUpdateError as e:
+                raise InvalidUpdateError(
+                    f"Invalid update for channel {chan} with values {vals}"
+                ) from e
+            if updated and get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+            updated_channels.add(chan)
+
+    # Channels that weren't updated in this step are notified of a new step
+    for chan in channels:
+        if chan not in updated_channels:
+            if channels[chan].update([]) and get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version, channels[chan]
+                )
+
+
+@overload
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    for_execution: Literal[False],
+    manager: Literal[None] = None,
+) -> list[PregelTaskDescription]:
+    ...
+
+
+@overload
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    for_execution: Literal[True],
+    manager: Union[None, ParentRunManager, AsyncParentRunManager],
+) -> list[PregelExecutableTask]:
+    ...
+
+
+def prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, PregelNode],
+    channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
+    config: RunnableConfig,
+    step: int,
+    *,
+    for_execution: bool,
+    manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
+) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
+    tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
+    # Consume pending packets
+    for packet in checkpoint["pending_sends"]:
+        if not isinstance(packet, Send):
+            logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends")
+            continue
+        if for_execution:
+            proc = processes[packet.node]
+            if node := proc.get_node():
+                triggers = [TASKS]
+                metadata = {
+                    "langgraph_step": step,
+                    "langgraph_node": packet.node,
+                    "langgraph_triggers": triggers,
+                    "langgraph_task_idx": len(tasks),
+                }
+                task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
+                writes = deque()
+                tasks.append(
+                    PregelExecutableTask(
+                        packet.node,
+                        packet.arg,
+                        node,
+                        writes,
+                        patch_config(
+                            merge_configs(
+                                config,
+                                processes[packet.node].config,
+                                {"metadata": metadata},
+                            ),
+                            run_name=packet.node,
+                            callbacks=(
+                                manager.get_child(f"graph:step:{step}")
+                                if manager
+                                else None
+                            ),
+                            configurable={
+                                # deque.extend is thread-safe
+                                CONFIG_KEY_SEND: partial(
+                                    local_write, writes.extend, processes, channels
+                                ),
+                                CONFIG_KEY_READ: partial(
+                                    local_read,
+                                    checkpoint,
+                                    channels,
+                                    PregelTaskWrites(packet.node, writes, triggers),
+                                    config,
+                                ),
+                            },
+                        ),
+                        triggers,
+                        proc.retry_policy,
+                        task_id,
+                    )
+                )
+        else:
+            tasks.append(PregelTaskDescription(packet.node, packet.arg))
+    # Check if any processes should be run in next step
+    # If so, prepare the values to be passed to them
+    version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
+    null_version = version_type()
+    if null_version is None:
+        return tasks
+    for name, proc in processes.items():
+        seen = checkpoint["versions_seen"].get(name, {})
+        # If any of the channels read by this process were updated
+        if triggers := sorted(
+            chan
+            for chan in proc.triggers
+            if not isinstance(
+                read_channel(channels, chan, return_exception=True), EmptyChannelError
+            )
+            and checkpoint["channel_versions"].get(chan, null_version)
+            > seen.get(chan, null_version)
+        ):
+            try:
+                val = next(_proc_input(step, name, proc, managed, channels))
+            except StopIteration:
+                continue
+
+            if for_execution:
+                if node := proc.get_node():
+                    metadata = {
+                        "langgraph_step": step,
+                        "langgraph_node": name,
+                        "langgraph_triggers": triggers,
+                        "langgraph_task_idx": len(tasks),
+                    }
+                    task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
+                    writes = deque()
+                    tasks.append(
+                        PregelExecutableTask(
+                            name,
+                            val,
+                            node,
+                            writes,
+                            patch_config(
+                                merge_configs(
+                                    config,
+                                    proc.config,
+                                    {"metadata": metadata},
+                                ),
+                                run_name=name,
+                                callbacks=(
+                                    manager.get_child(f"graph:step:{step}")
+                                    if manager
+                                    else None
+                                ),
+                                configurable={
+                                    # deque.extend is thread-safe
+                                    CONFIG_KEY_SEND: partial(
+                                        local_write, writes.extend, processes, channels
+                                    ),
+                                    CONFIG_KEY_READ: partial(
+                                        local_read,
+                                        checkpoint,
+                                        channels,
+                                        PregelTaskWrites(name, writes, triggers),
+                                        config,
+                                    ),
+                                },
+                            ),
+                            triggers,
+                            proc.retry_policy,
+                            task_id,
+                        )
+                    )
+            else:
+                tasks.append(PregelTaskDescription(name, val))
+    return tasks
+
+
+def _proc_input(
+    step: int,
+    name: str,
+    proc: PregelNode,
+    managed: ManagedValueMapping,
+    channels: Mapping[str, BaseChannel],
+) -> Iterator[Any]:
+    # If all trigger channels subscribed by this process are not empty
+    # then invoke the process with the values of all non-empty channels
+    if isinstance(proc.channels, dict):
+        try:
+            val: dict = {
+                k: read_channel(
+                    channels,
+                    chan,
+                    catch=chan not in proc.triggers,
+                )
+                for k, chan in proc.channels.items()
+                if isinstance(chan, str)
+            }
+
+            managed_values = {}
+            for key, chan in proc.channels.items():
+                if is_managed_value(chan):
+                    managed_values[key] = managed[key](
+                        step, PregelTaskDescription(name, val)
+                    )
+
+            val.update(managed_values)
+        except EmptyChannelError:
+            return
+    elif isinstance(proc.channels, list):
+        for chan in proc.channels:
+            try:
+                val = read_channel(channels, chan, catch=False)
+                break
+            except EmptyChannelError:
+                pass
+        else:
+            return
+    else:
+        raise RuntimeError(
+            "Invalid channels type, expected list or dict, got {proc.channels}"
+        )
+
+    # If the process has a mapper, apply it to the value
+    if proc.mapper is not None:
+        val = proc.mapper(val)
+
+    yield val

libs/langgraph/langgraph/pregel/executor.py

@@ -6,6 +6,7 @@ from contextvars import copy_context
 from types import TracebackType
 from typing import (
     AsyncContextManager,
+    Awaitable,
     Callable,
     Iterator,
     Optional,
@@ -78,7 +79,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
 
     def submit(
         self,
-        fn: Callable[P, T],
+        fn: Callable[P, Awaitable[T]],
         *args: P.args,
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
@@ -101,7 +102,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         else:
             self.tasks.pop(task)
 
-    async def __aenter__(self) -> Submit:
+    async def __aenter__(self) -> "submit":
         return self.submit
 
     async def exit(self) -> None:

libs/langgraph/langgraph/pregel/loop.py

@@ -0,0 +1,423 @@
+import asyncio
+from collections import deque
+from contextlib import AsyncExitStack, ExitStack
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    Callable,
+    ContextManager,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+
+from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import Self
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.manager import (
+    AsyncChannelsManager,
+    ChannelsManager,
+    create_checkpoint,
+)
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    PendingWrite,
+    copy_checkpoint,
+    empty_checkpoint,
+)
+from langgraph.constants import INPUT, INTERRUPT
+from langgraph.managed.base import (
+    AsyncManagedValuesManager,
+    ManagedValueMapping,
+    ManagedValuesManager,
+)
+from langgraph.pregel.algo import (
+    PregelTaskWrites,
+    apply_writes,
+    increment,
+    prepare_next_tasks,
+    should_interrupt,
+)
+from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
+from langgraph.pregel.executor import (
+    AsyncBackgroundExecutor,
+    BackgroundExecutor,
+    Submit,
+)
+from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
+from langgraph.pregel.types import PregelExecutableTask
+
+if TYPE_CHECKING:
+    from langgraph.pregel import Pregel
+
+
+V = TypeVar("V")
+INPUT_DONE = object()
+
+
+class PregelLoop:
+    input: Optional[Any]
+    config: RunnableConfig
+    checkpointer: Optional[BaseCheckpointSaver]
+    checkpointer_get_next_version: Callable[[Optional[V]], V]
+    checkpointer_put_writes: Optional[
+        Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
+    ]
+    checkpointer_put: Optional[
+        Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any]
+    ]
+    graph: "Pregel"
+
+    submit: Submit
+    channels: Mapping[str, BaseChannel]
+    managed: ManagedValueMapping
+    checkpoint: Checkpoint
+    checkpoint_config: RunnableConfig
+    checkpoint_metadata: CheckpointMetadata
+    checkpoint_pending_writes: Optional[List[PendingWrite]]
+
+    step: int
+    status: Literal[
+        "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
+    ]
+    tasks: Sequence[PregelExecutableTask]
+    stream: deque[Tuple[str, Any]]
+
+    # public
+
+    def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
+        """Mark tasks as scheduled, to be used by queue-based executors."""
+        raise NotImplementedError
+
+    def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
+        """Put writes for a task, to be read by the next tick."""
+        self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
+        if self.checkpointer_put_writes is not None:
+            self.submit(
+                self.checkpointer_put_writes,
+                {
+                    **self.checkpoint_config,
+                    "configurable": {
+                        **self.checkpoint_config["configurable"],
+                        "thread_ts": self.checkpoint["id"],
+                    },
+                },
+                writes,
+                task_id,
+            )
+
+    def tick(
+        self,
+        *,
+        output_keys: Union[str, Sequence[str]] = None,
+        interrupt_after: Optional[Sequence[str]] = None,
+        interrupt_before: Optional[Sequence[str]] = None,
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
+    ) -> bool:
+        """Execute a single iteration of the Pregel loop.
+        Returns True if more iterations are needed."""
+
+        if self.status != "pending":
+            raise RuntimeError("Cannot tick when status is no longer 'pending'")
+
+        if self.input is not INPUT_DONE:
+            self._first()
+        elif all(task.writes for task in self.tasks):
+            writes = [w for t in self.tasks for w in t.writes]
+            # all tasks have finished
+            apply_writes(
+                self.checkpoint,
+                self.channels,
+                self.tasks,
+                self.checkpointer_get_next_version,
+            )
+            # produce values output
+            self.stream.extend(
+                ("values", v)
+                for v in map_output_values(output_keys, writes, self.channels)
+            )
+            # clear pending writes
+            self.checkpoint_pending_writes.clear()
+            # save checkpoint
+            self._put_checkpoint(
+                {
+                    "source": "loop",
+                    "writes": single(
+                        map_output_updates(output_keys, self.tasks)
+                        if self.graph.stream_mode == "updates"
+                        else map_output_values(output_keys, writes, self.channels)
+                    ),
+                }
+            )
+            # after execution, check if we should interrupt
+            if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
+                self.status = "interrupt_after"
+                return False
+        else:
+            return False
+
+        # check if iteration limit is reached
+        if self.step > self.config["recursion_limit"]:
+            self.status = "out_of_steps"
+            return False
+
+        # prepare next tasks
+        self.tasks = prepare_next_tasks(
+            self.checkpoint,
+            self.graph.nodes,
+            self.channels,
+            self.managed,
+            self.config,
+            self.step,
+            for_execution=True,
+            manager=manager,
+        )
+
+        # if no more tasks, we're done
+        if not self.tasks:
+            self.status = "done"
+            return False
+
+        # if there are pending writes from a previous loop, apply them
+        if self.checkpoint_pending_writes:
+            for tid, k, v in self.checkpoint_pending_writes:
+                if task := next((t for t in self.tasks if t.id == tid), None):
+                    task.writes.append((k, v))
+
+        # if all tasks have finished, re-tick
+        if all(task.writes for task in self.tasks):
+            return self.tick()
+
+        # before execution, check if we should interrupt
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+            self.status = "interrupt_before"
+            return False
+
+        # produce debug output
+        self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks))
+
+        return True
+
+    # private
+
+    def _first(self) -> None:
+        # map inputs to channel updates
+        if input_writes := deque(map_input(self.graph.input_channels, self.input)):
+            # discard any unfinished tasks from previous checkpoint
+            discard_tasks = prepare_next_tasks(
+                self.checkpoint,
+                self.graph.nodes,
+                self.channels,
+                self.managed,
+                self.config,
+                self.step,
+                for_execution=True,
+            )
+            # apply input writes
+            apply_writes(
+                self.checkpoint,
+                self.channels,
+                discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
+                self.checkpointer_get_next_version,
+            )
+            # save input checkpoint
+            self._put_checkpoint({"source": "input", "writes": self.input})
+        else:
+            # no input is taken as signal to proceed past previous interrupt
+            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
+        # done with input
+        self.input = INPUT_DONE
+
+    def _put_checkpoint(
+        self,
+        metadata: CheckpointMetadata,
+    ) -> None:
+        # assign step
+        metadata["step"] = self.step
+        # bail if no checkpointer
+        if self.checkpointer_put is not None:
+            # create new checkpoint
+            self.checkpoint_metadata = metadata
+            self.checkpoint = create_checkpoint(
+                self.checkpoint, self.channels, self.step
+            )
+            # save it, without blocking
+            self.submit(
+                self.checkpointer_put,
+                self.checkpoint_config,
+                copy_checkpoint(self.checkpoint),
+                self.checkpoint_metadata,
+            )
+            self.checkpoint_config = {
+                **self.checkpoint_config,
+                "configurable": {
+                    **self.checkpoint_config["configurable"],
+                    "thread_ts": self.checkpoint["id"],
+                },
+            }
+            # produce debug output
+            self.stream.extend(
+                ("debug", v)
+                for v in map_debug_checkpoint(
+                    self.step,
+                    self.checkpoint_config,
+                    self.channels,
+                    self.graph.stream_channels_asis,
+                    self.checkpoint_metadata,
+                )
+            )
+        # increment step
+        self.step += 1
+
+
+class SyncPregelLoop(PregelLoop, ContextManager):
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.stack = ExitStack()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.checkpointer_get_next_version = (
+            checkpointer.get_next_version if checkpointer else increment
+        )
+        self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
+        self.checkpointer_put = checkpointer.put if checkpointer else None
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+
+    # context manager
+
+    def __enter__(self) -> Self:
+        saved = (
+            self.checkpointer.get_tuple(self.config) if self.checkpointer else None
+        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        self.checkpoint_config = {
+            **self.config,
+            **saved.config,
+            "configurable": {
+                **self.config.get("configurable", {}),
+                **saved.config.get("configurable", {}),
+            },
+        }
+        self.checkpoint = copy_checkpoint(saved.checkpoint)
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = saved.pending_writes
+
+        self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
+        self.channels = self.stack.enter_context(
+            ChannelsManager(self.graph.channels, self.checkpoint, self.config)
+        )
+        self.managed = self.stack.enter_context(
+            ManagedValuesManager(
+                self.graph.managed_values_dict, self.config, self.graph
+            )
+        )
+        self.status = "pending"
+        self.step = self.checkpoint_metadata["step"] + 1
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        del self.graph
+        return self.stack.__exit__(exc_type, exc_value, traceback)
+
+
+class AsyncPregelLoop(PregelLoop, AsyncContextManager):
+    def __init__(
+        self,
+        input: Optional[Any],
+        *,
+        config: RunnableConfig,
+        checkpointer: Optional[BaseCheckpointSaver],
+        graph: "Pregel",
+    ) -> None:
+        self.stream = deque()
+        self.stack = AsyncExitStack()
+        self.input = input
+        self.config = config
+        self.checkpointer = checkpointer
+        self.checkpointer_get_next_version = (
+            checkpointer.get_next_version if checkpointer else increment
+        )
+        self.checkpointer_put_writes = (
+            checkpointer.aput_writes if checkpointer else None
+        )
+        self.checkpointer_put = checkpointer.aput if checkpointer else None
+        self.graph = graph
+        # TODO if managed values no longer needs graph we can replace with
+        # managed_specs, channel_specs
+
+    # context manager
+
+    async def __aenter__(self) -> Self:
+        saved = (
+            await self.checkpointer.aget_tuple(self.config)
+            if self.checkpointer
+            else None
+        ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
+        self.checkpoint_config = {
+            **self.config,
+            **saved.config,
+            "configurable": {
+                **self.config.get("configurable", {}),
+                **saved.config.get("configurable", {}),
+            },
+        }
+        self.checkpoint = copy_checkpoint(saved.checkpoint)
+        self.checkpoint_metadata = saved.metadata
+        self.checkpoint_pending_writes = saved.pending_writes
+
+        self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
+        self.channels = await self.stack.enter_async_context(
+            AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
+        )
+        self.managed = await self.stack.enter_async_context(
+            AsyncManagedValuesManager(
+                self.graph.managed_values_dict, self.config, self.graph
+            )
+        )
+        self.status = "pending"
+        self.step = self.checkpoint_metadata["step"] + 1
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        del self.graph
+        return await asyncio.shield(
+            self.stack.__aexit__(exc_type, exc_value, traceback)
+        )

libs/langgraph/langgraph/pregel/types.py

@@ -88,3 +88,12 @@ class StateSnapshot(NamedTuple):
 
 
 All = Literal["*"]
+
+StreamMode = Literal["values", "updates", "debug"]
+"""How the stream method should emit outputs.
+
+- 'values': Emit all values of the state for each step.
+- 'updates': Emit only the node name(s) and updates
+    that were returned by the node(s) **after** each step.
+- 'debug': Emit debug events for each step.
+"""

libs/langgraph/poetry.lock

@@ -747,6 +747,20 @@ files = [
 [package.extras]
 test = ["pytest (>=6)"]
 
+[[package]]
+name = "execnet"
+version = "2.1.1"
+description = "execnet: rapid multi-Python deployment"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "execnet-2.1.1-py3-none-any.whl", hash = "sha256:26dee51f1b80cebd6d0ca8e74dd8745419761d3bef34163928cbebbdc4749fdc"},
+    {file = "execnet-2.1.1.tar.gz", hash = "sha256:5189b52c6121c24feae288166ab41b32549c7e2348652736540b9e6e7d4e72e3"},
+]
+
+[package.extras]
+testing = ["hatch", "pre-commit", "pytest", "tox"]
+
 [[package]]
 name = "executing"
 version = "2.0.1"
@@ -1746,13 +1760,13 @@ langchain-core = ">=0.2.2rc1,<0.3"
 
 [[package]]
 name = "langchain-core"
-version = "0.2.19"
+version = "0.2.22"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.8.1"
 files = [
-    {file = "langchain_core-0.2.19-py3-none-any.whl", hash = "sha256:5b3cd34395be274c89e822c84f0e03c4da14168c177a83921c5b9414ac7a0651"},
-    {file = "langchain_core-0.2.19.tar.gz", hash = "sha256:13043a83e5c9ab58b9f5ce2a56896e7e88b752e8891b2958960a98e71801471e"},
+    {file = "langchain_core-0.2.22-py3-none-any.whl", hash = "sha256:7731a86440c0958b3186c003fb9b26b2d5a682a6344bda7bfb9174e2898f8b43"},
+    {file = "langchain_core-0.2.22.tar.gz", hash = "sha256:582d6f929a43b830139444e4124123cd415331ad62f25757b1406252958cdcac"},
 ]
 
 [package.dependencies]
@@ -2784,6 +2798,27 @@ files = [
 tomli = {version = ">=2.0.1,<3.0.0", markers = "python_version < \"3.11\""}
 watchdog = ">=2.0.0"
 
+[[package]]
+name = "pytest-xdist"
+version = "3.6.1"
+description = "pytest xdist plugin for distributed testing, most importantly across multiple CPUs"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pytest_xdist-3.6.1-py3-none-any.whl", hash = "sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7"},
+    {file = "pytest_xdist-3.6.1.tar.gz", hash = "sha256:ead156a4db231eec769737f57668ef58a2084a34b2e55c4a8fa20d861107300d"},
+]
+
+[package.dependencies]
+execnet = ">=2.1"
+psutil = {version = ">=3.0", optional = true, markers = "extra == \"psutil\""}
+pytest = ">=7.0.0"
+
+[package.extras]
+psutil = ["psutil (>=3.0)"]
+setproctitle = ["setproctitle"]
+testing = ["filelock"]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -4130,4 +4165,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "170eaa0e542a02d5f2fb0d42d1f04c5d010bd3735a44b927b28e6b742c689eb0"
+content-hash = "5fb6190a1b01d0cd351ea9a0023c8c8d6acf4fe831101ab87f9fc41308f74b74"

libs/langgraph/pyproject.toml

@@ -31,6 +31,7 @@ langchainhub = "^0.1.14"
 langchain-openai = ">=0.1.2"
 langchain-anthropic = ">=0.1.8"
 dataclasses-json = "^0.6.7"
+pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
 
 [tool.poetry.group.dev]
 optional = true
@@ -61,7 +62,7 @@ omit = ["tests/*"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["-x", "--ff", "-vv", "--snapshot-update"]
+runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
 patterns = ["*.py"]
 
 [build-system]

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -509,10 +509,10 @@
   '''
 # ---
 # name: test_conditional_state_graph
-  '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
+  '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
 # ---
 # name: test_conditional_state_graph.1
-  '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
+  '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
 # ---
 # name: test_conditional_state_graph.2
   '''

libs/langgraph/tests/test_prebuilt.py

@@ -22,11 +22,14 @@ from langchain_core.tools import BaseTool
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel as BaseModelV2
 
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.prebuilt import (
     ToolNode,
     ValidationNode,
     create_react_agent,
 )
+from tests.any_str import AnyStr
+from tests.memory_assert import MemorySaverAssertImmutable
 
 
 class FakeToolCallingModel(BaseChatModel):
@@ -56,14 +59,117 @@ class FakeToolCallingModel(BaseChatModel):
         return self
 
 
-def test_no_modifier():
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        None,
+    ],
+    ids=[
+        "memory",
+        "none",
+    ],
+)
+def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
     model = FakeToolCallingModel()
-    agent = create_react_agent(model, [])
+    agent = create_react_agent(model, [], checkpointer=checkpointer)
     inputs = [HumanMessage("hi?")]
-    response = agent.invoke({"messages": inputs})
+    thread = {"configurable": {"thread_id": "123"}}
+    response = agent.invoke({"messages": inputs}, thread, debug=True)
     expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
     assert response == expected_response
 
+    if checkpointer:
+        saved = checkpointer.get_tuple(thread)
+        assert saved is not None
+        assert saved.checkpoint == {
+            "v": 1,
+            "ts": AnyStr(),
+            "id": AnyStr(),
+            "channel_values": {
+                "messages": [
+                    HumanMessage(content="hi?", id=AnyStr()),
+                    AIMessage(content="hi?", id="0"),
+                ],
+                "agent": "agent",
+            },
+            "channel_versions": {
+                "__start__": 2,
+                "messages": 3,
+                "start:agent": 3,
+                "agent": 3,
+            },
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {"__start__": 1},
+                "agent": {"start:agent": 2},
+            },
+            "pending_sends": [],
+            "current_tasks": {},
+        }
+        assert saved.metadata == {
+            "source": "loop",
+            "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
+            "step": 1,
+        }
+        assert saved.pending_writes == []
+
+
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        None,
+    ],
+    ids=[
+        "memory",
+        "none",
+    ],
+)
+async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
+    model = FakeToolCallingModel()
+    agent = create_react_agent(model, [], checkpointer=checkpointer)
+    inputs = [HumanMessage("hi?")]
+    thread = {"configurable": {"thread_id": "123"}}
+    response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
+    expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
+    assert response == expected_response
+
+    if checkpointer:
+        saved = await checkpointer.aget_tuple(thread)
+        assert saved is not None
+        assert saved.checkpoint == {
+            "v": 1,
+            "ts": AnyStr(),
+            "id": AnyStr(),
+            "channel_values": {
+                "messages": [
+                    HumanMessage(content="hi?", id=AnyStr()),
+                    AIMessage(content="hi?", id="0"),
+                ],
+                "agent": "agent",
+            },
+            "channel_versions": {
+                "__start__": 2,
+                "messages": 3,
+                "start:agent": 3,
+                "agent": 3,
+            },
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {"__start__": 1},
+                "agent": {"start:agent": 2},
+            },
+            "pending_sends": [],
+            "current_tasks": {},
+        }
+        assert saved.metadata == {
+            "source": "loop",
+            "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
+            "step": 1,
+        }
+        assert saved.pending_writes == []
+
 
 def test_passing_two_modifiers():
     model = FakeToolCallingModel()

libs/langgraph/tests/test_pregel.py

@@ -520,10 +520,8 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     assert app.invoke(2) == 4
 
-    assert app.invoke(2, input_keys="inbox") == 3
-
     with pytest.raises(GraphRecursionError):
-        app.invoke(2, {"recursion_limit": 1})
+        app.invoke(2, {"recursion_limit": 1}, debug=1)
 
     graph = Graph()
     graph.add_node("add_one", add_one)
@@ -535,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     assert gapp.invoke(2) == 4
 
-    for step, values in enumerate(gapp.stream(2), start=1):
+    for step, values in enumerate(gapp.stream(2, debug=1), start=1):
         if step == 1:
             assert values == {
                 "add_one": 3,
@@ -6168,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             "my_key": "value ⛰️",
             "market": "DE",
         }
+        assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
+            {
+                "source": "loop",
+                "step": 0,
+                "writes": None,
+            },
+            {
+                "source": "input",
+                "step": -1,
+                "writes": {"my_key": "value ⛰️", "market": "DE"},
+            },
+        ]
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value ⛰️", "market": "DE"},
             next=("tool_two_slow",),
@@ -6877,7 +6887,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
         },
     )
 
-    assert [c for c in app_w_interrupt.stream(None, config)] == [
+    assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
         {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
     ]
 

libs/langgraph/tests/test_pregel_async.py

@@ -623,8 +623,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
 
     assert await app.ainvoke(2) == 4
 
-    assert await app.ainvoke(2, input_keys="inbox") == 3
-
     with pytest.raises(GraphRecursionError):
         await app.ainvoke(2, {"recursion_limit": 1})
 
@@ -4737,6 +4735,18 @@ async def test_start_branch_then() -> None:
             "my_key": "value",
             "market": "DE",
         }
+        assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
+            {
+                "source": "loop",
+                "step": 0,
+                "writes": None,
+            },
+            {
+                "source": "input",
+                "step": -1,
+                "writes": {"my_key": "value", "market": "DE"},
+            },
+        ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value", "market": "DE"},
             next=("tool_two_slow",),