Merge pull request #2346 from langchain-ai/nc/4nov/send-eager

lib: Execute Sends in the superstep that originated them (feature-flagged)

.github/workflows/_test_langgraph.yml

@@ -19,14 +19,19 @@ jobs:
           - "3.13"
         core-version:
           - "latest"
+        ff-send-v2:
+          - "false"
         include:
           - python-version: "3.11"
             core-version: ">=0.2.42,<0.3.0"
+          - python-version: "3.11"
+            core-version: "latest"
+            ff-send-v2: "true"
 
     defaults:
       run:
         working-directory: libs/langgraph
-    name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }})"
+    name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }}, ff-send-v2: ${{ matrix.ff-send-v2 }})"
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
@@ -52,6 +57,8 @@ jobs:
 
       - name: Run tests
         shell: bash
+        env:
+          LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }}
         run: |
           make test
 

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -24,6 +24,8 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.checkpoint.serde.types import (
     ERROR,
+    INTERRUPT,
+    RESUME,
     SCHEDULED,
     ChannelProtocol,
     SendProtocol,
@@ -449,4 +451,4 @@ Special writes (e.g. errors) map to negative indices, to avoid those writes from
 conflicting with regular writes.
 Each Checkpointer implementation should use this mapping in put_writes.
 """
-WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2}
+WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3, RESUME: -4}

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -25,7 +25,7 @@ from langchain_core.load.serializable import Serializable
 from zoneinfo import ZoneInfo
 
 from langgraph.checkpoint.serde.base import SerializerProtocol
-from langgraph.checkpoint.serde.types import CommandProtocol, SendProtocol
+from langgraph.checkpoint.serde.types import SendProtocol
 from langgraph.store.base import Item
 
 LC_REVIVER = Reviver()
@@ -122,11 +122,6 @@ class JsonPlusSerializer(SerializerProtocol):
             return self._encode_constructor_args(
                 obj.__class__, kwargs={"node": obj.node, "arg": obj.arg}
             )
-        elif isinstance(obj, CommandProtocol):
-            return self._encode_constructor_args(
-                obj.__class__,
-                kwargs={k: getattr(obj, k) for k in obj.__all_slots__},
-            )
         elif isinstance(obj, (bytes, bytearray)):
             return self._encode_constructor_args(
                 obj.__class__, method="fromhex", args=(obj.hex(),)
@@ -407,17 +402,6 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]:
                 (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)),
             ),
         )
-    elif isinstance(obj, CommandProtocol):
-        return msgpack.ExtType(
-            EXT_CONSTRUCTOR_KW_ARGS,
-            _msgpack_enc(
-                (
-                    obj.__class__.__module__,
-                    obj.__class__.__name__,
-                    {k: getattr(obj, k) for k in obj.__all_slots__},
-                ),
-            ),
-        )
     elif dataclasses.is_dataclass(obj):
         # doesn't use dataclasses.asdict to avoid deepcopy and recursion
         return msgpack.ExtType(

libs/checkpoint/langgraph/checkpoint/serde/types.py

@@ -4,7 +4,6 @@ from typing import (
     Protocol,
     Sequence,
     TypeVar,
-    Union,
     runtime_checkable,
 )
 
@@ -12,6 +11,8 @@ from typing_extensions import Self
 
 ERROR = "__error__"
 SCHEDULED = "__scheduled__"
+INTERRUPT = "__interrupt__"
+RESUME = "__resume__"
 TASKS = "__pregel_tasks"
 
 Value = TypeVar("Value", covariant=True)
@@ -49,11 +50,3 @@ class SendProtocol(Protocol):
     def __repr__(self) -> str: ...
 
     def __eq__(self, value: object) -> bool: ...
-
-
-@runtime_checkable
-class CommandProtocol(Protocol):
-    # Mirrors langgraph.types.Command
-    update: Optional[dict[str, Any]]
-    send: Union[Any, Sequence[Any]]
-    __all_slots__: set[str]

libs/langgraph/Makefile

@@ -49,7 +49,7 @@ test:
 	exit $$EXIT_CODE
 
 test_watch:
-	make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
+	make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
 	EXIT_CODE=$$?; \
 	make stop-postgres; \
 	exit $$EXIT_CODE

libs/langgraph/langgraph/constants.py

@@ -1,4 +1,5 @@
 import sys
+from os import getenv
 from types import MappingProxyType
 from typing import Any, Literal, Mapping, cast
 
@@ -10,6 +11,7 @@ from langgraph.types import Interrupt, Send  # noqa: F401
 # --- Empty read-only containers ---
 EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
 EMPTY_SEQ: tuple[str, ...] = tuple()
+MISSING = object()
 
 # --- Public constants ---
 TAG_NOSTREAM = sys.intern("langsmith:nostream")
@@ -28,6 +30,8 @@ INPUT = sys.intern("__input__")
 # for values passed as input to the graph
 INTERRUPT = sys.intern("__interrupt__")
 # for dynamic interrupts raised by nodes
+RESUME = sys.intern("__resume__")
+# for values passed to resume a node after an interrupt
 ERROR = sys.intern("__error__")
 # for errors raised by nodes
 NO_WRITES = sys.intern("__no_writes__")
@@ -69,6 +73,8 @@ CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
 # holds the current checkpoint_ns, "" for root graph
 CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
 # callback to be called when a node is finished
+CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
+# holds the value that "answers" an interrupt() call
 
 # --- Other constants ---
 PUSH = sys.intern("__pregel_push")
@@ -81,12 +87,17 @@ NS_END = sys.intern(":")
 # for checkpoint_ns, for each level, separates the namespace from the task_id
 CONF = cast(Literal["configurable"], sys.intern("configurable"))
 # key for the configurable dict in RunnableConfig
+FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true"
+# temporary flag to enable new Send semantics
+NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
+# the task_id to use for writes that are not associated with a task
 
 RESERVED = {
     TAG_HIDDEN,
     # reserved write keys
     INPUT,
     INTERRUPT,
+    RESUME,
     ERROR,
     NO_WRITES,
     SCHEDULED,

libs/langgraph/langgraph/errors.py

@@ -70,7 +70,7 @@ class NodeInterrupt(GraphInterrupt):
     """Raised by a node to interrupt execution."""
 
     def __init__(self, value: Any) -> None:
-        super().__init__([Interrupt(value)])
+        super().__init__([Interrupt(value=value)])
 
 
 class GraphDelegate(Exception):

libs/langgraph/langgraph/graph/state.py

@@ -1,3 +1,4 @@
+import dataclasses
 import inspect
 import logging
 import typing
@@ -14,7 +15,6 @@ from typing import (
     Optional,
     Sequence,
     Type,
-    TypeVar,
     Union,
     cast,
     get_args,
@@ -50,15 +50,13 @@ from langgraph.managed.base import (
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.types import All, Checkpointer, Command, RetryPolicy
+from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy
 from langgraph.utils.fields import get_field_default
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
 
-N = TypeVar("N")
-
 
 def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None:
     if isinstance(schema, type):
@@ -81,20 +79,20 @@ def _get_node_name(node: RunnableLike) -> str:
         raise TypeError(f"Unsupported node type: {type(node)}")
 
 
-class GraphCommand(Command, Generic[N]):
+@dataclasses.dataclass(**_DC_KWARGS)
+class GraphCommand(Generic[N], Command[N]):
     """One or more commands to update a StateGraph's state and go to, or send messages to nodes."""
 
-    __slots__ = ("goto",)
+    goto: Union[str, Sequence[str]] = ()
 
-    def __init__(
-        self,
-        *,
-        update: Optional[dict[str, Any]] = None,
-        goto: Union[str, Sequence[str]] = (),
-        send: Union[Send, Sequence[Send]] = (),
-    ) -> None:
-        super().__init__(update=update, send=send)
-        self.goto = goto
+    def __repr__(self) -> str:
+        # get all non-None values
+        contents = ", ".join(
+            f"{key}={value!r}"
+            for key, value in dataclasses.asdict(self).items()
+            if value
+        )
+        return f"Command({contents})"
 
 
 class StateNodeSpec(NamedTuple):
@@ -389,7 +387,7 @@ class StateGraph(Graph):
                             input = input_hint
                 if (
                     (rtn := hints.get("return"))
-                    and get_origin(rtn) is GraphCommand
+                    and get_origin(rtn) in (Command, GraphCommand)
                     and (rargs := get_args(rtn))
                     and get_origin(rargs[0]) is Literal
                     and (vals := get_args(rargs[0]))

libs/langgraph/langgraph/pregel/__init__.py

@@ -66,9 +66,12 @@ from langgraph.constants import (
     CONFIG_KEY_STREAM_WRITER,
     CONFIG_KEY_TASK_ID,
     ERROR,
+    INPUT,
     INTERRUPT,
     NS_END,
     NS_SEP,
+    NULL_TASK_ID,
+    PUSH,
     SCHEDULED,
 )
 from langgraph.errors import (
@@ -98,7 +101,13 @@ from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_version
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.types import All, Checkpointer, LoopProtocol, StateSnapshot, StreamMode
+from langgraph.types import (
+    All,
+    Checkpointer,
+    LoopProtocol,
+    StateSnapshot,
+    StreamMode,
+)
 from langgraph.utils.config import (
     ensure_config,
     merge_configs,
@@ -468,6 +477,7 @@ class Pregel(PregelProtocol):
             # tasks for this checkpoint
             next_tasks = prepare_next_tasks(
                 saved.checkpoint,
+                saved.pending_writes or [],
                 self.nodes,
                 channels,
                 managed,
@@ -511,6 +521,15 @@ class Pregel(PregelProtocol):
                         config, subgraphs=True
                     )
             # apply pending writes
+            if null_writes := [
+                w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
+            ]:
+                apply_writes(
+                    saved.checkpoint,
+                    channels,
+                    [PregelTaskWrites((), INPUT, null_writes, [])],
+                    None,
+                )
             if apply_pending_writes and saved.pending_writes:
                 for tid, k, v in saved.pending_writes:
                     if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -570,6 +589,7 @@ class Pregel(PregelProtocol):
             # tasks for this checkpoint
             next_tasks = prepare_next_tasks(
                 saved.checkpoint,
+                saved.pending_writes or [],
                 self.nodes,
                 channels,
                 managed,
@@ -613,6 +633,15 @@ class Pregel(PregelProtocol):
                         config, subgraphs=True
                     )
             # apply pending writes
+            if null_writes := [
+                w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
+            ]:
+                apply_writes(
+                    saved.checkpoint,
+                    channels,
+                    [PregelTaskWrites((), INPUT, null_writes, [])],
+                    None,
+                )
             if apply_pending_writes and saved.pending_writes:
                 for tid, k, v in saved.pending_writes:
                     if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -878,6 +907,7 @@ class Pregel(PregelProtocol):
                     # tasks for this checkpoint
                     next_tasks = prepare_next_tasks(
                         checkpoint,
+                        saved.pending_writes or [],
                         self.nodes,
                         channels,
                         managed,
@@ -888,6 +918,18 @@ class Pregel(PregelProtocol):
                         checkpointer=self.checkpointer or None,
                         manager=None,
                     )
+                    # apply null writes
+                    if null_writes := [
+                        w[1:]
+                        for w in saved.pending_writes or []
+                        if w[0] == NULL_TASK_ID
+                    ]:
+                        apply_writes(
+                            saved.checkpoint,
+                            channels,
+                            [PregelTaskWrites((), INPUT, null_writes, [])],
+                            None,
+                        )
                     # apply writes from tasks that already ran
                     for tid, k, v in saved.pending_writes or []:
                         if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -922,6 +964,7 @@ class Pregel(PregelProtocol):
                 # tasks for this checkpoint
                 next_tasks = prepare_next_tasks(
                     checkpoint,
+                    saved.pending_writes,
                     self.nodes,
                     channels,
                     managed,
@@ -932,6 +975,16 @@ class Pregel(PregelProtocol):
                     checkpointer=self.checkpointer or None,
                     manager=None,
                 )
+                # apply null writes
+                if null_writes := [
+                    w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
+                ]:
+                    apply_writes(
+                        saved.checkpoint,
+                        channels,
+                        [PregelTaskWrites((), INPUT, null_writes, [])],
+                        None,
+                    )
                 # apply writes
                 for tid, k, v in saved.pending_writes:
                     if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -1001,8 +1054,14 @@ class Pregel(PregelProtocol):
                 ),
             )
             # save task writes
-            if saved:
-                checkpointer.put_writes(checkpoint_config, task.writes, task_id)
+            # channel writes are saved to current checkpoint
+            # push writes are saved to next checkpoint
+            channel_writes, push_writes = (
+                [w for w in task.writes if w[0] != PUSH],
+                [w for w in task.writes if w[0] == PUSH],
+            )
+            if saved and channel_writes:
+                checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
             # apply to checkpoint and save
             mv_writes = apply_writes(
                 checkpoint, channels, [task], checkpointer.get_next_version
@@ -1023,6 +1082,8 @@ class Pregel(PregelProtocol):
                     checkpoint_previous_versions, checkpoint["channel_versions"]
                 ),
             )
+            if push_writes:
+                checkpointer.put_writes(next_config, push_writes, task_id)
             return patch_checkpoint_map(next_config, saved.metadata if saved else None)
 
     async def aupdate_state(
@@ -1088,6 +1149,7 @@ class Pregel(PregelProtocol):
                     # tasks for this checkpoint
                     next_tasks = prepare_next_tasks(
                         checkpoint,
+                        saved.pending_writes or [],
                         self.nodes,
                         channels,
                         managed,
@@ -1098,6 +1160,18 @@ class Pregel(PregelProtocol):
                         checkpointer=self.checkpointer or None,
                         manager=None,
                     )
+                    # apply null writes
+                    if null_writes := [
+                        w[1:]
+                        for w in saved.pending_writes or []
+                        if w[0] == NULL_TASK_ID
+                    ]:
+                        apply_writes(
+                            saved.checkpoint,
+                            channels,
+                            [PregelTaskWrites((), INPUT, null_writes, [])],
+                            None,
+                        )
                     # apply writes from tasks that already ran
                     for tid, k, v in saved.pending_writes or []:
                         if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -1132,6 +1206,7 @@ class Pregel(PregelProtocol):
                 # tasks for this checkpoint
                 next_tasks = prepare_next_tasks(
                     checkpoint,
+                    saved.pending_writes,
                     self.nodes,
                     channels,
                     managed,
@@ -1142,6 +1217,16 @@ class Pregel(PregelProtocol):
                     checkpointer=self.checkpointer or None,
                     manager=None,
                 )
+                # apply null writes
+                if null_writes := [
+                    w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
+                ]:
+                    apply_writes(
+                        saved.checkpoint,
+                        channels,
+                        [PregelTaskWrites((), INPUT, null_writes, [])],
+                        None,
+                    )
                 for tid, k, v in saved.pending_writes:
                     if k in (ERROR, INTERRUPT, SCHEDULED):
                         continue
@@ -1208,14 +1293,23 @@ class Pregel(PregelProtocol):
                 ),
             )
             # save task writes
