When using Command.PARENT also pass existing subgraph state to parent graph (#3134)

Co-authored-by: Vadym Barda <vadym@langchain.dev>

libs/langgraph/langgraph/pregel/io.py

@@ -74,7 +74,7 @@ def map_command(
 ) -> Iterator[tuple[str, str, Any]]:
     """Map input chunk to a sequence of pending writes in the form (channel, value)."""
     if cmd.graph == Command.PARENT:
-        raise InvalidUpdateError("There is not parent graph")
+        raise InvalidUpdateError("There is no parent graph")
     if cmd.goto:
         if isinstance(cmd.goto, (tuple, list)):
             sends = cmd.goto

libs/langgraph/langgraph/pregel/loop.py

@@ -2,6 +2,7 @@ import asyncio
 import concurrent.futures
 from collections import defaultdict, deque
 from contextlib import AsyncExitStack, ExitStack
+from dataclasses import replace
 from inspect import signature
 from types import TracebackType
 from typing import (
@@ -66,6 +67,7 @@ from langgraph.errors import (
     EmptyInputError,
     GraphDelegate,
     GraphInterrupt,
+    ParentCommand,
 )
 from langgraph.managed.base import (
     ManagedValueMapping,
@@ -736,6 +738,16 @@ class PregelLoop(LoopProtocol):
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
+        # add current state to parent command
+        if isinstance(exc_value, ParentCommand):
+            cmd = exc_value.args[0]
+            state = (
+                [(self.output_keys, read_channels(self.channels, self.output_keys))]
+                if isinstance(self.output_keys, str)
+                else list(read_channels(self.channels, self.output_keys).items())
+            )
+            exc_value.args = (replace(cmd, update=[*state, *cmd._update_as_tuples()]),)
+        # suppress interrupt
         suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
         if suppress:
             # emit one last "values" event, with pending writes applied

libs/langgraph/tests/test_pregel.py

@@ -4828,6 +4828,13 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str)
             "source": "loop",
             "writes": {
                 "alice": {
+                    "messages": [
+                        _AnyIdHumanMessage(
+                            content="get user name",
+                            additional_kwargs={},
+                            response_metadata={},
+                        ),
+                    ],
                     "user_name": "Meow",
                 }
             },
@@ -6167,6 +6174,88 @@ def test_multiple_subgraphs_checkpointer(
     ]
 
 
+def test_merging_updates_command_parent():
+    # simple reducer
+    def append_unique(left, right):
+        combined = list(left)
+        for item in right:
+            if item in combined:
+                continue
+            else:
+                combined.append(item)
+        return combined
+
+    class State(TypedDict):
+        foo: str
+        bar: Annotated[list[str], append_unique]
+
+    # Define subgraph
+    def subgraph_node_1(state: State):
+        return Command(
+            goto="subgraph_node_2",
+            update={
+                "foo": "foo",
+                "bar": ["subgraph_node_1"],
+            },
+        )
+
+    def subgraph_node_2(state: State):
+        return Command(
+            goto="node_3",
+            update={"bar": ["subgraph_node_2"]},
+            graph=Command.PARENT,
+        )
+
+    subgraph_builder = StateGraph(State)
+    subgraph_builder.add_node(subgraph_node_1)
+    subgraph_builder.add_node(subgraph_node_2)
+    subgraph_builder.add_edge(START, "subgraph_node_1")
+
+    # Define main graph
+    def node_1(state: State):
+        return Command(
+            goto="node_2",
+            update={"bar": ["node_1"]},
+        )
+
+    def node_3(state: State, store):
+        return Command(
+            update={"bar": ["node_3"]},
+        )
+
+    main_builder = StateGraph(State)
+    main_builder.add_node("node_1", node_1)
+    main_builder.add_node("node_2", subgraph_builder.compile())
+    main_builder.add_node("node_3", node_3)
+    main_builder.add_edge(START, "node_1")
+    main_builder.add_edge("node_2", "node_3")
+    main_graph = main_builder.compile()
+
+    assert main_graph.invoke({"foo": ""}) == {
+        "foo": "foo",
+        "bar": ["node_1", "subgraph_node_1", "subgraph_node_2", "node_3"],
+    }
+
+    assert list(
+        main_graph.stream({"foo": ""}, stream_mode="updates", subgraphs=True)
+    ) == [
+        ((), {"node_1": {"bar": ["node_1"]}}),
+        (
+            (AnyStr("node_2:"),),
+            {"subgraph_node_1": {"foo": "foo", "bar": ["subgraph_node_1"]}},
+        ),
+        (
+            (),
+            {
+                "node_2": [
+                    {"foo": "foo"},
+                    {"bar": ["node_1", "subgraph_node_1"]},
+                    {"bar": ["subgraph_node_2"]},
+                ]
+            },
+        ),
+        ((), {"node_3": {"bar": ["node_3"]}}),
+    ]
 def test_entrypoint_output_schema_with_return_and_save() -> None:
     """Test output schema inference with entrypoint.final."""
 

libs/langgraph/tests/test_pregel_async.py

@@ -6148,6 +6148,13 @@ async def test_parent_command(checkpointer_name: str) -> None:
                 "source": "loop",
                 "writes": {
                     "alice": {
+                        "messages": [
+                            _AnyIdHumanMessage(
+                                content="get user name",
+                                additional_kwargs={},
+                                response_metadata={},
+                            ),
+                        ],
                         "user_name": "Meow",
                     }
                 },