langgraph: fix drawing graph with __root__ channel (#4695)

libs/langgraph/langgraph/pregel/draw.py

@@ -108,7 +108,12 @@ def draw_graph(
                 for w in task.writers:
                     # apply regular writes
                     if isinstance(w, ChannelWrite):
-                        w.invoke(None, task.config)
+                        empty_input = (
+                            cast(BaseChannel, specs["__root__"]).ValueType()
+                            if "__root__" in specs
+                            else None
+                        )
+                        w.invoke(empty_input, task.config)
                     # apply conditional writes declared for static analysis, only once
                     if w not in static_seen:
                         static_seen.add(w)
@@ -120,7 +125,7 @@ def draw_graph(
                                     edges.add((task.name, t[0], True, t[2]))
                             writes = [t for t in writes if t[0] != END]
                             conditionals.update(
-                                {(task.name, *t[:2]): t[2] for t in writes}
+                                {(task.name, t[0], t[1] or None): t[2] for t in writes}
                             )
                             task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
             # collect sources
@@ -128,8 +133,8 @@ def draw_graph(
                 task.name: {
                     (
                         w[0],
-                        (task.name, *w) in conditionals,
-                        conditionals.get((task.name, *w)),
+                        (task.name, w[0], w[1] or None) in conditionals,
+                        conditionals.get((task.name, w[0], w[1] or None)),
                     )
                     for w in task.writes
                 }

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(list)
+    child_builder.add_node("child_node", lambda x: x)
+    child_builder.add_edge(START, "child_node")
+    child_graph = child_builder.compile()
+
+    graph_builder = StateGraph(list)
+    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

libs/langgraph/tests/test_remote_graph.py

@@ -863,7 +863,9 @@ async def test_ainvoke():
     assert result == {"messages": [{"type": "human", "content": "world"}]}
 
 
-@pytest.mark.skip("Unskip this test to manually test the LangGraph Platform integration")
+@pytest.mark.skip(
+    "Unskip this test to manually test the LangGraph Platform integration"
+)
 @pytest.mark.anyio
 async def test_langgraph_cloud_integration():
     from langgraph_sdk.client import get_client, get_sync_client