Merge pull request #164 from langchain-ai/nc/29feb/waiting-edge

Add StateGraph.add_edge(string[], string)

langgraph/channels/any_value.py

@@ -38,6 +38,10 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     def update(self, values: Sequence[Value]) -> None:
         if len(values) == 0:
+            try:
+                del self.value
+            except AttributeError:
+                pass
             return
 
         self.value = values[-1]

langgraph/channels/named_barrier_value.py

@@ -0,0 +1,58 @@
+from contextlib import contextmanager
+from typing import Generator, Generic, Optional, Sequence, Type
+
+from typing_extensions import Self
+
+from langgraph.channels.base import (
+    BaseChannel,
+    EmptyChannelError,
+    InvalidUpdateError,
+    Value,
+)
+
+
+class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """A channel that waits until all named values are received before making the value available."""
+
+    def __init__(self, typ: Type[Value], names: set[str]) -> None:
+        self.typ = typ
+        self.names = names
+        self.seen = set()
+
+    @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
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ, self.names)
+        if checkpoint is not None:
+            empty.seen = checkpoint
+
+        try:
+            yield empty
+        finally:
+            pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        if self.seen == self.names:
+            self.seen = set()
+        for value in values:
+            if value in self.names:
+                self.seen.add(value)
+            else:
+                raise InvalidUpdateError(f"Value {value} not in {self.names}")
+
+    def get(self) -> Value:
+        if self.seen != self.names:
+            raise EmptyChannelError()
+        return None
+
+    def checkpoint(self) -> Value:
+        return self.seen

langgraph/graph/graph.py

@@ -50,6 +50,10 @@ class Graph:
         self.entry_point: Optional[str] = None
         self.entry_point_branch: Optional[Branch] = None
 
+    @property
+    def _all_edges(self) -> set[tuple[str, str]]:
+        return self.edges
+
     def add_node(self, key: str, action: RunnableLike) -> None:
         if self.compiled:
             logger.warning(
@@ -147,7 +151,9 @@ class Graph:
         return self.add_edge(key, END)
 
     def validate(self, interrupt: Optional[Sequence[str]] = None) -> None:
-        all_starts = {src for src, _ in self.edges} | {src for src in self.branches}
+        all_starts = {src for src, _ in self._all_edges} | {
+            src for src in self.branches
+        }
         for node in self.nodes:
             if node not in all_starts:
                 raise ValueError(f"Node `{node}` is a dead-end")
@@ -158,7 +164,7 @@ class Graph:
         if self.entry_point_branch is not None:
             branches.append(self.entry_point_branch)
 
-        all_hard_ends = {end for _, end in self.edges}
+        all_hard_ends = {end for _, end in self._all_edges}
         if self.entry_point is not None:
             all_hard_ends.add(self.entry_point)
 
@@ -275,7 +281,7 @@ class CompiledGraph(Pregel):
                 n = graph.add_node(node, key)
                 start_nodes[key] = n
                 end_nodes[key] = n
-        for start, end in self.graph.edges:
+        for start, end in self.graph._all_edges:
             graph.add_edge(start_nodes[start], end_nodes[end])
         for start, branches in self.graph.branches.items():
             for i, branch in enumerate(branches):

langgraph/graph/state.py

@@ -1,7 +1,8 @@
+import logging
 from collections import defaultdict
 from functools import partial
 from inspect import signature
-from typing import Any, Optional, Sequence, Type
+from typing import Any, Optional, Sequence, Type, Union
 
 from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.base import RunnableLike
@@ -11,12 +12,15 @@ from langgraph.channels.base import BaseChannel, InvalidUpdateError
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
+from langgraph.channels.named_barrier_value import NamedBarrierValue
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, START, CompiledGraph, Graph
 from langgraph.pregel import Channel
 from langgraph.pregel.read import ChannelInvoke
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 
+logger = logging.getLogger(__name__)
+
 
 class StateGraph(Graph):
     def __init__(self, schema: Type[Any]) -> None:
@@ -25,6 +29,13 @@ class StateGraph(Graph):
         self.channels = _get_channels(schema)
         if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()):
             self.support_multiple_edges = True
+        self.w_edges: set[tuple[tuple[str, ...], str]] = set()
+
+    @property
+    def _all_edges(self) -> set[tuple[str, str]]:
+        return self.edges | {
+            (start, end) for starts, end in self.w_edges for start in starts
+        }
 
     def add_node(self, key: str, action: RunnableLike) -> None:
         if key in self.channels:
@@ -34,6 +45,27 @@ class StateGraph(Graph):
             )
         return super().add_node(key, action)
 
