fix: Fix for drawing subgraphs with multiple sinks (#1962)

* Fix for drawing subgraphs with multiple sinks

* Expand error message

libs/langgraph/langgraph/graph/graph.py

@@ -574,9 +574,13 @@ class CompiledGraph(Pregel):
                 subgraph.trim_last_node()
                 if len(subgraph.nodes) > 1:
                     e, s = graph.extend(subgraph, prefix=key)
-                    if s is None or e is None:
-                        raise ValueError(f"Could not extend subgraph {key}")
-                    end_nodes[key], start_nodes[key] = e, s
+                    if e is None:
+                        raise ValueError(
+                            f"Could not extend subgraph '{key}' due to missing entrypoint"
+                        )
+                    if s is not None:
+                        start_nodes[key] = s
+                    end_nodes[key] = e
                 else:
                     nn = graph.add_node(node, key, metadata=metadata or None)
                     start_nodes[key] = nn

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -4535,6 +4535,29 @@
   
   '''
 # ---
+# name: test_multiple_sinks_subgraphs
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__([<p>__start__</p>]):::first
+  	uno(uno)
+  	dos(dos)
+  	subgraph_one(one)
+  	subgraph_two(two)
+  	subgraph_three(three)
+  	__start__ --> uno;
+  	uno -.-> dos;
+  	uno -.-> subgraph_one;
+  	subgraph subgraph
+  	subgraph_one -.-> subgraph_two;
+  	subgraph_one -.-> subgraph_three;
+  	end
+  	classDef default fill:#f2f0ff,line-height:1.2
+  	classDef first fill-opacity:0
+  	classDef last fill:#bfb6fc
+  
+  '''
+# ---
 # name: test_nested_graph
   '''
   graph TD;

libs/langgraph/tests/test_pregel.py

@@ -11408,6 +11408,29 @@ def test_xray_bool(snapshot: SnapshotAssertion) -> None:
     assert app.get_graph(xray=True).draw_mermaid() == snapshot
 
 
+def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None:
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    subgraph_builder = StateGraph(State)
+    subgraph_builder.add_node("one", lambda x: x)
+    subgraph_builder.add_node("two", lambda x: x)
+    subgraph_builder.add_node("three", lambda x: x)
+    subgraph_builder.add_edge("__start__", "one")
+    subgraph_builder.add_conditional_edges("one", lambda x: "two", ["two", "three"])
+    subgraph = subgraph_builder.compile()
+
+    builder = StateGraph(State)
+    builder.add_node("uno", lambda x: x)
+    builder.add_node("dos", lambda x: x)
+    builder.add_node("subgraph", subgraph)
+    builder.add_edge("__start__", "uno")
+    builder.add_conditional_edges("uno", lambda x: "dos", ["dos", "subgraph"])
+
+    app = builder.compile()
+    assert app.get_graph(xray=True).draw_mermaid() == snapshot
+
+
 def test_subgraph_retries():
     class State(TypedDict):
         count: int