Deferred Node (#4269)

libs/langgraph/langgraph/channels/base.py

@@ -53,16 +53,7 @@ class BaseChannel(Generic[Value, Update, C], ABC):
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
 
-    # state methods
-
-    @abstractmethod
-    def update(self, values: Sequence[Update]) -> bool:
-        """Update the channel's value with the given sequence of updates.
-        The order of the updates in the sequence is arbitrary.
-        This method is called by Pregel for all channels at the end of each step.
-        If there are no updates, it is called with an empty sequence.
-        Raises InvalidUpdateError if the sequence of updates is invalid.
-        Returns True if the channel was updated, False otherwise."""
+    # read methods
 
     @abstractmethod
     def get(self) -> Value:
@@ -70,13 +61,6 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
         Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
-    def consume(self) -> bool:
-        """Mark the current value of the channel as consumed. By default, no-op.
-        This is called by Pregel before the start of the next step, for all
-        channels that triggered a node. If the channel was updated, return True.
-        """
-        return False
-
     def is_available(self) -> bool:
         """Return True if the channel is available (not empty), False otherwise.
         Subclasses should override this method to provide a more efficient
@@ -88,6 +72,34 @@ class BaseChannel(Generic[Value, Update, C], ABC):
         except EmptyChannelError:
             return False
 
+    # write methods
+
+    @abstractmethod
+    def update(self, values: Sequence[Update]) -> bool:
+        """Update the channel's value with the given sequence of updates.
+        The order of the updates in the sequence is arbitrary.
+        This method is called by Pregel for all channels at the end of each step.
+        If there are no updates, it is called with an empty sequence.
+        Raises InvalidUpdateError if the sequence of updates is invalid.
+        Returns True if the channel was updated, False otherwise."""
+
+    def consume(self) -> bool:
+        """Notify the channel that a subscribed task ran. By default, no-op.
+        A channel can use this method to modify its state, preventing the value
+        from being consumed again.
+
+        Returns True if the channel was updated, False otherwise.
+        """
+        return False
+
+    def finish(self) -> bool:
+        """Notify the channel that the Pregel run is finishing. By default, no-op.
+        A channel can use this method to modify its state, preventing finish.
+
+        Returns True if the channel was updated, False otherwise.
+        """
+        return False
+
 
 __all__ = [
     "BaseChannel",

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -81,12 +81,9 @@ class DynamicBarrierValue(
             updated = False
             for value in values:
                 assert not isinstance(value, WaitForNames)
-                if value in self.names:
-                    if value not in self.seen:
-                        self.seen.add(value)
-                        updated = True
-                else:
-                    raise InvalidUpdateError(f"Value {value} not in {self.names}")
+                if value in self.names and value not in self.seen:
+                    self.seen.add(value)
+                    updated = True
             return updated
 
     def get(self) -> Value:
@@ -103,3 +100,107 @@ class DynamicBarrierValue(
             self.names = None
             return True
         return False
+
+
+class DynamicBarrierValueAfterFinish(
+    Generic[Value], BaseChannel[Value, Union[Value, WaitForNames], Set[Value]]
+):
+    """A channel that switches between two states
+
+    - in the "priming" state it can't be read from.
+        - if it receives a WaitForNames update, it switches to the "waiting" state.
+    - in the "waiting" state it collects named values until all are received.
+        - once all named values are received, and the finished flag is set, it can be read once, and it switches
+          back to the "priming" state.
+    """
+
+    __slots__ = ("names", "seen", "finished")
+
+    names: Optional[Set[Value]]
+    seen: set[Value]
+    finished: bool
+
+    def __init__(self, typ: type[Value]) -> None:
+        super().__init__(typ)
+        self.names = None
+        self.seen = set()
+        self.finished = False
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, DynamicBarrierValueAfterFinish)
+            and value.names == self.names
+        )
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ)
+        empty.key = self.key
+        empty.names = self.names
+        empty.seen = self.seen.copy()
+        empty.finished = self.finished
+        return empty
+
+    def checkpoint(self) -> tuple[Optional[Set[Value]], set[Value], bool]:
+        return (self.names, self.seen, self.finished)
+
+    def from_checkpoint(
+        self, checkpoint: tuple[Optional[Set[Value]], set[Value], bool]
+    ) -> Self:
+        empty = self.__class__(self.typ)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            names, seen, finished = checkpoint
+            empty.names = names if names is not None else None
+            empty.seen = seen
+            empty.finished = finished
+        return empty
+
+    def update(self, values: Sequence[Union[Value, WaitForNames]]) -> bool:
+        if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
+            if len(wait_for_names) > 1:
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Received multiple WaitForNames updates in the same step."
+                )
+            self.names = wait_for_names[0].names
+            return True
+        elif self.names is not None:
+            updated = False
+            for value in values:
+                assert not isinstance(value, WaitForNames)
+                if value in self.names and value not in self.seen:
+                    self.seen.add(value)
+                    updated = True
+            return updated
+
+    def get(self) -> Value:
+        if not self.finished and self.seen != self.names:
+            raise EmptyChannelError()
+        return None
+
+    def is_available(self) -> bool:
+        return self.seen == self.names and self.finished
+
+    def consume(self) -> bool:
+        if self.finished and self.seen == self.names:
+            self.seen = set()
+            self.names = None
+            return True
+        return False
+
+    def finish(self) -> bool:
+        if not self.finished and self.seen == self.names:
+            self.finished = True
+            return True
+        else:
+            return False

libs/langgraph/langgraph/channels/last_value.py

@@ -70,3 +70,73 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     def checkpoint(self) -> Value:
         return self.value
+
+
+class LastValueAfterFinish(
+    Generic[Value], BaseChannel[Value, Value, tuple[Value, bool]]
+):
+    """Stores the last value received, but only made available after finish().
+    Once made available, clears the value."""
+
+    __slots__ = ("value", "finished")
+
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+        self.finished = False
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, LastValueAfterFinish)
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def checkpoint(self) -> tuple[Value, bool]:
+        if self.value is MISSING:
+            return MISSING
+        return (self.value, self.finished)
+
+    def from_checkpoint(self, checkpoint: tuple[Value, bool]) -> Self:
+        empty = self.__class__(self.typ)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.value, empty.finished = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        if len(values) == 0:
+            return False
+
+        self.finished = False
+        self.value = values[-1]
+        return True
+
+    def consume(self) -> bool:
+        if self.finished:
+            self.finished = False
+            self.value = MISSING
+            return True
+
+        return False
+
+    def finish(self) -> bool:
+        if not self.finished and self.value is not MISSING:
+            self.finished = True
+            return True
+        else:
+            return False
+
+    def get(self) -> Value:
+        if self.value is MISSING or not self.finished:
+            raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING and self.finished

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -77,3 +77,89 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
             self.seen = set()
             return True
         return False
+
+
+class NamedBarrierValueAfterFinish(
+    Generic[Value], BaseChannel[Value, Value, set[Value]]
+):
+    """A channel that waits until all named values are received before making the value ready to be made available. It is only made available after finish() is called."""
+
+    __slots__ = ("names", "seen", "finished")
+
+    names: set[Value]
+    seen: set[Value]
+
+    def __init__(self, typ: type[Value], names: set[Value]) -> None:
+        super().__init__(typ)
+        self.names = names
+        self.seen: set[str] = set()
+        self.finished = False
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, NamedBarrierValueAfterFinish)
+            and value.names == self.names
+        )
+
+    @property
+    def ValueType(self) -> type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        empty.seen = self.seen.copy()
+        empty.finished = self.finished
+        return empty
+
+    def checkpoint(self) -> tuple[set[Value], bool]:
+        return (self.seen, self.finished)
+
+    def from_checkpoint(self, checkpoint: tuple[set[Value], bool]) -> Self:
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        if checkpoint is not MISSING:
+            empty.seen, empty.finished = checkpoint
+        return empty
+
+    def update(self, values: Sequence[Value]) -> bool:
+        updated = False
+        for value in values:
+            if value in self.names:
+                if value not in self.seen:
+                    self.seen.add(value)
+                    updated = True
+            else:
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
+        return updated
+
+    def get(self) -> Value:
+        if not self.finished or self.seen != self.names:
+            raise EmptyChannelError()
+        return None
+
+    def is_available(self) -> bool:
+        return self.finished and self.seen == self.names
+
+    def consume(self) -> bool:
+        if self.finished and self.seen == self.names:
+            self.finished = False
+            self.seen = set()
+            return True
+        return False
+
+    def finish(self) -> bool:
+        if not self.finished and self.seen == self.names:
+            self.finished = True
+            return True
+        else:
+            return False

libs/langgraph/langgraph/graph/state.py

@@ -28,10 +28,17 @@ from typing_extensions import Self
 from langgraph._api.deprecation import LangGraphDeprecationWarning
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
-from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitForNames
+from langgraph.channels.dynamic_barrier_value import (
+    DynamicBarrierValue,
+    DynamicBarrierValueAfterFinish,
+    WaitForNames,
+)
 from langgraph.channels.ephemeral_value import EphemeralValue
-from langgraph.channels.last_value import LastValue
-from langgraph.channels.named_barrier_value import NamedBarrierValue
+from langgraph.channels.last_value import LastValue, LastValueAfterFinish
+from langgraph.channels.named_barrier_value import (
+    NamedBarrierValue,
+    NamedBarrierValueAfterFinish,
+)
 from langgraph.checkpoint.base import Checkpoint
 from langgraph.constants import (
     EMPTY_SEQ,
@@ -107,6 +114,7 @@ class StateNodeSpec(NamedTuple):
     input: type[Any]
     retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]]
     ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ
+    defer: bool = False
 
 
 class StateGraph(Graph):
@@ -247,6 +255,7 @@ class StateGraph(Graph):
         self,
         node: RunnableLike,
         *,
+        defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
@@ -263,6 +272,7 @@ class StateGraph(Graph):
         node: str,
         action: RunnableLike,
         *,
+        defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
@@ -276,6 +286,7 @@ class StateGraph(Graph):
         node: Union[str, RunnableLike],
         action: Optional[RunnableLike] = None,
         *,
+        defer: bool = False,
         metadata: Optional[dict[str, Any]] = None,
         input: Optional[type[Any]] = None,
         retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
@@ -421,6 +432,7 @@ class StateGraph(Graph):
             input=input or self.schema,
             retry_policy=retry,
             ends=ends,
+            defer=defer,
         )
         return self
 
@@ -784,7 +796,11 @@ class CompiledStateGraph(CompiledGraph):
                 self.schema_to_mapper[input_schema] = mapper
 
             branch_channel = CHANNEL_BRANCH_TO.format(key)
-            self.channels[branch_channel] = EphemeralValue(Any, guard=False)
+            self.channels[branch_channel] = (
+                LastValueAfterFinish(Any)
+                if node.defer
+                else EphemeralValue(Any, guard=False)
+            )
             self.nodes[key] = PregelNode(
                 triggers=[branch_channel],
                 # read state keys and managed values
@@ -812,7 +828,12 @@ class CompiledStateGraph(CompiledGraph):
         elif end != END:
             channel_name = f"join:{'+'.join(starts)}:{end}"
             # register channel
-            self.channels[channel_name] = NamedBarrierValue(str, set(starts))
+            if self.builder.nodes[end].defer:
+                self.channels[channel_name] = NamedBarrierValueAfterFinish(
+                    str, set(starts)
+                )
+            else:
+                self.channels[channel_name] = NamedBarrierValue(str, set(starts))
             # subscribe to channel
             self.nodes[end].triggers.append(channel_name)
             # publish to channel
@@ -889,7 +910,10 @@ class CompiledStateGraph(CompiledGraph):
                 else [node for node in self.builder.nodes if node != branch.then]
             )
             channel_name = f"branch:{start}:{name}::then"
-            self.channels[channel_name] = DynamicBarrierValue(str)
+            if self.builder.nodes[branch.then].defer:
+                self.channels[channel_name] = DynamicBarrierValueAfterFinish(str)
+            else:
+                self.channels[channel_name] = DynamicBarrierValue(str)
             self.nodes[branch.then].triggers.append(channel_name)
             for end in ends:
                 if end != END:

libs/langgraph/langgraph/pregel/__init__.py

@@ -506,7 +506,7 @@ class Pregel(PregelProtocol):
 
     name: str = "LangGraph"
 
-    trigger_to_nodes: Mapping[str, Sequence[str]] | None = None
+    trigger_to_nodes: Mapping[str, Sequence[str]]
 
     def __init__(
         self,
@@ -552,7 +552,7 @@ class Pregel(PregelProtocol):
         self.config_type = config_type
         self.input_model = input_model
         self.config = config
-        self.trigger_to_nodes = trigger_to_nodes
+        self.trigger_to_nodes = trigger_to_nodes or {}
         self.name = name
         if auto_validate:
             self.validate()
@@ -949,6 +949,7 @@ class Pregel(PregelProtocol):
                     channels,
                     [PregelTaskWrites((), INPUT, null_writes, [])],
                     None,
+                    self.trigger_to_nodes,
                 )
             if apply_pending_writes and saved.pending_writes:
                 for tid, k, v in saved.pending_writes:
@@ -958,7 +959,9 @@ class Pregel(PregelProtocol):
                         continue
                     next_tasks[tid].writes.append((k, v))
                 if tasks := [t for t in next_tasks.values() if t.writes]:
-                    apply_writes(saved.checkpoint, channels, tasks, None)
+                    apply_writes(
+                        saved.checkpoint, channels, tasks, None, self.trigger_to_nodes
+                    )
             tasks_with_writes = tasks_w_writes(
                 next_tasks.values(),
                 saved.pending_writes,
@@ -1071,6 +1074,7 @@ class Pregel(PregelProtocol):
                     channels,
                     [PregelTaskWrites((), INPUT, null_writes, [])],
                     None,
+                    self.trigger_to_nodes,
                 )
             if apply_pending_writes and saved.pending_writes:
                 for tid, k, v in saved.pending_writes:
@@ -1080,7 +1084,9 @@ class Pregel(PregelProtocol):
                         continue
                     next_tasks[tid].writes.append((k, v))
                 if tasks := [t for t in next_tasks.values() if t.writes]:
-                    apply_writes(saved.checkpoint, channels, tasks, None)
+                    apply_writes(
+                        saved.checkpoint, channels, tasks, None, self.trigger_to_nodes
+                    )
 
             tasks_with_writes = tasks_w_writes(
                 next_tasks.values(),
@@ -1407,6 +1413,7 @@ class Pregel(PregelProtocol):
                                 channels,
                                 [PregelTaskWrites((), INPUT, null_writes, [])],
                                 None,
+                                self.trigger_to_nodes,
                             )
                         # apply writes from tasks that already ran
                         for tid, k, v in saved.pending_writes or []:
@@ -1416,7 +1423,13 @@ class Pregel(PregelProtocol):
                                 continue
                             next_tasks[tid].writes.append((k, v))
                         # clear all current tasks
-                        apply_writes(checkpoint, channels, next_tasks.values(), None)
+                        apply_writes(
+                            checkpoint,
+                            channels,
+                            next_tasks.values(),
+                            None,
+                            self.trigger_to_nodes,
+                        )
                     # save checkpoint
                     next_config = checkpointer.put(
                         checkpoint_config,
@@ -1475,6 +1488,7 @@ class Pregel(PregelProtocol):
                             channels,
                             [PregelTaskWrites((), INPUT, input_writes, [])],
                             checkpointer.get_next_version,
+                            self.trigger_to_nodes,
                         )
 
                         # apply input write to channels
@@ -1575,6 +1589,7 @@ class Pregel(PregelProtocol):
                             channels,
                             [PregelTaskWrites((), INPUT, null_writes, [])],
                             None,
+                            self.trigger_to_nodes,
                         )
                     # apply writes
                     for tid, k, v in saved.pending_writes:
@@ -1584,7 +1599,9 @@ class Pregel(PregelProtocol):
                             continue
                         next_tasks[tid].writes.append((k, v))
                     if tasks := [t for t in next_tasks.values() if t.writes]:
-                        apply_writes(checkpoint, channels, tasks, None)
+                        apply_writes(
+                            checkpoint, channels, tasks, None, self.trigger_to_nodes
+                        )
             valid_updates: list[tuple[str, dict[str, Any] | None]] = []
             if len(updates) == 1:
                 values, as_node = updates[0]
@@ -1672,7 +1689,11 @@ class Pregel(PregelProtocol):
                     checkpointer.put_writes(checkpoint_config, channel_writes, task_id)
             # apply to checkpoint and save
             mv_writes, _ = apply_writes(
-                checkpoint, channels, run_tasks, checkpointer.get_next_version
+                checkpoint,
+                channels,
+                run_tasks,
+                checkpointer.get_next_version,
+                self.trigger_to_nodes,
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
@@ -1822,6 +1843,7 @@ class Pregel(PregelProtocol):
                                 channels,
                                 [PregelTaskWrites((), INPUT, null_writes, [])],
                                 None,
+                                self.trigger_to_nodes,
                             )
                         # apply writes from tasks that already ran
                         for tid, k, v in saved.pending_writes or []:
@@ -1831,7 +1853,13 @@ class Pregel(PregelProtocol):
                                 continue
                             next_tasks[tid].writes.append((k, v))
                         # clear all current tasks
-                        apply_writes(checkpoint, channels, next_tasks.values(), None)
+                        apply_writes(
+                            checkpoint,
+                            channels,
+                            next_tasks.values(),
+                            None,
+                            self.trigger_to_nodes,
+                        )
                     # save checkpoint
                     next_config = await checkpointer.aput(
                         checkpoint_config,
@@ -1890,6 +1918,7 @@ class Pregel(PregelProtocol):
                             channels,
                             [PregelTaskWrites((), INPUT, input_writes, [])],
                             checkpointer.get_next_version,
+                            self.trigger_to_nodes,
                         )
 
                         # apply input write to channels
@@ -1990,6 +2019,7 @@ class Pregel(PregelProtocol):
                             channels,
                             [PregelTaskWrites((), INPUT, null_writes, [])],
                             None,
+                            self.trigger_to_nodes,
                         )
                     for tid, k, v in saved.pending_writes:
                         if k in (ERROR, INTERRUPT, SCHEDULED):
@@ -1998,7 +2028,9 @@ class Pregel(PregelProtocol):
                             continue
                         next_tasks[tid].writes.append((k, v))
                     if tasks := [t for t in next_tasks.values() if t.writes]:
-                        apply_writes(checkpoint, channels, tasks, None)
+                        apply_writes(
+                            checkpoint, channels, tasks, None, self.trigger_to_nodes
+                        )
             valid_updates: list[tuple[str, dict[str, Any] | None]] = []
             if len(updates) == 1:
                 values, as_node = updates[0]
@@ -2084,7 +2116,11 @@ class Pregel(PregelProtocol):
                     )
             # apply to checkpoint and save
             mv_writes, _ = apply_writes(
-                checkpoint, channels, run_tasks, checkpointer.get_next_version
+                checkpoint,
+                channels,
+                run_tasks,
+                checkpointer.get_next_version,
+                self.trigger_to_nodes,
             )
             assert not mv_writes, "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)

libs/langgraph/langgraph/pregel/algo.py

@@ -232,6 +232,7 @@ def apply_writes(
     channels: Mapping[str, BaseChannel],
     tasks: Iterable[WritesProtocol],
     get_next_version: Optional[GetNextVersion],
+    trigger_to_nodes: Mapping[str, Sequence[str]],
 ) -> tuple[dict[str, list[Any]], set[str]]:
     """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
@@ -317,7 +318,9 @@ def apply_writes(
                     max_version,
                     channels[chan],
                 )
-            updated_channels.add(chan)
+                # unavailable channels can't trigger tasks, so don't add them
+                if channels[chan].is_available():
+                    updated_channels.add(chan)
 
     # Channels that weren't updated in this step are notified of a new step
     if bump_step:
@@ -328,6 +331,26 @@ def apply_writes(
                         max_version,
                         channels[chan],
                     )
+                    # unavailable channels can't trigger tasks, so don't add them
+                    if channels[chan].is_available():
+                        updated_channels.add(chan)
+
+    # If this is (tentatively) the last superstep, notify all channels of finish
+    if (
+        bump_step
+        and not checkpoint["pending_sends"]
+        and updated_channels.isdisjoint(trigger_to_nodes)
+    ):
+        for chan in channels:
+            if channels[chan].finish() and get_next_version is not None:
+                checkpoint["channel_versions"][chan] = get_next_version(
+                    max_version,
+                    channels[chan],
+                )
+                # unavailable channels can't trigger tasks, so don't add them
+                if channels[chan].is_available():
+                    updated_channels.add(chan)
+
     # Return managed values writes to be applied externally
     return pending_writes_by_managed, updated_channels
 

libs/langgraph/langgraph/pregel/draw.py

@@ -31,7 +31,7 @@ def draw_graph(
     input_channels: Union[str, Sequence[str]],
     interrupt_after_nodes: Union[All, Sequence[str]],
     interrupt_before_nodes: Union[All, Sequence[str]],
-    trigger_to_nodes: Optional[Mapping[str, Sequence[str]]],
+    trigger_to_nodes: Mapping[str, Sequence[str]],
     checkpointer: Checkpointer,
     subgraphs: dict[str, Graph],
     limit: int = 250,
@@ -79,6 +79,7 @@ def draw_graph(
                 PregelTaskWrites((), INPUT, input_writes, []),
             ],
             get_next_version,
+            trigger_to_nodes,
         )
         # prepare first tasks
         tasks = prepare_next_tasks(
@@ -98,7 +99,7 @@ def draw_graph(
         )
         start_tasks = tasks
         # run the pregel loop
-        for _ in range(limit):
+        for step in range(step, limit):
             if not tasks:
                 break
             conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
@@ -144,7 +145,7 @@ def draw_graph(
                     trigger_to_sources[trigger].add((src, cond, label))
             # apply writes
             _, updated_channels = apply_writes(
-                checkpoint, channels, tasks.values(), get_next_version
+                checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes
             )
             # prepare next tasks
             tasks = prepare_next_tasks(
@@ -164,9 +165,19 @@ def draw_graph(
             )
             # collect edges
             for task in tasks.values():
+                added = False
                 for trigger in task.triggers:
                     for src, cond, label in sorted(trigger_to_sources[trigger]):
                         edges.add((src, task.name, cond, label))
+                        # if the edge is from this step, skip adding the implicit edges
+                        if (trigger, cond, label) in step_sources.get(src, set()):
+                            added = True
+                        else:
+                            sources[src].discard((trigger, cond, label))
+                # if no edges from this step, add implicit edges from all previous tasks
+                if not added:
+                    for src in step_sources:
+                        edges.add((src, task.name, True, None))
         # assemble the graph
         graph = Graph()
         # add nodes

libs/langgraph/langgraph/pregel/loop.py

@@ -211,13 +211,13 @@ class PregelLoop(LoopProtocol):
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]],
         stream_keys: Union[str, Sequence[str]],
+        trigger_to_nodes: Mapping[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,
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
-        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(
@@ -483,6 +483,7 @@ class PregelLoop(LoopProtocol):
                 self.channels,
                 self.tasks.values(),
                 self.checkpointer_get_next_version,
+                self.trigger_to_nodes,
             )
             # apply writes to managed values
             for key, values in mv_writes.items():
@@ -679,6 +680,7 @@ class PregelLoop(LoopProtocol):
                 self.channels,
                 [PregelTaskWrites((), INPUT, null_writes, [])],
                 self.checkpointer_get_next_version,
+                self.trigger_to_nodes,
             )
             for key, values in mv_writes.items():
                 self._update_mv(key, values)
@@ -731,6 +733,7 @@ class PregelLoop(LoopProtocol):
                     PregelTaskWrites((), INPUT, input_writes, []),
                 ],
                 self.checkpointer_get_next_version,
+                self.trigger_to_nodes,
             )
             assert not mv_writes, "Can't write to SharedValues in graph input"
             # save input checkpoint