-            if saved:
-                await checkpointer.aput_writes(checkpoint_config, writes, task_id)
+            # channel writes are saved to current checkpoint
+            # push writes are saved to next checkpoint
+            channel_writes, push_writes = (
+                [w for w in task.writes if w[0] != PUSH],
+                [w for w in task.writes if w[0] == PUSH],
+            )
+            if saved and channel_writes:
+                await checkpointer.aput_writes(
+                    checkpoint_config, channel_writes, task_id
+                )
             # apply to checkpoint and save
             mv_writes = apply_writes(
                 checkpoint, channels, [task], checkpointer.get_next_version
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
+            # save checkpoint, after applying writes
             next_config = await checkpointer.aput(
                 checkpoint_config,
                 checkpoint,
@@ -1230,6 +1324,9 @@ class Pregel(PregelProtocol):
                     checkpoint_previous_versions, checkpoint["channel_versions"]
                 ),
             )
+            # save push writes
+            if push_writes:
+                await checkpointer.aput_writes(next_config, push_writes, task_id)
             return patch_checkpoint_map(next_config, saved.metadata if saved else None)
 
     def _defaults(
@@ -1432,12 +1529,16 @@ class Pregel(PregelProtocol):
                 specs=self.channels,
                 output_keys=output_keys,
                 stream_keys=self.stream_channels_asis,
+                interrupt_before=interrupt_before_,
+                interrupt_after=interrupt_after_,
+                manager=run_manager,
                 debug=debug,
             ) as loop:
                 # create runner
                 runner = PregelRunner(
                     submit=loop.submit,
                     put_writes=loop.put_writes,
+                    schedule_task=loop.accept_push,
                     node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
                 )
                 # enable subgraph streaming
@@ -1468,12 +1569,7 @@ class Pregel(PregelProtocol):
                 # channel updates from step N are only visible in step N+1
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
-                while loop.tick(
-                    input_keys=self.input_channels,
-                    interrupt_before=interrupt_before_,
-                    interrupt_after=interrupt_after_,
-                    manager=run_manager,
-                ):
+                while loop.tick(input_keys=self.input_channels):
                     for _ in runner.tick(
                         loop.tasks.values(),
                         timeout=self.step_timeout,
@@ -1654,11 +1750,15 @@ class Pregel(PregelProtocol):
                 specs=self.channels,
                 output_keys=output_keys,
                 stream_keys=self.stream_channels_asis,
+                interrupt_before=interrupt_before_,
+                interrupt_after=interrupt_after_,
+                manager=run_manager,
             ) as loop:
                 # create runner
                 runner = PregelRunner(
                     submit=loop.submit,
                     put_writes=loop.put_writes,
+                    schedule_task=loop.accept_push,
                     use_astream=do_stream is not None,
                     node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
                 )
@@ -1678,12 +1778,7 @@ class Pregel(PregelProtocol):
                 # channel updates from step N are only visible in step N+1
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
-                while loop.tick(
-                    input_keys=self.input_channels,
-                    interrupt_before=interrupt_before_,
-                    interrupt_after=interrupt_after_,
-                    manager=run_manager,
-                ):
+                while loop.tick(input_keys=self.input_channels):
                     async for _ in runner.atick(
                         loop.tasks.values(),
                         timeout=self.step_timeout,

libs/langgraph/langgraph/pregel/algo.py

@@ -25,6 +25,7 @@ from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
+    PendingWrite,
     V,
     copy_checkpoint,
 )
@@ -35,17 +36,21 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_CHECKPOINTER,
     CONFIG_KEY_READ,
+    CONFIG_KEY_RESUME_VALUE,
     CONFIG_KEY_SEND,
     CONFIG_KEY_STORE,
     CONFIG_KEY_TASK_ID,
     EMPTY_SEQ,
     INTERRUPT,
+    MISSING,
     NO_WRITES,
     NS_END,
     NS_SEP,
+    NULL_TASK_ID,
     PULL,
     PUSH,
     RESERVED,