+    def add_edge(self, start_key: Union[str, list[str]], end_key: str) -> None:
+        if isinstance(start_key, str):
+            return super().add_edge(start_key, end_key)
+
+        if self.compiled:
+            logger.warning(
+                "Adding an edge to a graph that has already been compiled. This will "
+                "not be reflected in the compiled graph."
+            )
+        for start in start_key:
+            if start == END:
+                raise ValueError("END cannot be a start node")
+            if start not in self.nodes:
+                raise ValueError(f"Need to add_node `{start}` first")
+        if end_key == END:
+            raise ValueError("END cannot be an end node")
+        if end_key not in self.nodes:
+            raise ValueError(f"Need to add_node `{end_key}` first")
+
+        self.w_edges.add((tuple(start_key), end_key))
+
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
@@ -68,14 +100,27 @@ class StateGraph(Graph):
             else None
         )
 
+        waiting_edges = {
+            (f"{starts}:{end}", starts, end) for starts, end in self.w_edges
+        }
+        waiting_edge_channels = {
+            key: NamedBarrierValue(str, set(starts)) for key, starts, _ in waiting_edges
+        }
+
         outgoing_edges = defaultdict(list)
         for start, end in self.edges:
             outgoing_edges[start].append(f"{end}:inbox" if end != END else END)
+        for key, starts, end in waiting_edges:
+            for start in starts:
+                outgoing_edges[start].append(key)
 
         nodes = {
             key: (
                 ChannelInvoke(
-                    triggers=[f"{key}:inbox"],
+                    triggers=[
+                        f"{key}:inbox",
+                        *[chan for chan, _, end in waiting_edges if end == key],
+                    ],
                     channels=state_channels,
                     mapper=coerce_state,
                 )
@@ -142,11 +187,15 @@ class StateGraph(Graph):
                 **self.channels,
                 **node_inboxes,
                 **node_outboxes,
+                **waiting_edge_channels,
                 END: LastValue(self.schema),
             },
             input=f"{START}:inbox",
             output=END,
-            hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,
+            hidden=[f"{node}:inbox" for node in self.nodes]
+            + [START]
+            + state_keys
+            + [key for key, _, _ in waiting_edges],
             snapshot_channels=state_keys_read,
             checkpointer=checkpointer,
             interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before],

langgraph/pregel/__init__.py

@@ -1018,12 +1018,18 @@ def _should_interrupt(
 
 
 def _read_channel(
-    channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
+    channels: Mapping[str, BaseChannel],
+    chan: str,
+    *,
+    catch: bool = True,
+    return_exception: bool = False,
 ) -> Any:
     try:
         return channels[chan].get()
-    except EmptyChannelError:
-        if catch:
+    except EmptyChannelError as exc:
+        if return_exception:
+            return exc
+        elif catch:
             return None
         else:
             raise
@@ -1090,7 +1096,7 @@ def _prepare_next_tasks(
     channels: Mapping[str, BaseChannel],
     update_seen: bool = True,
 ) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]:
