langgraph: fix drawing graph with __root__ channel

libs/langgraph/langgraph/graph/state.py

@@ -1114,7 +1114,7 @@ def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
             else:
                 updates.append(("__root__", i))
         return updates
-    elif input is not None:
+    else:
         return [("__root__", input)]
 
 

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -396,6 +396,61 @@
   
   '''
 # ---
+# name: test_get_graph_root_channel
+  '''
+  {
+    "nodes": [
+      {
+        "id": "__start__",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langchain",
+            "schema",
+            "runnable",
+            "RunnablePassthrough"
+          ],
+          "name": "__start__"
+        }
+      },
+      {
+        "id": "child",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langgraph",
+            "graph",
+            "state",
+            "CompiledStateGraph"
+          ],
+          "name": "child"
+        }
+      },
+      {
+        "id": "__end__"
+      }
+    ],
+    "edges": [
+      {
+        "source": "__start__",
+        "target": "child"
+      },
+      {
+        "source": "child",
+        "target": "__end__"
+      }
+    ]
+  }
+  '''
+# ---
+# name: test_get_graph_root_channel.1
+  '''
+  graph TD;
+  	__start__ --> child;
+  	child --> __end__;
+  
+  '''
+# ---
 # name: test_get_graph_self_loop
   '''
   {

libs/langgraph/tests/test_pregel.py

@@ -8754,3 +8754,18 @@ def test_get_graph_self_loop(snapshot: SnapshotAssertion) -> None:
 
     assert json.dumps(self_loop_graph.get_graph().to_json(), indent=2) == snapshot
     assert self_loop_graph.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+
+def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None:
+    child_builder = StateGraph(str)
+    child_builder.add_node("child_node", lambda x: x)
+    child_builder.add_edge(START, "child_node")
+    child_graph = child_builder.compile()
+
+    graph_builder = StateGraph(str)
+    graph_builder.add_node("child", child_graph)
+    graph_builder.add_edge(START, "child")
+    graph = graph_builder.compile()
+
+    assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot
+    assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot