From 254a38560e0f036531ac2cde7f55ce93a2cc3ff0 Mon Sep 17 00:00:00 2001
From: vbarda <vadym@langchain.dev>
Date: Wed, 14 May 2025 20:35:10 -0400
Subject: [PATCH] langgraph: fix drawing graph with __root__ channel

---
 libs/langgraph/langgraph/graph/state.py       |  2 +-
 .../tests/__snapshots__/test_pregel.ambr      | 55 +++++++++++++++++++
 libs/langgraph/tests/test_pregel.py           | 15 +++++
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index f6a29198b..68a75fa8e 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/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)]
 
 
diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr
index 4caad8801..9559651f8 100644
--- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr
+++ b/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
   '''
   {
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index b251cf465..461791e00 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/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