diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index ffe35e8bc..9331e5163 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -418,6 +418,15 @@ class StateGraph(Graph): if not isinstance(val, Context) and not is_managed_value(val) ] ) + stream_channels = ( + "__root__" + if len(self.channels) == 1 and "__root__" in self.channels + else [ + key + for key, val in self.channels.items() + if not isinstance(val, Context) and not is_managed_value(val) + ] + ) compiled = CompiledStateGraph( builder=self, @@ -427,7 +436,7 @@ class StateGraph(Graph): input_channels=START, stream_mode="updates", output_channels=output_channels, - stream_channels=output_channels, + stream_channels=stream_channels, checkpointer=checkpointer, interrupt_before_nodes=interrupt_before, interrupt_after_nodes=interrupt_after, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 08cc53fbd..e9fbf8991 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -310,6 +310,20 @@ def test_node_schemas_custom_output() -> None: "messages": [_AnyIdHumanMessage(content="hello")], } + assert [ + c + for c in graph.stream( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) + ] == [ + {"b": {"hello": "again", "now": 123}}, + ] + def test_reducer_before_first_node() -> None: class State(TypedDict): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 23065b3d2..73ab714d3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -493,6 +493,20 @@ async def test_node_schemas_custom_output() -> None: "messages": [_AnyIdHumanMessage(content="hello")], } + assert [ + c + async for c in graph.astream( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) + ] == [ + {"b": {"hello": "again", "now": 123}}, + ] + async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1)