Fix output schema affecting stream output (#1239)

* Fix output schema affecting stream output

* Remove file

* Remove file

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,

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):

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)