langgraph: support streaming messages from Command.update (#4250)

libs/langgraph/langgraph/pregel/messages.py

@@ -18,7 +18,7 @@ from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatGenerationChunk, LLMResult
 
 from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
-from langgraph.types import StreamChunk
+from langgraph.types import Command, StreamChunk
 
 try:
     from langchain_core.tracers._streaming import _StreamingCallbackHandler
@@ -153,6 +153,17 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
         **kwargs: Any,
     ) -> Any:
         if meta := self.metadata.pop(run_id, None):
+            if isinstance(response, Command):
+                response = response.update
+
+            if isinstance(response, Sequence) and any(
+                isinstance(value, Command) for value in response
+            ):
+                response = [
+                    value.update if isinstance(value, Command) else value
+                    for value in response
+                ]
+
             if isinstance(response, BaseMessage):
                 self._emit(meta, response, dedupe=True)
             elif isinstance(response, Sequence):

libs/langgraph/tests/test_pregel.py

@@ -7168,6 +7168,52 @@ def test_tags_stream_mode_messages() -> None:
     ]
 
 
+def test_stream_mode_messages_command() -> None:
+    from langchain_core.messages import HumanMessage
+
+    def my_node(state):
+        return {"messages": HumanMessage(content="foo")}
+
+    def my_other_node(state):
+        return Command(update={"messages": HumanMessage(content="bar")})
+
+    graph = (
+        StateGraph(MessagesState)
+        .add_sequence([my_node, my_other_node])
+        .add_edge(START, "my_node")
+        .compile()
+    )
+    assert list(
+        graph.stream(
+            {
+                "messages": [],
+            },
+            stream_mode="messages",
+        )
+    ) == [
+        (
+            _AnyIdHumanMessage(content="foo"),
+            {
+                "langgraph_step": 1,
+                "langgraph_node": "my_node",
+                "langgraph_triggers": ("branch:to:my_node",),
+                "langgraph_path": ("__pregel_pull", "my_node"),
+                "langgraph_checkpoint_ns": AnyStr("my_node:"),
+            },
+        ),
+        (
+            _AnyIdHumanMessage(content="bar"),
+            {
+                "langgraph_step": 2,
+                "langgraph_node": "my_other_node",
+                "langgraph_triggers": ("branch:to:my_other_node",),
+                "langgraph_path": ("__pregel_pull", "my_other_node"),
+                "langgraph_checkpoint_ns": AnyStr("my_other_node:"),
+            },
+        ),
+    ]
+
+
 def test_node_destinations() -> None:
     class State(TypedDict):
         foo: Annotated[str, operator.add]

libs/langgraph/tests/test_pregel_async.py

@@ -7894,6 +7894,53 @@ async def test_tags_stream_mode_messages() -> None:
     ]
 
 
+async def test_stream_mode_messages_command() -> None:
+    from langchain_core.messages import HumanMessage
+
+    async def my_node(state):
+        return {"messages": HumanMessage(content="foo")}
+
+    async def my_other_node(state):
+        return Command(update={"messages": HumanMessage(content="bar")})
+
+    graph = (
+        StateGraph(MessagesState)
+        .add_sequence([my_node, my_other_node])
+        .add_edge(START, "my_node")
+        .compile()
+    )
+    assert [
+        c
+        async for c in graph.astream(
+            {
+                "messages": [],
+            },
+            stream_mode="messages",
+        )
+    ] == [
+        (
+            _AnyIdHumanMessage(content="foo"),
+            {
+                "langgraph_step": 1,
+                "langgraph_node": "my_node",
+                "langgraph_triggers": ("branch:to:my_node",),
+                "langgraph_path": ("__pregel_pull", "my_node"),
+                "langgraph_checkpoint_ns": AnyStr("my_node:"),
+            },
+        ),
+        (
+            _AnyIdHumanMessage(content="bar"),
+            {
+                "langgraph_step": 2,
+                "langgraph_node": "my_other_node",
+                "langgraph_triggers": ("branch:to:my_other_node",),
+                "langgraph_path": ("__pregel_pull", "my_other_node"),
+                "langgraph_checkpoint_ns": AnyStr("my_other_node:"),
+            },
+        ),
+    ]
+
+
 async def test_stream_messages_dedupe_inputs() -> None:
     from langchain_core.messages import AIMessage