-    checkpoint = copy_checkpoint(checkpoint) if update_seen else checkpoint
+    checkpoint = copy_checkpoint(checkpoint)
     tasks: list[tuple[Runnable, Any, str]] = []
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
@@ -1098,7 +1104,11 @@ def _prepare_next_tasks(
         seen = checkpoint["versions_seen"][name]
         # If any of the channels read by this process were updated
         if any(
-            checkpoint["channel_versions"][chan] > seen[chan] for chan in proc.triggers
+            checkpoint["channel_versions"][chan] > seen[chan]
+            for chan in proc.triggers
+            if not isinstance(
+                _read_channel(channels, chan, return_exception=True), EmptyChannelError
+            )
         ):
             # If all trigger channels subscribed by this process are not empty
             # then invoke the process with the values of all non-empty channels

langgraph/pregel/io.py

@@ -1,9 +1,21 @@
 from typing import Any, Iterator, Mapping, Optional, Sequence, Union
 
-from langgraph.channels.base import BaseChannel
+from langgraph.channels.base import BaseChannel, EmptyChannelError
 from langgraph.pregel.log import logger
 
 
+def _read_channel(
+    channels: Mapping[str, BaseChannel], chan: str, catch: bool = True
+) -> Any:
+    try:
+        return channels[chan].get()
+    except EmptyChannelError:
+        if catch:
+            return None
+        else:
+            raise
+
+
 def map_input(
     input_channels: Union[str, Sequence[str]],
     chunk: Optional[Union[dict[str, Any], Any]],
@@ -31,8 +43,8 @@ def map_output(
     """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
     if isinstance(output_channels, str):
         if any(chan == output_channels for chan, _ in pending_writes):
-            return channels[output_channels].get()
+            return _read_channel(channels, output_channels)
     else:
         if updated := {c for c, _ in pending_writes if c in output_channels}:
-            return {chan: channels[chan].get() for chan in updated}
+            return {chan: _read_channel(channels, chan) for chan in updated}
     return None

tests/conftest.py

@@ -10,3 +10,6 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture:
         UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000)
     )
     return mocker.patch("uuid.uuid4", side_effect=side_effect)
+
+
+pytest.register_assert_rewrite("tests.memory_assert")

tests/test_pregel.py

@@ -2783,3 +2783,358 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
             }
         },
     ]
+
+
+def test_in_one_fan_out_state_graph_waiting_edge() -> 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]
+
+    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:
+        return {"docs": ["doc3", "doc4"]}
+
+    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)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert app.get_graph().draw_ascii() == (
+        """              +-----------+              
+              | __start__ |              
+              +-----------+              
+                    *                    
+                    *                    
+                    *                    
+            +---------------+            
+            | rewrite_query |            
+            +---------------+            
+             ***         ***             
+            *               *            
+          **                 ***         
++--------------+                *        
+| analyzer_one |                *        
++--------------+                *        
+        *                       *        
+        *                       *        
+        *                       *        
++---------------+      +---------------+ 
+| retriever_one |      | retriever_two | 
++---------------+      +---------------+ 
+             ***         ***             
+                *       *                
+                 **   **                 
+                 +----+                  
+                 | qa |                  
+                 +----+                  
+                    *                    
+                    *                    
+                    *                    
+              +---------+                
+              | __end__ |                
+              +---------+                """
+    )
+
+    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"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_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"]},
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+    ]
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+
+def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> 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]
+
+    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:
+        return {"docs": ["doc3", "doc4"]}
+
+    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)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    workflow.set_finish_point("qa")
+
+    # silly edge, to make sure having been triggered before doesn't break
+    # semantics of named barrier (== waiting edges)
+    workflow.add_edge("rewrite_query", "qa")
+
+    app = workflow.compile()
+
+    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"]},
+            "qa": {"answer": ""},
+        },
+        {
+            "__end__": {
+                "answer": "",
+                "docs": ["doc3", "doc4"],
+                "query": "analyzed: query: what is weather in sf",
+            }
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_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"]},
+            "qa": {"answer": ""},
+        },
+        {
+            "__end__": {
+                "answer": "",
+                "docs": ["doc3", "doc4"],
+                "query": "analyzed: query: what is weather in sf",
+            }
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+    ]
+
+    assert [c for c in app_w_interrupt.stream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+
+def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> 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]
+
+    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:
+        return {"docs": ["doc3", "doc4"]}
+
+    def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    def decider(data: State) -> None:
+        return None
+
+    def decider_cond(data: State) -> str:
+        if data["query"].count("analyzed") > 1:
+            return "qa"
+        else:
+            return "rewrite_query"
+
+    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("decider", decider)
+    workflow.add_node("qa", qa)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
+    workflow.add_conditional_edges("decider", decider_cond)
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert app.invoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: analyzed: query: what is weather in sf",
+        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "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"]}},
+        {"decider": None},
+        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "analyzer_one": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf"
+            },
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"decider": None},
+        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf",
+                "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+                "docs": [
+                    "doc1",
+                    "doc1",
+                    "doc2",
+                    "doc2",
+                    "doc3",
+                    "doc3",
+                    "doc4",
+                    "doc4",
+                ],
+            }
+        },
+    ]

tests/test_pregel_async.py

@@ -2812,3 +2812,328 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
             }
         },
     ]
+
+
+async def test_in_one_fan_out_state_graph_waiting_edge() -> 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:
+        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)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"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 [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"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(), 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"]}},
+    ]
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+
+async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> 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:
+        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)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    workflow.set_finish_point("qa")
+
+    # silly edge, to make sure having been triggered before doesn't break
+    # semantics of named barrier (== waiting edges)
+    workflow.add_edge("rewrite_query", "qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"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 [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"]},
+            "qa": {"answer": ""},
+        },
+        {
+            "__end__": {
+                "answer": "",
+                "docs": ["doc3", "doc4"],
+                "query": "analyzed: query: what is weather in sf",
+            }
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(), 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"]},
+            "qa": {"answer": ""},
+        },
+        {
+            "__end__": {
+                "answer": "",
+                "docs": ["doc3", "doc4"],
+                "query": "analyzed: query: what is weather in sf",
+            }
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+    ]
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]
+
+
+async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> 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:
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    async def decider(data: State) -> None:
+        return None
+
+    def decider_cond(data: State) -> str:
+        if data["query"].count("analyzed") > 1:
+            return "qa"
+        else:
+            return "rewrite_query"
+
+    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("decider", decider)
+    workflow.add_node("qa", qa)
+
+    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")
+    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
+    workflow.add_conditional_edges("decider", decider_cond)
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: analyzed: query: what is weather in sf",
+        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "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"]}},
+        {"decider": None},
+        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "analyzer_one": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf"
+            },
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"decider": None},
+        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
+        {
+            "__end__": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf",
+                "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+                "docs": [
+                    "doc1",
+                    "doc1",
+                    "doc2",
+                    "doc2",
+                    "doc3",
+                    "doc3",
+                    "doc4",
+                    "doc4",
+                ],
+            }
+        },
+    ]