+    RESUME,
     TAG_HIDDEN,
     TASKS,
     Send,
@@ -68,7 +73,7 @@ class WritesProtocol(Protocol):
     Implemented by PregelTaskWrites and PregelExecutableTask."""
 
     @property
-    def path(self) -> tuple[Union[str, int], ...]: ...
+    def path(self) -> tuple[Union[str, int, tuple], ...]: ...
 
     @property
     def name(self) -> str: ...
@@ -84,7 +89,7 @@ class PregelTaskWrites(NamedTuple):
     """Simplest implementation of WritesProtocol, for usage with writes that
     don't originate from a runnable task, eg. graph input, update_state, etc."""
 
-    path: tuple[Union[str, int], ...]
+    path: tuple[Union[str, int, tuple], ...]
     name: str
     writes: Sequence[tuple[str, Any]]
     triggers: Sequence[str]
@@ -172,7 +177,7 @@ def local_write(
     """Function injected under CONFIG_KEY_SEND in task config, to write to channels.
     Validates writes and forwards them to `commit` function."""
     for chan, value in writes:
-        if chan == TASKS:
+        if chan in (PUSH, TASKS):
             if not isinstance(value, Send):
                 raise InvalidUpdateError(f"Expected Send, got {value}")
             if value.node not in process_keys:
@@ -194,8 +199,13 @@ def apply_writes(
     """Apply writes from a set of tasks (usually the tasks from a Pregel step)
     to the checkpoint and channels, and return managed values writes to be applied
     externally."""
-    # sort tasks on path
-    tasks = sorted(tasks, key=lambda t: t.path)
+    # sort tasks on path, to ensure deterministic order for update application
+    # any path parts after the 3rd are ignored for sorting
+    # (we use them for eg. task ids which aren't good for sorting)
+    tasks = sorted(tasks, key=lambda t: t.path[:3])
+    # if no task has triggers this is applying writes from the null task only
+    # so we don't do anything other than update the channels written to
+    bump_step = any(t.triggers for t in tasks)
 
     # update seen versions
     for task in tasks:
@@ -227,7 +237,7 @@ def apply_writes(
             )
 
     # clear pending sends
-    if checkpoint["pending_sends"]:
+    if checkpoint["pending_sends"] and bump_step:
         checkpoint["pending_sends"].clear()
 
     # Group writes by channel
@@ -235,9 +245,9 @@ def apply_writes(
     pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
-            if chan == NO_WRITES:
+            if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT):
                 pass
-            elif chan == TASKS:
+            elif chan == TASKS:  # TODO: remove branch in 1.0
                 checkpoint["pending_sends"].append(val)
             elif chan in channels:
                 pending_writes_by_channel[chan].append(val)
@@ -262,13 +272,14 @@ def apply_writes(
             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],
-                )
+    if bump_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],
+                    )
 
     # Return managed values writes to be applied externally
     return pending_writes_by_managed