@@ -868,6 +871,7 @@ class PregelLoop(LoopProtocol):
                     self.channels,
                     self.tasks.values(),
                     self.checkpointer_get_next_version,
+                    self.trigger_to_nodes,
                 )
                 for key, values in mv_writes.items():
                     self._update_mv(key, values)
@@ -966,6 +970,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        trigger_to_nodes: Mapping[str, Sequence[str]],
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
@@ -974,7 +979,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
-        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(
@@ -1116,6 +1120,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         checkpointer: Optional[BaseCheckpointSaver],
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        trigger_to_nodes: Mapping[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,
@@ -1124,7 +1129,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         input_model: Optional[type[BaseModel]] = None,
         debug: bool = False,
         migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None,
-        trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
         checkpoint_during: bool = True,
     ) -> None:
         super().__init__(

libs/langgraph/langgraph/types.py

@@ -313,8 +313,8 @@ class Command(Generic[N], ToolOutputMixin):
 
     graph: Optional[str] = None
     update: Optional[Any] = None
-    resume: Optional[Union[dict[str, Any], Any]] = None
-    goto: Union[Send, Sequence[Union[Send, str]], str] = ()
+    resume: Optional[Union[Any, dict[str, Any]]] = None
+    goto: Union[Send, Sequence[Union[Send, N]], N] = ()
 
     def __repr__(self) -> str:
         # get all non-None values

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -378,7 +378,7 @@
         "target": "agent"
       },
       {
-        "source": "human",
+        "source": "agent",
         "target": "__end__",
         "conditional": true
       }
@@ -392,7 +392,50 @@
   	__start__ --> human;
   	agent --> human;
   	human --> agent;
-  	human -.-> __end__;
+  	agent -.-> __end__;
+  
+  '''
+# ---
+# name: test_in_one_fan_out_state_graph_defer_node[memory-False]
+  '''
+  graph TD;
+  	__start__ --> rewrite_query;
+  	analyzer_one -.-> qa;
+  	retriever_one --> analyzer_one;
+  	retriever_one --> qa;
+  	retriever_two --> qa;
+  	rewrite_query --> retriever_one;
+  	rewrite_query --> retriever_two;
+  	qa --> __end__;
+  
+  '''
+# ---
+# name: test_in_one_fan_out_state_graph_defer_node[memory-True]
+  '''
+  graph TD;
+  	__start__ --> rewrite_query;
+  	analyzer_one -.-> qa;
+  	retriever_one --> analyzer_one;
+  	retriever_one --> qa;
+  	retriever_two --> qa;
+  	rewrite_query --> retriever_one;
+  	rewrite_query --> retriever_two;
+  	qa --> __end__;
+  
+  '''
+# ---
+# name: test_in_one_fan_out_state_graph_then_defer_node[memory-True]
+  '''
+  graph TD;
+  	__start__ --> rewrite_query;
+  	analyzer_one --> qa;
+  	analyzer_one --> retriever_one;
+  	retriever_one -.-> qa;
+  	retriever_two --> qa;
+  	rewrite_query -.-> analyzer_one;
+  	rewrite_query -.-> qa;
+  	rewrite_query -.-> retriever_two;
+  	qa --> __end__;
   
   '''
 # ---

libs/langgraph/tests/test_pregel.py

@@ -2500,6 +2500,563 @@ def test_in_one_fan_out_state_graph_waiting_edge(
     ]
 
 
+@pytest.mark.parametrize("use_waiting_edge", (True, False))
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_in_one_fan_out_state_graph_defer_node(
+    snapshot: SnapshotAssertion,
+    request: pytest.FixtureRequest,
+    checkpointer_name: str,
+    use_waiting_edge: bool,
+) -> None:
+    checkpointer: BaseCheckpointSaver = request.getfixturevalue(
+        f"checkpointer_{checkpointer_name}"
+    )
+
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    workflow = StateGraph(State)
+
+    @workflow.add_node
+    def rewrite_query(data: State) -> State:
+        return {"query": f"query: {data['query']}"}
+
+    def analyzer_one(data: State) -> State:
+        return {"query": f"analyzed: {data['query']}"}
+
+    def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    def retriever_two(data: State) -> State:
+        time.sleep(0.1)  # to ensure stream order
+        return {"docs": ["doc3", "doc4"]}
+
+    def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow.add_node(analyzer_one)
+    workflow.add_node(retriever_one)
+    workflow.add_node(retriever_two)
+    workflow.add_node(qa, defer=True)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_edge("rewrite_query", "retriever_one")
+    workflow.add_edge("retriever_one", "analyzer_one")
+    workflow.add_edge("rewrite_query", "retriever_two")
+    if use_waiting_edge:
+        workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    else:
+        workflow.add_edge("retriever_one", "qa")
+        workflow.add_edge("retriever_two", "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    if checkpointer_name == "memory":
+        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+    assert app.invoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [*app.stream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 1,
+            "payload": {
+                "id": AnyStr(),
+                "name": "rewrite_query",
+                "input": {"query": "what is weather in sf", "docs": []},
+                "triggers": ("branch:to:rewrite_query",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 1,
+            "payload": {
+                "id": AnyStr(),
+                "name": "rewrite_query",
+                "error": None,
+                "result": [("query", "query: what is weather in sf")],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_one",
+                "input": {"query": "query: what is weather in sf", "docs": []},
+                "triggers": ("branch:to:retriever_one",),
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_two",
+                "input": {"query": "query: what is weather in sf", "docs": []},
+                "triggers": ("branch:to:retriever_two",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_one",
+                "error": None,
+                "result": [("docs", ["doc1", "doc2"])],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_two",
+                "error": None,
+                "result": [("docs", ["doc3", "doc4"])],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 3,
+            "payload": {
+                "id": AnyStr(),
+                "name": "analyzer_one",
+                "input": {
+                    "query": "query: what is weather in sf",
+                    "docs": ["doc1", "doc2", "doc3", "doc4"],
+                },
+                "triggers": ("branch:to:analyzer_one",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 3,
+            "payload": {
+                "id": AnyStr(),
+                "name": "analyzer_one",
+                "error": None,
+                "result": [("query", "analyzed: query: what is weather in sf")],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 4,
+            "payload": {
+                "id": AnyStr(),
+                "name": "qa",
+                "input": {
+                    "query": "analyzed: query: what is weather in sf",
+                    "docs": ["doc1", "doc2", "doc3", "doc4"],
+                },
+                "triggers": ("branch:to:qa", "join:retriever_one+retriever_two:qa")
+                if use_waiting_edge
+                else ("branch:to:qa",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 4,
+            "payload": {
+                "id": AnyStr(),
+                "name": "qa",
+                "error": None,
+                "result": [("answer", "doc1,doc2,doc3,doc4")],
+                "interrupts": [],
+            },
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer,
+        interrupt_after=["analyzer_one"],
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"__interrupt__": ()},
+    ]
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer,
+        interrupt_before=["qa"],
+    )
+    config = {"configurable": {"thread_id": "2"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"__interrupt__": ()},
+    ]
+
+    app_w_interrupt.update_state(config, {"docs": ["doc5"]})
+    expected_parent_config = (
+        None
+        if "shallow" in checkpointer_name
+        else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
+    )
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "query": "analyzed: query: what is weather in sf",
+            "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
+        },
+        tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
+        next=("qa",),
+        config={
+            "configurable": {
+                "thread_id": "2",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        created_at=AnyStr(),
+        metadata={
+            "parents": {},
+            "source": "update",
+            "step": 4,
+            "writes": {"analyzer_one": {"docs": ["doc5"]}},
+            "thread_id": "2",
+        },
+        parent_config=expected_parent_config,
+        interrupts=(),
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
+    ]
+
+
+@pytest.mark.parametrize("with_path_map", (True, False))
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_in_one_fan_out_state_graph_then_defer_node(
+    snapshot: SnapshotAssertion,
+    request: pytest.FixtureRequest,
+    checkpointer_name: str,
+    with_path_map: bool,
+) -> None:
+    checkpointer: BaseCheckpointSaver = request.getfixturevalue(
+        f"checkpointer_{checkpointer_name}"
+    )
+
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    workflow = StateGraph(State)
+
+    @workflow.add_node
+    def rewrite_query(data: State) -> State:
+        return {"query": f"query: {data['query']}"}
+
+    def analyzer_one(data: State) -> State:
+        return {"query": f"analyzed: {data['query']}"}
+
+    def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    def retriever_two(data: State) -> State:
+        time.sleep(0.1)  # to ensure stream order
+        return {"docs": ["doc3", "doc4"]}
+
+    def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow.add_node(analyzer_one)
+    workflow.add_node(retriever_one)
+    workflow.add_node(retriever_two)
+    workflow.add_node(qa, defer=True)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_conditional_edges(
+        "rewrite_query",
+        lambda _: ["analyzer_one", "retriever_two"],
+        ["analyzer_one", "retriever_two"] if with_path_map else None,
+        then="qa",
+    )
+    workflow.add_edge("analyzer_one", "retriever_one")
+
+    app = workflow.compile()
+
+    if checkpointer_name == "memory" and with_path_map:
+        assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+    assert app.invoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [*app.stream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 1,
+            "payload": {
+                "id": AnyStr(),
+                "name": "rewrite_query",
+                "input": {"query": "what is weather in sf", "docs": []},
+                "triggers": ("branch:to:rewrite_query",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 1,
+            "payload": {
+                "id": AnyStr(),
+                "name": "rewrite_query",
+                "error": None,
+                "result": [("query", "query: what is weather in sf")],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "analyzer_one",
+                "input": {
+                    "query": "query: what is weather in sf",
+                    "docs": [],
+                },
+                "triggers": ("branch:to:analyzer_one",),
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_two",
+                "input": {"query": "query: what is weather in sf", "docs": []},
+                "triggers": ("branch:to:retriever_two",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "analyzer_one",
+                "error": None,
+                "result": [("query", "analyzed: query: what is weather in sf")],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 2,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_two",
+                "error": None,
+                "result": [("docs", ["doc3", "doc4"])],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 3,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_one",
+                "input": {
+                    "query": "analyzed: query: what is weather in sf",
+                    "docs": ["doc3", "doc4"],
+                },
+                "triggers": ("branch:to:retriever_one",),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 3,
+            "payload": {
+                "id": AnyStr(),
+                "name": "retriever_one",
+                "error": None,
+                "result": [("docs", ["doc1", "doc2"])],
+                "interrupts": [],
+            },
+        },
+        {
+            "type": "task",
+            "timestamp": AnyStr(),
+            "step": 4,
+            "payload": {
+                "id": AnyStr(),
+                "name": "qa",
+                "input": {
+                    "query": "analyzed: query: what is weather in sf",
+                    "docs": ["doc1", "doc2", "doc3", "doc4"],
+                },
+                "triggers": ("branch:rewrite_query:condition::then", "branch:to:qa"),
+            },
+        },
+        {
+            "type": "task_result",
+            "timestamp": AnyStr(),
+            "step": 4,
+            "payload": {
+                "id": AnyStr(),
+                "name": "qa",
+                "error": None,
+                "result": [("answer", "doc1,doc2,doc3,doc4")],
+                "interrupts": [],
+            },
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer,
+        interrupt_after=["analyzer_one"],
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"__interrupt__": ()},
+    ]
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer,
+        interrupt_before=["qa"],
+    )
+    config = {"configurable": {"thread_id": "2"}}
+
+    assert [
+        c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"__interrupt__": ()},
+    ]
+
+    app_w_interrupt.update_state(config, {"docs": ["doc5"]})
+    expected_parent_config = (
+        None
+        if "shallow" in checkpointer_name
+        else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
+    )
+    assert app_w_interrupt.get_state(config) == StateSnapshot(
+        values={
+            "query": "analyzed: query: what is weather in sf",
+            "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
+        },
+        tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
+        next=("qa",),
+        config={
+            "configurable": {
+                "thread_id": "2",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        created_at=AnyStr(),
+        metadata={
+            "parents": {},
+            "source": "update",
+            "step": 4,
+            "writes": {"retriever_one": {"docs": ["doc5"]}},
+            "thread_id": "2",
+        },
+        parent_config=expected_parent_config,
+        interrupts=(),
+    )
+
+    assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}},
+    ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str

libs/langgraph/tests/test_pregel_async.py

@@ -4572,6 +4572,194 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -
         ]
 
 
+@pytest.mark.parametrize("use_waiting_edge", (True, False))
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_in_one_fan_out_state_graph_defer_node(
+    checkpointer_name: str, use_waiting_edge: bool
+) -> None:
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    async def rewrite_query(data: State) -> State:
+        return {"query": f"query: {data['query']}"}
+
+    async def analyzer_one(data: State) -> State:
+        return {"query": f"analyzed: {data['query']}"}
+
+    async def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    async def retriever_two(data: State) -> State:
+        await asyncio.sleep(0.1)
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("analyzer_one", analyzer_one)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("qa", qa, defer=True)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_edge("rewrite_query", "analyzer_one")
+    workflow.add_edge("analyzer_one", "retriever_one")
+    workflow.add_edge("rewrite_query", "retriever_two")
+    if use_waiting_edge:
+        workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    else:
+        workflow.add_edge("retriever_one", "qa")
+        workflow.add_edge("retriever_two", "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == {
+        "query": "analyzed: query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        app_w_interrupt = workflow.compile(
+            checkpointer=checkpointer,
+            interrupt_after=["retriever_one"],
+        )
+        config = {"configurable": {"thread_id": "1"}}
+
+        assert [
+            c
+            async for c in app_w_interrupt.astream(
+                {"query": "what is weather in sf"}, config
+            )
+        ] == [
+            {"rewrite_query": {"query": "query: what is weather in sf"}},
+            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+            {"retriever_two": {"docs": ["doc3", "doc4"]}},
+            {"retriever_one": {"docs": ["doc1", "doc2"]}},
+            {"__interrupt__": ()},
+        ]
+
+        assert [c async for c in app_w_interrupt.astream(None, config)] == [
+            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        ]
+
+
+@pytest.mark.parametrize("with_path_map", (True, False))
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_in_one_fan_out_state_graph_then_defer_node(
+    checkpointer_name: str, with_path_map: bool
+) -> None:
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    async def rewrite_query(data: State) -> State:
+        return {"query": f"query: {data['query']}"}
+
+    async def analyzer_one(data: State) -> State:
+        return {"query": f"analyzed: {data['query']}"}
+
+    async def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    async def retriever_two(data: State) -> State:
+        await asyncio.sleep(0.1)
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("analyzer_one", analyzer_one)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("qa", qa, defer=True)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_conditional_edges(
+        "rewrite_query",
+        lambda _: ["analyzer_one", "retriever_two"],
+        ["analyzer_one", "retriever_two"] if with_path_map else None,
+        then="qa",
+    )
+    workflow.add_edge("analyzer_one", "retriever_one")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == {
+        "query": "analyzed: query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+        {"retriever_two": {"docs": ["doc3", "doc4"]}},
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        app_w_interrupt = workflow.compile(
+            checkpointer=checkpointer,
+            interrupt_after=["retriever_one"],
+        )
+        config = {"configurable": {"thread_id": "1"}}
+
+        assert [
+            c
+            async for c in app_w_interrupt.astream(
+                {"query": "what is weather in sf"}, config
+            )
+        ] == [
+            {"rewrite_query": {"query": "query: what is weather in sf"}},
+            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+            {"retriever_two": {"docs": ["doc3", "doc4"]}},
+            {"retriever_one": {"docs": ["doc1", "doc2"]}},
+            {"__interrupt__": ()},
+        ]
+
+        assert [c async for c in app_w_interrupt.astream(None, config)] == [
+            {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     snapshot: SnapshotAssertion, checkpointer_name: str

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

@@ -159,6 +159,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
             stream_keys=graph.stream_channels,
             interrupt_after=graph.interrupt_after_nodes,
             interrupt_before=graph.interrupt_before_nodes,
+            trigger_to_nodes=graph.trigger_to_nodes,
         ) as loop:
             if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved
@@ -345,6 +346,7 @@ class KafkaOrchestrator(AbstractContextManager):
             stream_keys=graph.stream_channels,
             interrupt_after=graph.interrupt_after_nodes,
             interrupt_before=graph.interrupt_before_nodes,
+            trigger_to_nodes=graph.trigger_to_nodes,
         ) as loop:
             if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved