fix(langgraph): support emitting messages from subgraphs when messages mode explicitly requested (#5836)

Reproduces:
https://github.com/langchain-ai/langgraph/issues/5249#issuecomment-3156519635
Caused after this change:
https://github.com/langchain-ai/langgraph/pull/4843

Fix to allow emitting messages from subgraphs if the subgraphs
explicitly used a stream mode "messages".

```python

def node_in_parent(...):
   # subgraph was called as a function.
   # messages are explicitly requested.
   for event in subgraph.stream(..., stream_mode="messages"):
      # something is done with `event`
   return ...

# subgraphs = False!
parent_graph.invoke(..., subgraphs=False)
```

The code above should continue to work correctly regardless of the value
of subgraphs as streaming messages was requested explicitly in the
parent node!

---------

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>

libs/langgraph/langgraph/pregel/_messages.py

@@ -29,16 +29,51 @@ Meta = tuple[tuple[str, ...], dict[str, Any]]
 
 class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
     """A callback handler that implements stream_mode=messages.
-    Collects messages from (1) chat model stream events and (2) node outputs."""
+
+    Collects messages from:
+    (1) chat model stream events; and
+    (2) node outputs.
+    """
 
     run_inline = True
-    """We want this callback to run in the main thread, to avoid order/locking issues."""
+    """We want this callback to run in the main thread to avoid order/locking issues."""
 
-    def __init__(self, stream: Callable[[StreamChunk], None], subgraphs: bool):
+    def __init__(
+        self,
+        stream: Callable[[StreamChunk], None],
+        subgraphs: bool,
+        *,
+        parent_ns: tuple[str, ...] | None = None,
+    ) -> None:
+        """Configure the handler to stream messages from LLMs and nodes.
+
+        Args:
+            stream: A callable that takes a StreamChunk and emits it.
+            subgraphs: Whether to emit messages from subgraphs.
+            parent_ns: The namespace where the handler was created.
+                We keep track of this namespace to allow calls to subgraphs that
+                were explicitly requested as a stream with `messages` mode
+                configured.
+
+        Example:
+            parent_ns is used to handle scenarios where the subgraph is explicitly
+            streamed with `stream_mode="messages"`.
+
+            ```python
+            def parent_graph_node():
+                # This node is in the parent graph.
+                async for event in some_subgraph(..., stream_mode="messages"):
+                    do something with event # <-- these events will be emitted
+                return ...
+
+            parent_graph.invoke(subgraphs=False)
+            ```
+        """
         self.stream = stream
         self.subgraphs = subgraphs
         self.metadata: dict[UUID, Meta] = {}
         self.seen: set[int | str] = set()
+        self.parent_ns = parent_ns
 
     def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None:
         if dedupe and message.id in self.seen:
@@ -100,7 +135,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
             ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
                 :-1
             ]
-            if not self.subgraphs and len(ns) > 0:
+            if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
                 return
             if tags:
                 if filtered_tags := [t for t in tags if not t.startswith("seq:step")]:

libs/langgraph/langgraph/pregel/main.py

@@ -11,7 +11,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from dataclasses import is_dataclass
 from functools import partial
 from inspect import isclass
-from typing import Any, Callable, Generic, Union, cast, get_type_hints
+from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints
 from uuid import UUID, uuid5
 
 from langchain_core.globals import get_debug
@@ -2534,8 +2534,13 @@ class Pregel(
                 config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns)
             # set up messages stream mode
             if "messages" in stream_modes:
+                ns_ = cast(Optional[str], config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
                 run_manager.inheritable_handlers.append(
-                    StreamMessagesHandler(stream.put, subgraphs)
+                    StreamMessagesHandler(
+                        stream.put,
+                        subgraphs,
+                        parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
+                    )
                 )
 
             # set up custom stream mode
@@ -2814,8 +2819,14 @@ class Pregel(
                 config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns)
             # set up messages stream mode
             if "messages" in stream_modes:
+                # namespace can be None in a root level graph?
+                ns_ = cast(Optional[str], config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
                 run_manager.inheritable_handlers.append(
-                    StreamMessagesHandler(stream_put, subgraphs)
+                    StreamMessagesHandler(
+                        stream_put,
+                        subgraphs,
+                        parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
+                    )
                 )
 
             # set up custom stream mode

libs/langgraph/tests/test_pregel.py

@@ -27,7 +27,7 @@ from langsmith import traceable
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
 from langgraph.cache.base import BaseCache
@@ -8266,3 +8266,53 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) ->
             ],
         ],
     ]
+
+
+def test_subgraph_streaming_sync() -> None:
+    """Test subgraph streaming when used as a node in sync version"""
+
+    # Create a fake chat model that returns a simple response
+    model = GenericFakeChatModel(messages=iter(["The weather is sunny today."]))
+
+    # Create a subgraph that uses the fake chat model
+    def call_model_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
+        """Node that calls the model with the last message."""
+        messages = state["messages"]
+        last_message = messages[-1].content if messages else ""
+        response = model.invoke([("user", last_message)], config)
+        return {"messages": [response]}
+
+    # Build the subgraph
+    subgraph = StateGraph(MessagesState)
+    subgraph.add_node("call_model", call_model_node)
+    subgraph.add_edge(START, "call_model")
+    compiled_subgraph = subgraph.compile()
+
+    class SomeCustomState(TypedDict):
+        last_chunk: NotRequired[str]
+        num_chunks: NotRequired[int]
+
+    # Will invoke a subgraph as a function
+    def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict:
+        """Node that runs the subgraph."""
+        msgs = {"messages": [("user", "What is the weather in Tokyo?")]}
+        events = []
+        for event in compiled_subgraph.stream(msgs, config, stream_mode="messages"):
+            events.append(event)
+        ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events]
+        return {
+            "last_chunk": ai_msg_chunks[-1],
+            "num_chunks": len(ai_msg_chunks),
+        }
+
+    # Build the main workflow
+    workflow = StateGraph(SomeCustomState)
+    workflow.add_node("subgraph", parent_node)
+    workflow.add_edge(START, "subgraph")
+    compiled_workflow = workflow.compile()
+
+    # Test the basic functionality
+    result = compiled_workflow.invoke({})
+
+    assert result["last_chunk"].content == "today."
+    assert result["num_chunks"] == 9

libs/langgraph/tests/test_pregel_async.py

@@ -26,7 +26,7 @@ from langchain_core.utils.aiter import aclosing
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
 from langgraph.cache.base import BaseCache
@@ -9053,3 +9053,57 @@ async def test_fork_and_update_task_results(
             ],
         ],
     ]
+
+
+async def test_subgraph_streaming_async() -> None:
+    """Test subgraph streaming when used as a node in async version"""
+
+    # Create a fake chat model that returns a simple response
+    model = GenericFakeChatModel(messages=iter(["The weather is sunny today."]))
+
+    # Create a subgraph that uses the fake chat model
+    async def call_model_node(
+        state: MessagesState, config: RunnableConfig
+    ) -> MessagesState:
+        """Node that calls the model with the last message."""
+        messages = state["messages"]
+        last_message = messages[-1].content if messages else ""
+        response = await model.ainvoke([("user", last_message)], config)
+        return {"messages": [response]}
+
+    # Build the subgraph
+    subgraph = StateGraph(MessagesState)
+    subgraph.add_node("call_model", call_model_node)
+    subgraph.add_edge(START, "call_model")
+    compiled_subgraph = subgraph.compile()
+
+    class SomeCustomState(TypedDict):
+        last_chunk: NotRequired[str]
+        num_chunks: NotRequired[int]
+
+    # Will invoke a subgraph as a function
+    async def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict:
+        """Node that runs the subgraph."""
+        msgs = {"messages": [("user", "What is the weather in Tokyo?")]}
+        events = []
+        async for event in compiled_subgraph.astream(
+            msgs, config, stream_mode="messages"
+        ):
+            events.append(event)
+        ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events]
+        return {
+            "last_chunk": ai_msg_chunks[-1],
+            "num_chunks": len(ai_msg_chunks),
+        }
+
+    # Build the main workflow
+    workflow = StateGraph(SomeCustomState)
+    workflow.add_node("subgraph", parent_node)
+    workflow.add_edge(START, "subgraph")
+    compiled_workflow = workflow.compile()
+
+    # Test the basic functionality
+    result = await compiled_workflow.ainvoke({})
+
+    assert result["last_chunk"].content == "today."
+    assert result["num_chunks"] == 9

libs/prebuilt/tests/test_react_agent.py

@@ -24,7 +24,7 @@ from langchain_core.messages import (
     ToolCall,
     ToolMessage,
 )
-from langchain_core.runnables import RunnableLambda
+from langchain_core.runnables import RunnableConfig, RunnableLambda
 from langchain_core.tools import InjectedToolCallId, ToolException
 from langchain_core.tools import tool as dec_tool
 from pydantic import BaseModel, Field
@@ -1236,6 +1236,190 @@ def test_tool_node_stream_writer() -> None:
     ]
 
 
+@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
+def test_react_agent_subgraph_streaming_sync(version: Literal["v1", "v2"]) -> None:
+    """Test React agent streaming when used as a subgraph node sync version"""
+
+    @dec_tool
+    def get_weather(city: str) -> str:
+        """Get the weather of a city."""
+        return f"The weather of {city} is sunny."
+
+    # Create a React agent
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"args": {"city": "Tokyo"}, "id": "1", "name": "get_weather"}],
+            [],
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather],
+        prompt="You are a helpful travel assistant.",
+        version=version,
+    )
+
+    # Create a subgraph that uses the React agent as a node
+    def react_agent_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
+        """Node that runs the React agent and collects streaming output."""
+        collected_content = ""
+
+        # Stream the agent output and collect content
+        for msg_chunk, msg_metadata in agent.stream(
+            {"messages": [("user", state["messages"][-1].content)]},
+            config,
+            stream_mode="messages",
+        ):
+            if hasattr(msg_chunk, "content") and msg_chunk.content:
+                collected_content += msg_chunk.content
+
+        return {"messages": [("assistant", collected_content)]}
+
+    # Create the main workflow with the React agent as a subgraph node
+    workflow = StateGraph(MessagesState)
+    workflow.add_node("react_agent", react_agent_node)
+    workflow.add_edge(START, "react_agent")
+    workflow.add_edge("react_agent", "__end__")
+    compiled_workflow = workflow.compile()
+
+    # Test the streaming functionality
+    result = compiled_workflow.invoke(
+        {"messages": [("user", "What is the weather in Tokyo?")]}
+    )
+
+    # Verify the result contains expected structure
+    assert len(result["messages"]) == 2
+    assert result["messages"][0].content == "What is the weather in Tokyo?"
+    assert "assistant" in str(result["messages"][1])
+
+    # Test streaming with subgraphs = True
+    result = compiled_workflow.invoke(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        subgraphs=True,
+    )
+    assert len(result["messages"]) == 2
+
+    events = []
+    for event in compiled_workflow.stream(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        stream_mode="messages",
+        subgraphs=False,
+    ):
+        events.append(event)
+
+    assert len(events) == 0
+
+    events = []
+    for event in compiled_workflow.stream(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        stream_mode="messages",
+        subgraphs=True,
+    ):
+        events.append(event)
+
+    assert len(events) == 3
+    namespace, (msg, metadata) = events[0]
+    # FakeToolCallingModel returns a single AIMessage with tool calls
+    # The content of the AIMessage reflects the input message
+    assert msg.content.startswith("You are a helpful travel assistant")
+    namespace, (msg, metadata) = events[1]  # ToolMessage
+    assert msg.content.startswith("The weather of Tokyo is sunny.")
+
+
+@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
+async def test_react_agent_subgraph_streaming(version: Literal["v1", "v2"]) -> None:
+    """Test React agent streaming when used as a subgraph node."""
+
+    @dec_tool
+    def get_weather(city: str) -> str:
+        """Get the weather of a city."""
+        return f"The weather of {city} is sunny."
+
+    # Create a React agent
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"args": {"city": "Tokyo"}, "id": "1", "name": "get_weather"}],
+            [],
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather],
+        prompt="You are a helpful travel assistant.",
+        version=version,
+    )
+
+    # Create a subgraph that uses the React agent as a node
+    async def react_agent_node(
+        state: MessagesState, config: RunnableConfig
+    ) -> MessagesState:
+        """Node that runs the React agent and collects streaming output."""
+        collected_content = ""
+
+        # Stream the agent output and collect content
+        async for msg_chunk, msg_metadata in agent.astream(
+            {"messages": [("user", state["messages"][-1].content)]},
+            config,
+            stream_mode="messages",
+        ):
+            if hasattr(msg_chunk, "content") and msg_chunk.content:
+                collected_content += msg_chunk.content
+
+        return {"messages": [("assistant", collected_content)]}
+
+    # Create the main workflow with the React agent as a subgraph node
+    workflow = StateGraph(MessagesState)
+    workflow.add_node("react_agent", react_agent_node)
+    workflow.add_edge(START, "react_agent")
+    workflow.add_edge("react_agent", "__end__")
+    compiled_workflow = workflow.compile()
+
+    # Test the streaming functionality
+    result = await compiled_workflow.ainvoke(
+        {"messages": [("user", "What is the weather in Tokyo?")]}
+    )
+
+    # Verify the result contains expected structure
+    assert len(result["messages"]) == 2
+    assert result["messages"][0].content == "What is the weather in Tokyo?"
+    assert "assistant" in str(result["messages"][1])
+
+    # Test streaming with subgraphs = True
+    result = await compiled_workflow.ainvoke(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        subgraphs=True,
+    )
+    assert len(result["messages"]) == 2
+
+    events = []
+    async for event in compiled_workflow.astream(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        stream_mode="messages",
+        subgraphs=False,
+    ):
+        events.append(event)
+
+    assert len(events) == 0
+
+    events = []
+    async for event in compiled_workflow.astream(
+        {"messages": [("user", "What is the weather in Tokyo?")]},
+        stream_mode="messages",
+        subgraphs=True,
+    ):
+        events.append(event)
+
+    assert len(events) == 3
+    namespace, (msg, metadata) = events[0]
+    # FakeToolCallingModel returns a single AIMessage with tool calls
+    # The content of the AIMessage reflects the input message
+    assert msg.content.startswith("You are a helpful travel assistant")
+    namespace, (msg, metadata) = events[1]  # ToolMessage
+    assert msg.content.startswith("The weather of Tokyo is sunny.")
+
+
 @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
 def test_tool_node_node_interrupt(
     sync_checkpointer: BaseCheckpointSaver, version: str