@@ -277,6 +288,7 @@ def apply_writes(
 @overload
 def prepare_next_tasks(
     checkpoint: Checkpoint,
+    pending_writes: Sequence[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -293,6 +305,7 @@ def prepare_next_tasks(
 @overload
 def prepare_next_tasks(
     checkpoint: Checkpoint,
+    pending_writes: Sequence[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -308,6 +321,7 @@ def prepare_next_tasks(
 
 def prepare_next_tasks(
     checkpoint: Checkpoint,
+    pending_writes: Sequence[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -322,13 +336,14 @@ def prepare_next_tasks(
     """Prepare the set of tasks that will make up the next Pregel step.
     This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
     by edges)."""
-    tasks: dict[str, Union[PregelTask, PregelExecutableTask]] = {}
-    # Consume pending packets
-    for idx, _ in enumerate(checkpoint["pending_sends"]):
+    tasks: list[Union[PregelTask, PregelExecutableTask]] = []
+    # Consume pending_sends from previous step (legacy version of Send)
+    for idx, _ in enumerate(checkpoint["pending_sends"]):  # TODO: remove branch in 1.0
         if task := prepare_single_task(
             (PUSH, idx),
             None,
             checkpoint=checkpoint,
+            pending_writes=pending_writes,
             processes=processes,
             channels=channels,
             managed=managed,
@@ -339,7 +354,7 @@ def prepare_next_tasks(
             checkpointer=checkpointer,
             manager=manager,
         ):
-            tasks[task.id] = task
+            tasks.append(task)
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
     for name in processes:
@@ -347,6 +362,7 @@ def prepare_next_tasks(
             (PULL, name),
             None,
             checkpoint=checkpoint,
+            pending_writes=pending_writes,
             processes=processes,
             channels=channels,
             managed=managed,
@@ -357,15 +373,74 @@ def prepare_next_tasks(
             checkpointer=checkpointer,
             manager=manager,
         ):
-            tasks[task.id] = task
-    return tasks
+            tasks.append(task)
+    # Consume pending Sends from this step (new version of Send)
+    if any(c == PUSH for _, c, _ in pending_writes):
+        # group writes by task id
+        grouped_by_task = defaultdict(list)
+        for tid, c, _ in pending_writes:
+            grouped_by_task[tid].append(c)
+        # prepare send tasks from grouped writes
+        # 1. start from sends originating from existing tasks
+        tidx = 0
+        while tidx < len(tasks):
+            task = tasks[tidx]
+            if twrites := grouped_by_task.pop(task.id, None):
+                for idx, c in enumerate(twrites):
+                    if c != PUSH:
+                        continue
+                    if next_task := prepare_single_task(
+                        (PUSH, task.path, idx, task.id),
+                        None,
+                        checkpoint=checkpoint,
+                        pending_writes=pending_writes,
+                        processes=processes,
+                        channels=channels,
+                        managed=managed,
+                        config=config,
+                        step=step,
+                        for_execution=for_execution,
+                        store=store,
+                        checkpointer=checkpointer,
+                        manager=manager,
+                    ):
+                        tasks.append(next_task)
+            tidx += 1
+        # key tasks by id
+        task_map = {t.id: t for t in tasks}
+        # 2. create new tasks for remaining sends (eg. from update_state)
+        for tid, writes in grouped_by_task.items():
+            task = task_map.get(tid)
+            for idx, c in enumerate(writes):
+                if c != PUSH:
+                    continue
+                if next_task := prepare_single_task(
+                    (PUSH, task.path if task else (), idx, tid),
+                    None,
+                    checkpoint=checkpoint,
+                    pending_writes=pending_writes,
+                    processes=processes,
+                    channels=channels,
+                    managed=managed,
+                    config=config,
+                    step=step,
+                    for_execution=for_execution,
+                    store=store,
+                    checkpointer=checkpointer,
+                    manager=manager,
+                ):
+                    task_map[next_task.id] = next_task
+    else:
+        task_map = {t.id: t for t in tasks}
+    return task_map
 
 
 def prepare_single_task(
-    task_path: tuple[str, Union[int, str]],
+    task_path: tuple[Union[str, int, tuple], ...],
     task_id_checksum: Optional[str],
     *,
     checkpoint: Checkpoint,
+    pending_writes: Sequence[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -383,31 +458,74 @@ def prepare_single_task(
     parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
 
     if task_path[0] == PUSH:
-        idx = int(task_path[1])
-        if idx >= len(checkpoint["pending_sends"]):
-            return
-        packet = checkpoint["pending_sends"][idx]
-        if not isinstance(packet, Send):
-            logger.warning(
-                f"Ignoring invalid packet type {type(packet)} in pending sends"
+        if len(task_path) == 2:  # TODO: remove branch in 1.0
+            # legacy SEND tasks, executed in superstep n+1
+            # (PUSH, idx of pending send)
+            idx = cast(int, task_path[1])
+            if idx >= len(checkpoint["pending_sends"]):
+                return
+            packet = checkpoint["pending_sends"][idx]
+            if not isinstance(packet, Send):
+                logger.warning(
+                    f"Ignoring invalid packet type {type(packet)} in pending sends"
+                )
+                return
+            if packet.node not in processes:
+                logger.warning(
+                    f"Ignoring unknown node name {packet.node} in pending sends"
+                )
+                return
+            # create task id
+            triggers = [PUSH]
+            checkpoint_ns = (
+                f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
             )
+            task_id = _uuid5_str(
+                checkpoint_id,
+                checkpoint_ns,
+                str(step),
+                packet.node,
+                PUSH,
+                str(idx),
+            )
+        elif len(task_path) == 4:
+            # new PUSH tasks, executed in superstep n
+            # (PUSH, parent task path, idx of PUSH write, id of parent task)
+            task_path_t = cast(tuple[str, tuple, int, str], task_path)
+            writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
+            if task_path_t[2] >= len(writes_for_path):
+                logger.warning(
+                    f"Ignoring invalid write index {task_path[2]} in pending writes"
+                )
+                return
+            packet = writes_for_path[task_path_t[2]][2]
+            if not isinstance(packet, Send):
+                logger.warning(
+                    f"Ignoring invalid packet type {type(packet)} in pending writes"
+                )
+                return
+            if packet.node not in processes:
+                logger.warning(
+                    f"Ignoring unknown node name {packet.node} in pending writes"
+                )
+                return
+            # create task id
+            triggers = [PUSH]
+            checkpoint_ns = (
+                f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
+            )
+            task_id = _uuid5_str(
+                checkpoint_id,
+                checkpoint_ns,
+                str(step),
+                packet.node,
+                PUSH,
+                _tuple_str(task_path[1]),
+                str(task_path[2]),
+            )
+        else:
+            logger.warning(f"Ignoring invalid PUSH task path {task_path}")
             return
-        if packet.node not in processes:
-            logger.warning(f"Ignoring unknown node name {packet.node} in pending sends")
-            return
-        # create task id
-        triggers = [PUSH]
-        checkpoint_ns = (
-            f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
-        )
-        task_id = _uuid5_str(
-            checkpoint_id,
-            checkpoint_ns,
-            str(step),
-            packet.node,
-            PUSH,
-            str(idx),
-        )
         task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
         metadata = {
             "langgraph_step": step,
@@ -417,7 +535,7 @@ def prepare_single_task(
             "langgraph_checkpoint_ns": task_checkpoint_ns,
         }
         if task_id_checksum is not None:
-            assert task_id == task_id_checksum
+            assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
         if for_execution:
             proc = processes[packet.node]
             if node := proc.node:
@@ -469,6 +587,14 @@ def prepare_single_task(
                             },
                             CONFIG_KEY_CHECKPOINT_ID: None,
                             CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
+                            CONFIG_KEY_RESUME_VALUE: next(
+                                (
+                                    v
+                                    for tid, c, v in pending_writes
+                                    if tid in (NULL_TASK_ID, task_id) and c == RESUME
+                                ),
+                                MISSING,
+                            ),
                         },
                     ),
                     triggers,
@@ -481,6 +607,7 @@ def prepare_single_task(
         else:
             return PregelTask(task_id, packet.node, task_path)
     elif task_path[0] == PULL:
+        # (PULL, node name)
         name = cast(str, task_path[1])
         if name not in processes:
             return
@@ -577,6 +704,15 @@ def prepare_single_task(
                                 },
                                 CONFIG_KEY_CHECKPOINT_ID: None,
                                 CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
+                                CONFIG_KEY_RESUME_VALUE: next(
+                                    (
+                                        v
+                                        for tid, c, v in pending_writes
+                                        if tid in (NULL_TASK_ID, task_id)
+                                        and c == RESUME
+                                    ),
+                                    MISSING,
+                                ),
                             },
                         ),
                         triggers,
@@ -642,3 +778,12 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
     sha.update(b"".join(p.encode() for p in parts))
     hex = sha.hexdigest()
     return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}"
+
+
+def _tuple_str(tup: Union[str, int, tuple]) -> str:
+    """Generate a string representation of a tuple."""
+    return (
+        f"({', '.join(_tuple_str(x) for x in tup)})"
+        if isinstance(tup, (tuple, list))
+        else str(tup)
+    )

libs/langgraph/langgraph/pregel/debug.py

@@ -208,7 +208,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None:
     print(
         f"{get_colored_text(f'[{step}:tasks]', color='blue')} "
         + get_bolded_text(
-            f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n"
+            f"Starting {n_tasks} task{'s' if n_tasks != 1 else ''} for step {step}:\n"
         )
         + "\n".join(
             f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}"

libs/langgraph/langgraph/pregel/io.py

@@ -1,11 +1,31 @@
 from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union
+from uuid import UUID
 
 from langchain_core.runnables.utils import AddableDict
 
 from langgraph.channels.base import BaseChannel, EmptyChannelError
-from langgraph.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN
+from langgraph.constants import (
+    EMPTY_SEQ,
+    ERROR,
+    FF_SEND_V2,
+    INTERRUPT,
+    NULL_TASK_ID,
+    PUSH,
+    RESUME,
+    TAG_HIDDEN,
+    TASKS,
+)
 from langgraph.pregel.log import logger
-from langgraph.types import PregelExecutableTask
+from langgraph.types import Command, PregelExecutableTask, Send
+
+
+def is_task_id(task_id: str) -> bool:
+    """Check if a string is a valid task id."""
+    try:
+        UUID(task_id)
+    except ValueError:
+        return False
+    return True
 
 
 def read_channel(
@@ -44,6 +64,44 @@ def read_channels(
         return values
 
 
+def map_command(
+    cmd: Command,
+) -> Iterator[tuple[str, str, Any]]:
+    """Map input chunk to a sequence of pending writes in the form (channel, value)."""
+    if cmd.send:
+        if isinstance(cmd.send, (tuple, list)) and all(
+            isinstance(x, Send)
+            or isinstance(x, (list, tuple))
+            and len(x) == 2
+            and isinstance(x[0], str)
+            for x in cmd.send
+        ):
+            sends = cmd.send
+        else:
+            sends = [cmd.send]
+        for send in sends:
+            if isinstance(send, tuple) and len(send) == 2 and isinstance(send[0], str):
+                send = Send(*send)
+            if not isinstance(send, Send):
+                raise TypeError(
+                    f"In Command.send, expected Send, got {type(send).__name__}"
+                )
+            yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
+    if cmd.resume:
+        if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
+            for tid, resume in cmd.resume.items():
+                yield (tid, RESUME, resume)
+        else:
+            yield (NULL_TASK_ID, RESUME, cmd.resume)
+    if cmd.update:
+        if not isinstance(cmd.update, dict):
+            raise TypeError(
+                f"Expected cmd.update to be a dict mapping channel names to update values, got {type(cmd.update).__name__}"
+            )
+        for k, v in cmd.update.items():
+            yield (NULL_TASK_ID, k, v)
+
+
 def map_input(
     input_channels: Union[str, Sequence[str]],
     chunk: Optional[Union[dict[str, Any], Any]],

libs/langgraph/langgraph/pregel/loop.py

@@ -1,6 +1,6 @@
 import asyncio
 import concurrent.futures
-from collections import deque
+from collections import defaultdict, deque
 from contextlib import AsyncExitStack, ExitStack
 from types import TracebackType
 from typing import (
@@ -52,6 +52,9 @@ from langgraph.constants import (
     INPUT,
     INTERRUPT,
     NS_SEP,
+    NULL_TASK_ID,
+    PUSH,
+    RESUME,
     SCHEDULED,
     TAG_HIDDEN,
 )
@@ -74,6 +77,7 @@ from langgraph.pregel.algo import (
     apply_writes,
     increment,
     prepare_next_tasks,
+    prepare_single_task,
     should_interrupt,
 )
 from langgraph.pregel.debug import (
@@ -90,6 +94,7 @@ from langgraph.pregel.executor import (
     Submit,
 )
 from langgraph.pregel.io import (
+    map_command,
     map_input,
     map_output_updates,
     map_output_values,
@@ -100,7 +105,13 @@ from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.store.base import BaseStore
-from langgraph.types import All, LoopProtocol, PregelExecutableTask, StreamProtocol
+from langgraph.types import (
+    All,
+    Command,
+    LoopProtocol,
+    PregelExecutableTask,
+    StreamProtocol,
+)
 from langgraph.utils.config import patch_configurable
 
 V = TypeVar("V")
@@ -130,6 +141,9 @@ class PregelLoop(LoopProtocol):
     stream_keys: Union[str, Sequence[str]]
     skip_done_tasks: bool
     is_nested: bool
+    manager: Union[None, AsyncParentRunManager, ParentRunManager]
+    interrupt_after: Union[All, Sequence[str]]
+    interrupt_before: Union[All, Sequence[str]]
 
     checkpointer_get_next_version: GetNextVersion
     checkpointer_put_writes: Optional[
@@ -162,6 +176,7 @@ class PregelLoop(LoopProtocol):
         "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
     ]
     tasks: dict[str, PregelExecutableTask]
+    to_interrupt: list[PregelExecutableTask]
     output: Union[None, dict[str, Any], Any] = None
 
     # public
@@ -178,6 +193,9 @@ class PregelLoop(LoopProtocol):
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]],
         stream_keys: Union[str, Sequence[str]],
+        interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
+        interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
@@ -194,6 +212,9 @@ class PregelLoop(LoopProtocol):
         self.specs = specs
         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.skip_done_tasks = (
             CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -261,15 +282,60 @@ class PregelLoop(LoopProtocol):
                 task_id,
             )
         # output writes
-        self._output_writes(task_id, writes)
+        if hasattr(self, "tasks"):
+            self._output_writes(task_id, writes)
+
+    def accept_push(
+        self, task: PregelExecutableTask, write_idx: int
+    ) -> Optional[PregelExecutableTask]:
+        """Accept a PUSH from a task, potentially returning a new task to start."""
+        # don't start if an earlier PUSH has already triggered an interrupt
+        if self.to_interrupt:
+            return
+        # don't start if we should interrupt *after* the original task
+        if should_interrupt(self.checkpoint, self.interrupt_after, [task]):
+            self.to_interrupt.append(task)
+            return
+        if pushed := cast(
+            Optional[PregelExecutableTask],
+            prepare_single_task(
+                (PUSH, task.path, write_idx, task.id),
+                None,
+                checkpoint=self.checkpoint,
+                pending_writes=[(task.id, *w) for w in task.writes],
+                processes=self.nodes,
+                channels=self.channels,
+                managed=self.managed,
+                config=self.config,
+                step=self.step,
+                for_execution=True,
+                store=self.store,
+                checkpointer=self.checkpointer,
+                manager=self.manager,
+            ),
+        ):
+            # don't start if we should interrupt *before* the new task
+            if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]):
+                self.to_interrupt.append(pushed)
+                return
+            # produce debug output
+            self._emit("debug", map_debug_tasks, self.step, [pushed])
+            # debug flag
+            if self.debug:
+                print_step_tasks(self.step, [pushed])
+            # save the new task
+            self.tasks[pushed.id] = pushed
+            # match any pending writes to the new task
+            if self.skip_done_tasks:
+                self._match_writes({pushed.id: pushed})
+            # return the new task, to be started, if not run before
+            if not pushed.writes:
+                return pushed
 
     def tick(
         self,
         *,
         input_keys: Union[str, Sequence[str]],
-        interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
-        interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
-        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
     ) -> bool:
         """Execute a single iteration of the Pregel loop.
         Returns True if more iterations are needed."""
@@ -278,6 +344,10 @@ class PregelLoop(LoopProtocol):
 
         if self.input not in (INPUT_DONE, INPUT_RESUMING):
             self._first(input_keys=input_keys)
+        elif self.to_interrupt:
+            # if we need to interrupt, do so
+            self.status = "interrupt_before"
+            raise GraphInterrupt()
         elif all(task.writes for task in self.tasks.values()):
             writes = [w for t in self.tasks.values() for w in t.writes]
             # debug flag
@@ -322,7 +392,9 @@ class PregelLoop(LoopProtocol):
                 }
             )
             # after execution, check if we should interrupt
-            if should_interrupt(self.checkpoint, interrupt_after, self.tasks.values()):
+            if should_interrupt(
+                self.checkpoint, self.interrupt_after, self.tasks.values()
+            ):
                 self.status = "interrupt_after"
                 raise GraphInterrupt()
         else:
@@ -333,19 +405,33 @@ class PregelLoop(LoopProtocol):
             self.status = "out_of_steps"
             return False
 
+        # apply NULL writes
+        if null_writes := [
+            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
+        ]:
+            mv_writes = apply_writes(
+                self.checkpoint,
+                self.channels,
+                [PregelTaskWrites((), INPUT, null_writes, [])],
+                self.checkpointer_get_next_version,
+            )
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
+            self.checkpoint_pending_writes,
             self.nodes,
             self.channels,
             self.managed,
             self.config,
             self.step,
             for_execution=True,
-            manager=manager,
+            manager=self.manager,
             store=self.store,
             checkpointer=self.checkpointer,
         )
+        self.to_interrupt = []
 
         # produce debug output
         if self._checkpointer_put_after_previous is not None:
@@ -387,15 +473,12 @@ class PregelLoop(LoopProtocol):
 
         # if all tasks have finished, re-tick
         if all(task.writes for task in self.tasks.values()):
-            return self.tick(
-                input_keys=input_keys,
-                interrupt_after=interrupt_after,
-                interrupt_before=interrupt_before,
-                manager=manager,
-            )
+            return self.tick(input_keys=input_keys)
 
         # before execution, check if we should interrupt
-        if should_interrupt(self.checkpoint, interrupt_before, self.tasks.values()):
+        if should_interrupt(
+            self.checkpoint, self.interrupt_before, self.tasks.values()
+        ):
             self.status = "interrupt_before"
             raise GraphInterrupt()
 
@@ -417,7 +500,7 @@ class PregelLoop(LoopProtocol):
 
     def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None:
         for tid, k, v in self.checkpoint_pending_writes:
-            if k in (ERROR, INTERRUPT):
+            if k in (ERROR, INTERRUPT, RESUME):
                 continue
             if task := tasks.get(tid):
                 if k == SCHEDULED:
@@ -449,8 +532,20 @@ class PregelLoop(LoopProtocol):
             self._emit(
                 "values", map_output_values, self.output_keys, True, self.channels
             )
+        # map command to writes
+        elif isinstance(self.input, Command):
+            writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
+            # group writes by task ID
+            for tid, c, v in map_command(self.input):
+                writes[tid].append((c, v))
+            if not writes:
+                raise EmptyInputError("Received empty Command input")
+            # save writes
+            for tid, ws in writes.items():
+                self.put_writes(tid, ws)
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
+            # TODO shouldn't these writes be passed to put_writes too?
             # check if we should delegate (used by subgraphs in distributed mode)
             if self.config[CONF].get(CONFIG_KEY_DELEGATE):
                 raise GraphDelegate(
@@ -464,6 +559,7 @@ class PregelLoop(LoopProtocol):
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
+                self.checkpoint_pending_writes,
                 self.nodes,
                 self.channels,
                 self.managed,
@@ -577,11 +673,33 @@ class PregelLoop(LoopProtocol):
             # save final output
             self.output = read_channels(self.channels, self.output_keys)
         if suppress:
-            # suppress interrupt
+            # 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())
+            ):
+                mv_writes = apply_writes(
+                    self.checkpoint,
+                    self.channels,
+                    self.tasks.values(),
+                    self.checkpointer_get_next_version,
+                )
+                for key, values in mv_writes.items():
+                    self._update_mv(key, values)
+                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 event
             self._emit(
                 "updates",
                 lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
             )
+            # suppress interrupt
             return True
 
     def _emit(
@@ -635,6 +753,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
+        interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
+        interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         check_subgraphs: bool = True,
@@ -650,7 +771,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             specs=specs,
             output_keys=output_keys,
             stream_keys=stream_keys,
+            interrupt_after=interrupt_after,
+            interrupt_before=interrupt_before,
             check_subgraphs=check_subgraphs,
+            manager=manager,
             debug=debug,
         )
         self.stack = ExitStack()
@@ -761,6 +885,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
+        interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
+        manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         check_subgraphs: bool = True,
@@ -776,7 +903,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             specs=specs,
             output_keys=output_keys,
             stream_keys=stream_keys,
+            interrupt_after=interrupt_after,
+            interrupt_before=interrupt_before,
             check_subgraphs=check_subgraphs,
+            manager=manager,
             debug=debug,
         )
         self.stack = AsyncExitStack()

libs/langgraph/langgraph/pregel/retry.py

@@ -2,9 +2,15 @@ import asyncio
 import logging
 import random
 import time
-from typing import Optional, Sequence
+from functools import partial
+from typing import Any, Callable, Optional, Sequence
 
-from langgraph.constants import CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING
+from langgraph.constants import (
+    CONF,
+    CONFIG_KEY_CHECKPOINT_NS,
+    CONFIG_KEY_RESUMING,
+    CONFIG_KEY_SEND,
+)
 from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt
 from langgraph.types import PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
@@ -15,12 +21,17 @@ logger = logging.getLogger(__name__)
 def run_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
+    writer: Optional[
+        Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
+    ] = None,
 ) -> None:
     """Run a task with retries."""
     retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     config = task.config
+    if writer is not None:
+        config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
     while True:
         try:
             # clear any writes from previous attempts
@@ -84,12 +95,17 @@ async def arun_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
     stream: bool = False,
+    writer: Optional[
+        Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
+    ] = None,
 ) -> None:
     """Run a task asynchronously with retries."""
     retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     config = task.config
+    if writer is not None:
+        config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
     while True:
         try:
             # clear any writes from previous attempts

libs/langgraph/langgraph/pregel/runner.py

@@ -14,7 +14,15 @@ from typing import (
     cast,
 )
 
-from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN
+from langgraph.constants import (
+    CONF,
+    CONFIG_KEY_SEND,
+    ERROR,
+    INTERRUPT,
+    NO_WRITES,
+    PUSH,
+    TAG_HIDDEN,
+)
 from langgraph.errors import GraphDelegate, GraphInterrupt
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
@@ -31,6 +39,9 @@ class PregelRunner:
         *,
         submit: Submit,
         put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
+        schedule_task: Callable[
+            [PregelExecutableTask, int], Optional[PregelExecutableTask]
+        ],
         use_astream: bool = False,
         node_finished: Optional[Callable[[str], None]] = None,
     ) -> None:
@@ -38,6 +49,7 @@ class PregelRunner:
         self.put_writes = put_writes
         self.use_astream = use_astream
         self.node_finished = node_finished
+        self.schedule_task = schedule_task
 
     def tick(
         self,
@@ -48,27 +60,58 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
     ) -> Iterator[None]:
+        def writer(
+            task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
+        ) -> None:
+            prev_length = len(task.writes)
+            # delegate to the underlying writer
+            task.config[CONF][CONFIG_KEY_SEND](writes)
+            for idx, w in enumerate(task.writes):
+                # find the index for the newly inserted writes
+                if idx < prev_length:
+                    continue
+                assert writes[idx - prev_length] is w
+                # bail if not a PUSH write
+                if w[0] != PUSH:
+                    continue
+                # schedule the next task, if the callback returns one
+                if next_task := self.schedule_task(task, idx):
+                    # if the parent task was retried,
+                    # the next task might already be running
+                    if any(
+                        t == next_task.id for t in futures.values() if t is not None
+                    ):
+                        continue
+                    # schedule the next task
+                    futures[
+                        self.submit(
+                            run_with_retry,
+                            next_task,
+                            retry_policy,
+                            writer=writer,
+                            __reraise_on_exit__=reraise,
+                        )
+                    ] = next_task
+
         tasks = tuple(tasks)
+        futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
         # give control back to the caller
         yield
         # fast path if single task with no timeout and no waiter
         if len(tasks) == 1 and timeout is None and get_waiter is None:
             t = tasks[0]
             try:
-                run_with_retry(t, retry_policy)
+                run_with_retry(t, retry_policy, writer=writer)
                 self.commit(t, None)
             except Exception as exc:
                 self.commit(t, exc)
                 if reraise:
                     raise
-            return
+            if not futures:  # maybe `t` schuduled another task
+                return
         # add waiter task if requested
         if get_waiter is not None:
-            futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {
-                get_waiter(): None
-            }
-        else:
-            futures = {}
+            futures[get_waiter()] = None
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
@@ -79,10 +122,11 @@ class PregelRunner:
                         run_with_retry,
                         t,
                         retry_policy,
+                        writer=writer,
                         __reraise_on_exit__=reraise,
                     )
                 ] = t
-        all_futures = futures.copy()
+        done_futures: set[concurrent.futures.Future] = set()
         end_time = timeout + time.monotonic() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = concurrent.futures.wait(
@@ -99,6 +143,8 @@ class PregelRunner:
                     if inflight and get_waiter is not None:
                         futures[get_waiter()] = None
                 else:
+                    # store for panic check
+                    done_futures.add(fut)
                     # task finished, commit writes
                     self.commit(task, _exception(fut))
             else:
@@ -110,7 +156,10 @@ class PregelRunner:
             # give control back to the caller
             yield
         # panic on failure or timeout
-        _panic_or_proceed(all_futures, panic=reraise)
+        _panic_or_proceed(
+            done_futures.union(f for f, t in futures.items() if t is not None),
+            panic=reraise,
+        )
 
     async def atick(
         self,
@@ -121,28 +170,67 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
     ) -> AsyncIterator[None]:
+        def writer(
+            task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
+        ) -> None:
+            prev_length = len(task.writes)
+            # delegate to the underlying writer
+            task.config[CONF][CONFIG_KEY_SEND](writes)
+            for idx, w in enumerate(task.writes):
+                # find the index for the newly inserted writes
+                if idx < prev_length:
+                    continue
+                assert writes[idx - prev_length] is w
+                # bail if not a PUSH write
+                if w[0] != PUSH:
+                    continue
+                # schedule the next task, if the callback returns one
+                if next_task := self.schedule_task(task, idx):
+                    # if the parent task was retried,
+                    # the next task might already be running
+                    if any(
+                        t == next_task.id for t in futures.values() if t is not None
+                    ):
+                        continue
+                    # schedule the next task
+                    futures[
+                        cast(
+                            asyncio.Future,
+                            self.submit(
+                                arun_with_retry,
+                                next_task,
+                                retry_policy,
+                                stream=self.use_astream,
+                                writer=writer,
+                                __name__=t.name,
+                                __cancel_on_exit__=True,
+                                __reraise_on_exit__=reraise,
+                            ),
+                        )
+                    ] = next_task
+
         loop = asyncio.get_event_loop()
         tasks = tuple(tasks)
+        futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {}
         # give control back to the caller
         yield
         # fast path if single task with no waiter and no timeout
         if len(tasks) == 1 and get_waiter is None and timeout is None:
             t = tasks[0]
             try:
-                await arun_with_retry(t, retry_policy, stream=self.use_astream)
+                await arun_with_retry(
+                    t, retry_policy, stream=self.use_astream, writer=writer
+                )
                 self.commit(t, None)
             except Exception as exc:
                 self.commit(t, exc)
                 if reraise:
                     raise
-            return
+            if not futures:  # maybe `t` schuduled another task
+                return
         # add waiter task if requested
         if get_waiter is not None:
-            futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {
-                get_waiter(): None
-            }
-        else:
-            futures = {}
+            futures[get_waiter()] = None
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
@@ -156,13 +244,14 @@ class PregelRunner:
                             t,
                             retry_policy,
                             stream=self.use_astream,
+                            writer=writer,
                             __name__=t.name,
                             __cancel_on_exit__=True,
                             __reraise_on_exit__=reraise,
                         ),
                     )
                 ] = t
-        all_futures = futures.copy()
+        done_futures: set[asyncio.Future] = set()
         end_time = timeout + loop.time() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = await asyncio.wait(
@@ -179,6 +268,8 @@ class PregelRunner:
                     if inflight and get_waiter is not None:
                         futures[get_waiter()] = None
                 else:
+                    # store for panic check
+                    done_futures.add(fut)
                     # task finished, commit writes
                     self.commit(task, _exception(fut))
             else:
@@ -194,7 +285,9 @@ class PregelRunner:
             fut.cancel()
         # panic on failure or timeout
         _panic_or_proceed(
-            all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise
+            done_futures.union(f for f, t in futures.items() if t is not None),
+            timeout_exc_cls=asyncio.TimeoutError,
+            panic=reraise,
         )
 
     def commit(
@@ -250,10 +343,7 @@ def _exception(
 
 
 def _panic_or_proceed(
-    futs: Union[
-        dict[concurrent.futures.Future, Optional[PregelExecutableTask]],
-        dict[asyncio.Future, Optional[PregelExecutableTask]],
-    ],
+    futs: Union[set[concurrent.futures.Future], set[asyncio.Future]],
     *,
     timeout_exc_cls: Type[Exception] = TimeoutError,
     panic: bool = True,
@@ -261,10 +351,8 @@ def _panic_or_proceed(
     """Cancel remaining tasks if any failed, re-raise exception if panic is True."""
     done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
-    for fut, val in futs.items():
-        if val is None:
-            continue
-        elif fut.done():
+    for fut in futs:
+        if fut.done():
             done.add(fut)
         else:
             inflight.add(fut)

libs/langgraph/langgraph/pregel/write.py

@@ -14,7 +14,7 @@ from typing import (
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 
-from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
+from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send
 from langgraph.errors import InvalidUpdateError
 from langgraph.utils.runnable import RunnableCallable
 
@@ -112,14 +112,18 @@ class ChannelWrite(RunnableCallable):
         # validate
         for w in writes:
             if isinstance(w, ChannelWriteEntry):
-                if w.channel == TASKS:
+                if w.channel in (TASKS, PUSH):
                     raise InvalidUpdateError(
                         "Cannot write to the reserved channel TASKS"
                     )
                 if w.value is PASSTHROUGH:
                     raise InvalidUpdateError("PASSTHROUGH value must be replaced")
         # split packets and entries
-        sends = [(TASKS, packet) for packet in writes if isinstance(packet, Send)]
+        sends = [
+            (PUSH if FF_SEND_V2 else TASKS, packet)
+            for packet in writes
+            if isinstance(packet, Send)
+        ]
         entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
         # process entries into values
         values = [

libs/langgraph/langgraph/types.py

@@ -1,14 +1,18 @@
+import dataclasses
+import sys
 from collections import deque
-from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    Generic,
+    Hashable,
     Literal,
     NamedTuple,
     Optional,
     Sequence,
     Type,
+    TypeVar,
     Union,
     cast,
 )
@@ -44,6 +48,11 @@ StreamWriter = Callable[[Any], None]
 Always injected into nodes if requested as a keyword argument, but it's a no-op
 when not using stream_mode="custom"."""
 
+if sys.version_info >= (3, 10):
+    _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
+else:
+    _DC_KWARGS = {"frozen": True}
+
 
 def default_retry_on(exc: Exception) -> bool:
     import httpx
@@ -101,16 +110,18 @@ class CachePolicy(NamedTuple):
     pass
 
 
-@dataclass
+@dataclasses.dataclass(**_DC_KWARGS)
 class Interrupt:
     value: Any
+    resumable: bool = False
+    ns: Optional[Sequence[str]] = None
     when: Literal["during"] = "during"
 
 
 class PregelTask(NamedTuple):
     id: str
     name: str
-    path: tuple[Union[str, int], ...]
+    path: tuple[Union[str, int, tuple], ...]
     error: Optional[Exception] = None
     interrupts: tuple[Interrupt, ...] = ()
     state: Union[None, RunnableConfig, "StateSnapshot"] = None
@@ -127,7 +138,7 @@ class PregelExecutableTask(NamedTuple):
     retry_policy: Optional[RetryPolicy]
     cache_policy: Optional[CachePolicy]
     id: str
-    path: tuple[Union[str, int], ...]
+    path: tuple[Union[str, int, tuple], ...]
     scheduled: bool = False
 
 
@@ -221,51 +232,26 @@ class Send:
         )
 
 
-class Command:
+N = TypeVar("N", bound=Hashable)
+
+
+@dataclasses.dataclass(**_DC_KWARGS)
+class Command(Generic[N]):
     """One or more commands to update the graph's state and send messages to nodes."""
 
-    __slots__ = ("update", "send")
-
-    def __init__(
-        self,
-        *,
-        update: Optional[dict[str, Any]] = None,
-        send: Union[Send, Sequence[Send]] = (),
-    ) -> None:
-        self.update = update
-        self.send = send
-
-    @property
-    def __all_slots__(self) -> set[str]:
-        # get all slots from mro
-        slots = set()
-        for cls in type(self).__mro__:
-            if ss := getattr(cls, "__slots__", ()):
-                if isinstance(ss, str):
-                    slots.add(ss)
-                else:
-                    slots.update(ss)
-        return slots
+    update: Optional[dict[str, Any]] = None
+    send: Union[Send, Sequence[Send]] = ()
+    resume: Optional[Union[Any, dict[str, Any]]] = None
 
     def __repr__(self) -> str:
         # get all non-None values
         contents = ", ".join(
             f"{key}={value!r}"
-            for key in self.__all_slots__
-            if (value := getattr(self, key))
+            for key, value in dataclasses.asdict(self).items()
+            if value
         )
         return f"Command({contents})"
 
-    def __eq__(self, value: Any) -> bool:
-        return type(value) is type(self) and all(
-            getattr(self, key) == getattr(value, key) for key in self.__all_slots__
-        )
-
-    def copy(self, **kwargs: Any) -> Self:
-        for slot in self.__all_slots__:
-            kwargs.setdefault(slot, getattr(self, slot))
-        return self.__class__(**kwargs)
-
 
 StreamChunk = tuple[tuple[str, ...], str, Any]
 
@@ -307,3 +293,28 @@ class LoopProtocol:
         self.store = store
         self.step = step
         self.stop = stop
+
+
+def interrupt(value: Any) -> Any:
+    from langgraph.constants import (
+        CONFIG_KEY_CHECKPOINT_NS,
+        CONFIG_KEY_RESUME_VALUE,
+        MISSING,
+        NS_SEP,
+    )
+    from langgraph.errors import GraphInterrupt
+    from langgraph.utils.config import get_configurable
+
+    conf = get_configurable()
+    if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING:
+        return resume
+    else:
+        raise GraphInterrupt(
+            (
+                Interrupt(
+                    value=value,
+                    resumable=True,
+                    ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
+                ),
+            )
+        )

libs/langgraph/langgraph/utils/config.py

@@ -1,3 +1,5 @@
+import asyncio
+import sys
 from collections import ChainMap
 from typing import Any, Optional, Sequence
 
@@ -290,3 +292,18 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig:
         ):
             empty["metadata"][key] = value
     return empty
+
+
+def get_configurable() -> dict[str, Any]:
+    if sys.version_info < (3, 11):
+        try:
+            if asyncio.current_task():
+                raise RuntimeError(
+                    "Python 3.11 or later required to use this in an async context"
+                )
+        except RuntimeError:
+            pass
+    if var_config := var_child_runnable_config.get():
+        return var_config[CONF]
+    else:
+        raise RuntimeError("Called get_configurable outside of a runnable context")

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -5108,6 +5108,81 @@
   
   '''
 # ---
+# name: test_send_react_interrupt_control[memory]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres_pipe]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres_pool]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[sqlite]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
 # name: test_simple_multi_edge
   '''
   graph TD;

libs/langgraph/tests/__snapshots__/test_pregel_async.ambr

@@ -1302,6 +1302,81 @@
    +---------+   
   '''
 # ---
+# name: test_send_react_interrupt_control[memory]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres_aio]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres_aio_pipe]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[postgres_aio_pool]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
+# name: test_send_react_interrupt_control[sqlite_aio]
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	agent(agent)
+  	foo([foo]):::last
+  	__start__ --> agent;
+  	agent -.-> foo;
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
 # name: test_weather_subgraph[duckdb_aio]
   '''
   %%{init: {'flowchart': {'curve': 'linear'}}}%%

libs/langgraph/tests/conftest.py

@@ -327,7 +327,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
 ALL_CHECKPOINTERS_SYNC = [
     "memory",
     "sqlite",
-    "duckdb",
     "postgres",
     "postgres_pipe",
     "postgres_pool",
@@ -335,7 +334,6 @@ ALL_CHECKPOINTERS_SYNC = [
 ALL_CHECKPOINTERS_ASYNC = [
     "memory",
     "sqlite_aio",
-    "duckdb_aio",
     "postgres_aio",
     "postgres_aio_pipe",
     "postgres_aio_pool",

libs/langgraph/tests/test_algo.py

@@ -11,13 +11,21 @@ def test_prepare_next_tasks() -> None:
     with ChannelsManager({}, checkpoint, config) as (channels, managed):
         assert (
             prepare_next_tasks(
-                checkpoint, processes, channels, managed, config, 0, for_execution=False
+                checkpoint,
+                {},
+                processes,
+                channels,
+                managed,
+                config,
+                0,
+                for_execution=False,
             )
             == {}
         )
         assert (
             prepare_next_tasks(
                 checkpoint,
+                {},
                 processes,
                 channels,
                 managed,

libs/scheduler-kafka/Makefile

@@ -19,7 +19,7 @@ test:
 	exit $$EXIT_CODE
 
 test_watch:
-	make start-services && poetry run ptw . -- $(TEST_PATH); \
+	make start-services && poetry run ptw . -- -x $(TEST_PATH); \
 	EXIT_CODE=$$?; \
 	make stop-services; \
 	exit $$EXIT_CODE

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -38,7 +38,7 @@ from langgraph.scheduler.kafka.types import (
     Sendable,
     Topics,
 )
-from langgraph.types import LoopProtocol, RetryPolicy
+from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
 
@@ -198,6 +198,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 msg["task"]["path"],
                 msg["task"]["id"],
                 checkpoint=saved.checkpoint,
+                pending_writes=saved.pending_writes or [],
                 processes=graph.nodes,
                 channels=channels,
                 managed=managed,
@@ -211,6 +212,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 runner = PregelRunner(
                     submit=submit,
                     put_writes=partial(self._put_writes, submit, msg["config"]),
+                    schedule_task=self._schedule_task,
                 )
                 async for _ in runner.atick([task], reraise=False):
                     pass
@@ -239,6 +241,14 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
         )
         await fut
 
+    def _schedule_task(
+        self,
+        task: PregelExecutableTask,
+        idx: int,
+    ) -> None:
+        # will be scheduled by orchestrator when executor finishes
+        pass
+
     def _put_writes(
         self,
         submit: Submit,
@@ -400,6 +410,7 @@ class KafkaExecutor(AbstractContextManager):
                 msg["task"]["path"],
                 msg["task"]["id"],
                 checkpoint=saved.checkpoint,
+                pending_writes=saved.pending_writes or [],
                 processes=graph.nodes,
                 channels=channels,
                 managed=managed,
@@ -412,6 +423,7 @@ class KafkaExecutor(AbstractContextManager):
                 runner = PregelRunner(
                     submit=submit,
                     put_writes=partial(self._put_writes, submit, msg["config"]),
+                    schedule_task=self._schedule_task,
                 )
                 for _ in runner.tick([task], reraise=False):
                     pass
@@ -440,6 +452,14 @@ class KafkaExecutor(AbstractContextManager):
         )
         fut.result()
 
+    def _schedule_task(
+        self,
+        task: PregelExecutableTask,
+        idx: int,
+    ) -> None:
+        # will be scheduled by orchestrator when executor finishes
+        pass
+
     def _put_writes(
         self,
         submit: Submit,

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -161,18 +161,18 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
             specs=graph.channels,
             output_keys=graph.output_channels,
             stream_keys=graph.stream_channels,
+            interrupt_after=graph.interrupt_after_nodes,
+            interrupt_before=graph.interrupt_before_nodes,
             check_subgraphs=False,
         ) as loop:
-            if loop.tick(
-                input_keys=graph.input_channels,
-                interrupt_after=graph.interrupt_after_nodes,
-                interrupt_before=graph.interrupt_before_nodes,
-            ):
+            if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved
                 if hasattr(loop, "_put_checkpoint_fut"):
                     await loop._put_checkpoint_fut
                 # schedule any new tasks
-                if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
+                if new_tasks := [
+                    t for t in loop.tasks.values() if not t.scheduled and not t.writes
+                ]:
                     # send messages to executor
                     futures = await asyncio.gather(
                         *(
@@ -351,18 +351,18 @@ class KafkaOrchestrator(AbstractContextManager):
             specs=graph.channels,
             output_keys=graph.output_channels,
             stream_keys=graph.stream_channels,
+            interrupt_after=graph.interrupt_after_nodes,
+            interrupt_before=graph.interrupt_before_nodes,
             check_subgraphs=False,
         ) as loop:
-            if loop.tick(
-                input_keys=graph.input_channels,
-                interrupt_after=graph.interrupt_after_nodes,
-                interrupt_before=graph.interrupt_before_nodes,
-            ):
+            if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved
                 if hasattr(loop, "_put_checkpoint_fut"):
                     loop._put_checkpoint_fut.result()
                 # schedule any new tasks
-                if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]:
+                if new_tasks := [
+                    t for t in loop.tasks.values() if not t.scheduled and not t.writes
+                ]:
                     # send messages to executor
                     futures = [
                         self.producer.send(

libs/scheduler-kafka/langgraph/scheduler/kafka/types.py

@@ -24,8 +24,8 @@ class MessageToOrchestrator(TypedDict):
 
 
 class ExecutorTask(TypedDict):
-    id: str
-    path: tuple[str, ...]
+    id: Optional[str]
+    path: tuple[Union[str, int], ...]
 
 
 class MessageToExecutor(TypedDict):

libs/scheduler-kafka/tests/test_push.py

@@ -0,0 +1,208 @@
+import operator
+from typing import (
+    Annotated,
+    Literal,
+    Union,
+)
+
+import pytest
+from aiokafka import AIOKafkaProducer
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import FF_SEND_V2, START
+from langgraph.errors import NodeInterrupt
+from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
+from langgraph.scheduler.kafka import serde
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from langgraph.types import Send
+from tests.any import AnyDict
+from tests.drain import drain_topics_async
+
+pytestmark = pytest.mark.anyio
+
+
+def mk_push_graph(
+    checkpointer: BaseCheckpointSaver,
+) -> CompiledStateGraph:
+    # copied from test_send_dedupe_on_resume
+
+    class InterruptOnce:
+        ticks: int = 0
+
+        def __call__(self, state):
+            self.ticks += 1
+            if self.ticks == 1:
+                raise NodeInterrupt("Bahh")
+            return ["|".join(("flaky", str(state)))]
+
+    class Node:
+        def __init__(self, name: str):
+            self.name = name
+            self.ticks = 0
+            self.__name__ = name
+
+        def __call__(self, state):
+            self.ticks += 1
+            update = (
+                [self.name]
+                if isinstance(state, list)
+                else ["|".join((self.name, str(state)))]
+            )
+            if isinstance(state, GraphCommand):
+                return state.copy(update=update)
+            else:
+                return update
+
+    def send_for_fun(state):
+        return [
+            Send("2", GraphCommand(send=Send("2", 3))),
+            Send("2", GraphCommand(send=Send("flaky", 4))),
+            "3.1",
+        ]
+
+    def route_to_three(state) -> Literal["3"]:
+        return "3"
+
+    builder = StateGraph(Annotated[list, operator.add])
+    builder.add_node(Node("1"))
+    builder.add_node(Node("2"))
+    builder.add_node(Node("3"))
+    builder.add_node(Node("3.1"))
+    builder.add_node("flaky", InterruptOnce())
+    builder.add_edge(START, "1")
+    builder.add_conditional_edges("1", send_for_fun)
+    builder.add_conditional_edges("2", route_to_three)
+
+    return builder.compile(checkpointer=checkpointer)
+
+
+async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
+    if not FF_SEND_V2:
+        pytest.skip("Test requires FF_SEND_V2")
+
+    input = ["0"]
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_push_graph(acheckpointer)
+    graph_compare = mk_push_graph(acheckpointer)
+
+    # start a new run
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=input, config=config),
+        )
+
+    # drain topics
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
+
+    # check state
+    state = await graph.aget_state(config)
+    assert all(not t.error for t in state.tasks)
+    assert state.next == ("flaky",)
+    assert (
+        state.values
+        == await graph_compare.ainvoke(input, {"configurable": {"thread_id": "2"}})
+        == [
+            "0",
+            "1",
+            "2|Control(send=Send(node='2', arg=3))",
+            "2|Control(send=Send(node='flaky', arg=4))",
+            "2|3",
+        ]
+    )
+
+    # check history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 2
+
+    # check messages
+    assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "input": None,
+            "finally_send": None,
+        }
+        for c in reversed(history)
+        for _ in c.tasks
+    ]
+    assert exec_msgs == [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "task": {
+                "id": t.id,
+                "path": _convert_path(t.path),
+            },
+            "finally_send": None,
+        }
+        for c in reversed(history)
+        for t in c.tasks
+    ]
+
+    # resume the thread
+    async with AIOKafkaProducer(value_serializer=serde.dumps) as producer:
+        await producer.send_and_wait(
+            topics.orchestrator,
+            MessageToOrchestrator(input=None, config=config),
+        )
+
+    orch_msgs, exec_msgs = await drain_topics_async(topics, graph)
+
+    # check final state
+    state = await graph.aget_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == await graph_compare.ainvoke(None, {"configurable": {"thread_id": "2"}})
+        == [
+            "0",
+            "1",
+            "2|Control(send=Send(node='2', arg=3))",
+            "2|Control(send=Send(node='flaky', arg=4))",
+            "2|3",
+            "flaky|4",
+            "3",
+            "3.1",
+        ]
+    )
+
+    # check history
+    history = [c async for c in graph.aget_state_history(config)]
+    assert len(history) == 4
+
+    # check executions
+    # node "2" doesn't get called again, as we recover writes saved before
+    assert graph.builder.nodes["2"].runnable.func.ticks == 3
+    # node "flaky" gets called again, as it was interrupted
+    assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
+
+
+def _convert_path(
+    path: tuple[Union[str, int, tuple], ...],
+) -> list[Union[str, int, list]]:
+    return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)

libs/scheduler-kafka/tests/test_push_sync.py

@@ -0,0 +1,210 @@
+import operator
+from typing import (
+    Annotated,
+    Literal,
+    Union,
+)
+
+import pytest
+
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import FF_SEND_V2, START
+from langgraph.errors import NodeInterrupt
+from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph
+from langgraph.scheduler.kafka import serde
+from langgraph.scheduler.kafka.default_sync import DefaultProducer
+from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
+from langgraph.types import Send
+from tests.any import AnyDict
+from tests.drain import drain_topics
+
+pytestmark = pytest.mark.anyio
+
+
+def mk_push_graph(
+    checkpointer: BaseCheckpointSaver,
+) -> CompiledStateGraph:
+    # copied from test_send_dedupe_on_resume
+
+    class InterruptOnce:
+        ticks: int = 0
+
+        def __call__(self, state):
+            self.ticks += 1
+            if self.ticks == 1:
+                raise NodeInterrupt("Bahh")
+            return ["|".join(("flaky", str(state)))]
+
+    class Node:
+        def __init__(self, name: str):
+            self.name = name
+            self.ticks = 0
+            self.__name__ = name
+
+        def __call__(self, state):
+            self.ticks += 1
+            update = (
+                [self.name]
+                if isinstance(state, list)
+                else ["|".join((self.name, str(state)))]
+            )
+            if isinstance(state, GraphCommand):
+                return state.copy(update=update)
+            else:
+                return update
+
+    def send_for_fun(state):
+        return [
+            Send("2", GraphCommand(send=Send("2", 3))),
+            Send("2", GraphCommand(send=Send("flaky", 4))),
+            "3.1",
+        ]
+
+    def route_to_three(state) -> Literal["3"]:
+        return "3"
+
+    builder = StateGraph(Annotated[list, operator.add])
+    builder.add_node(Node("1"))
+    builder.add_node(Node("2"))
+    builder.add_node(Node("3"))
+    builder.add_node(Node("3.1"))
+    builder.add_node("flaky", InterruptOnce())
+    builder.add_edge(START, "1")
+    builder.add_conditional_edges("1", send_for_fun)
+    builder.add_conditional_edges("2", route_to_three)
+
+    return builder.compile(checkpointer=checkpointer)
+
+
+def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
+    if not FF_SEND_V2:
+        pytest.skip("Test requires FF_SEND_V2")
+
+    input = ["0"]
+    config = {"configurable": {"thread_id": "1"}}
+    graph = mk_push_graph(acheckpointer)
+    graph_compare = mk_push_graph(acheckpointer)
+
+    # start a new run
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=input, config=config)),
+        )
+        producer.flush()
+
+    # drain topics
+    orch_msgs, exec_msgs = drain_topics(topics, graph)
+
+    # check state
+    state = graph.get_state(config)
+    assert all(not t.error for t in state.tasks)
+    assert state.next == ("flaky",)
+    assert (
+        state.values
+        == graph_compare.invoke(input, {"configurable": {"thread_id": "2"}})
+        == [
+            "0",
+            "1",
+            "2|Control(send=Send(node='2', arg=3))",
+            "2|Control(send=Send(node='flaky', arg=4))",
+            "2|3",
+        ]
+    )
+
+    # check history
+    history = [c for c in graph.get_state_history(config)]
+    assert len(history) == 2
+
+    # check messages
+    assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "input": None,
+            "finally_send": None,
+        }
+        for c in reversed(history)
+        for _ in c.tasks
+    ]
+    assert exec_msgs == [
+        {
+            "config": {
+                "callbacks": None,
+                "configurable": {
+                    "__pregel_ensure_latest": True,
+                    "__pregel_dedupe_tasks": True,
+                    "__pregel_resuming": False,
+                    "checkpoint_id": c.config["configurable"]["checkpoint_id"],
+                    "checkpoint_ns": "",
+                    "thread_id": "1",
+                },
+                "metadata": AnyDict(),
+                "recursion_limit": 25,
+                "tags": [],
+            },
+            "task": {
+                "id": t.id,
+                "path": _convert_path(t.path),
+            },
+            "finally_send": None,
+        }
+        for c in reversed(history)
+        for t in c.tasks
+    ]
+
+    # resume the thread
+    with DefaultProducer() as producer:
+        producer.send(
+            topics.orchestrator,
+            value=serde.dumps(MessageToOrchestrator(input=None, config=config)),
+        )
+        producer.flush()
+
+    orch_msgs, exec_msgs = drain_topics(topics, graph)
+
+    # check final state
+    state = graph.get_state(config)
+    assert state.next == ()
+    assert (
+        state.values
+        == graph_compare.invoke(None, {"configurable": {"thread_id": "2"}})
+        == [
+            "0",
+            "1",
+            "2|Control(send=Send(node='2', arg=3))",
+            "2|Control(send=Send(node='flaky', arg=4))",
+            "2|3",
+            "flaky|4",
+            "3",
+            "3.1",
+        ]
+    )
+
+    # check history
+    history = [c for c in graph.get_state_history(config)]
+    assert len(history) == 4
+
+    # check executions
+    # node "2" doesn't get called again, as we recover writes saved before
+    assert graph.builder.nodes["2"].runnable.func.ticks == 3
+    # node "flaky" gets called again, as it was interrupted
+    assert graph.builder.nodes["flaky"].runnable.func.ticks == 2
+
+
+def _convert_path(
+    path: tuple[Union[str, int, tuple], ...],
+) -> list[Union[str, int, list]]:
+    return list(_convert_path(p) if isinstance(p, tuple) else p for p in path)

libs/scheduler-kafka/tests/test_subgraph.py

@@ -194,8 +194,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -258,8 +259,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -352,8 +354,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -456,8 +459,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -515,8 +519,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -630,8 +635,9 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]

libs/scheduler-kafka/tests/test_subgraph_sync.py

@@ -193,8 +193,9 @@ def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -255,10 +256,11 @@ def test_subgraph_w_interrupt(
                             "__pregel_read": None,
                             "__pregel_send": None,
                             "__pregel_ensure_latest": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -350,9 +352,10 @@ def test_subgraph_w_interrupt(
                             "__pregel_send": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_resuming": False,
                             "__pregel_task_id": history[0].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -453,9 +456,10 @@ def test_subgraph_w_interrupt(
                             "__pregel_send": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_resuming": True,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -512,9 +516,10 @@ def test_subgraph_w_interrupt(
                             "__pregel_send": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_resuming": True,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -628,8 +633,9 @@ def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
-                            '__pregel_store': None,
+                            "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
+                            "__pregel_resume_value": None,
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]