From bdef6b3f5d700d5cf8bbf0ffab06fa25f4792e61 Mon Sep 17 00:00:00 2001 From: Caspar Broekhuizen Date: Wed, 10 Sep 2025 08:28:33 -0700 Subject: [PATCH] fix(langgraph): get_graph generates unexpected conditional edge (#6122) ### Description * Fix `get_graph()` generating an unexpected conditional edge to `__end__` when the last step has a single (non-terminal) source and the graph is cyclic. ### Issue * There was a fallback path that was triggered in `draw_graph()` when, for a Pregel instance; no termini exist and there is only a single step source (the last one). In this case an edge was added: (last source) -> `__end__`, even when another node already had a valid edge: (node) -> `__end__`. See this example:
code ```python from langgraph.graph import END, START, StateGraph from pydantic import BaseModel class State(TypedDict): messages: list[str] def chatbot_node(state: State) -> State: return {"messages": state["messages"] + ["chatbot"]} def tools_node(state: State) -> State: return {"messages": state["messages"] + ["tools"]} def human_node(state: State) -> State: return {"messages": state["messages"] + ["human"]} def tools_condition(_: State) -> str: return "tools" def end_condition(_: State) -> str: return "chatbot" workflow = StateGraph(State) workflow.add_node("chatbot", chatbot_node) workflow.add_node("tools", tools_node) workflow.add_node("human", human_node) workflow.add_edge(START, "human") workflow.add_edge("tools", "chatbot") # graph_builder.add_edge("chatbot", "human") !!! workflow.add_conditional_edges( "chatbot", tools_condition, {"tools": "tools", "human": "human"} ) workflow.add_conditional_edges( "human", end_condition, {"chatbot": "chatbot", END: END} ) app = workflow.compile() mermaid = app.get_graph().draw_mermaid() ```
The code above, as-is, generates the graph on the left. There is an unexpected conditional edge: chatbot -> `__end__`. If you uncomment the commented line and introduce a static edge: chatbot -> human, `get_graph()` returns the correct representation: 1 Without `graph_builder.add_edge("chatbot", "human")` | 2 With `graph_builder.add_edge("chatbot", "human")` :-------------------------:|:-------------------------: ![](https://github.com/user-attachments/assets/aa3149c2-ceee-4c0c-9c0c-e999caf042f0) | ![](https://github.com/user-attachments/assets/ea53287f-1d68-47d6-8b36-9ec7ca1d52fa) * In case 1), the graph is cyclic so termini is empty, and the last `step_sources` set during the static walk contains only the chatbot node, so an edge is added: chatbot -> `__end__`. * In case 2), the graph is cyclic so termini is empty, and the last `step_sources` set during the static walk contains only the human node, so an edge is added: human -> `__end__`, but `add_edge()` dedups (the edge already exists) so the graph appears correct. ### Solution * Check that no valid edges: (node) -> `__end__` exist before triggering the fallback path and creating an edge. Before | After :-------------------------:|:-------------------------: ![](https://github.com/user-attachments/assets/aa3149c2-ceee-4c0c-9c0c-e999caf042f0) | ![](https://github.com/user-attachments/assets/9de4ab4f-6503-4894-bfda-37aba1d1be05) After: The graph is cyclic so termini is empty, and the last `step_sources` contains the chatbot node, but an edge already exists: human -> `__end__`, so no more edges are added. ### Tests * `test_get_graph_nonterminal_last_step_source()` which asserts no unexpected edge to `__end__` is produced from the last nonterminal step source. ### Issue Closes #4394 --------- Co-authored-by: Sydney Runkle Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> --- libs/langgraph/langgraph/pregel/_draw.py | 3 +- .../tests/__snapshots__/test_pregel.ambr | 93 +++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 40 ++++++++ 3 files changed, 135 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/_draw.py b/libs/langgraph/langgraph/pregel/_draw.py index b8ae73389..af6d03427 100644 --- a/libs/langgraph/langgraph/pregel/_draw.py +++ b/libs/langgraph/langgraph/pregel/_draw.py @@ -215,10 +215,11 @@ def draw_graph( termini = {d for _, d, _, _ in edges if d != END}.difference( s for s, _, _, _ in edges ) + end_edge_exists = any(d == END for _, d, _, _ in edges) if termini: for src in sorted(termini): add_edge(graph, src, END) - elif len(step_sources) == 1: + elif len(step_sources) == 1 and not end_edge_exists: for src in sorted(step_sources): add_edge(graph, src, END, conditional=True) # replace subgraphs diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 5167d8924..b7dfbcd0e 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -795,6 +795,99 @@ ''' # --- +# name: test_get_graph_nonterminal_last_step_source + ''' + { + "edges": [ + { + "source": "__start__", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "human" + }, + { + "conditional": true, + "source": "chatbot", + "target": "tools" + }, + { + "conditional": true, + "source": "human", + "target": "__end__" + }, + { + "conditional": true, + "source": "human", + "target": "chatbot" + }, + { + "source": "tools", + "target": "chatbot" + } + ], + "nodes": [ + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "__start__" + }, + "id": "__start__", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "chatbot" + }, + "id": "chatbot", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "id": "tools", + "type": "runnable" + }, + { + "data": { + "id": [ + "langgraph", + "_internal", + "_runnable", + "RunnableCallable" + ], + "name": "human" + }, + "id": "human", + "type": "runnable" + }, + { + "id": "__end__" + } + ] + } + ''' +# --- # name: test_repeat_condition ''' graph TD; diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index cbf6d03b6..7cc67bd35 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8343,6 +8343,46 @@ def test_subgraph_streaming_sync() -> None: assert result["num_chunks"] == 9 +def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: list[str] + + def chatbot_node(state: State) -> State: + return {"messages": state["messages"] + ["chatbot"]} + + def tools_node(state: State) -> State: + return {"messages": state["messages"] + ["tools"]} + + def human_node(state: State) -> State: + return {"messages": state["messages"] + ["human"]} + + def tools_condition(_: State) -> str: + return "tools" + + def end_condition(_: State) -> str: + return "chatbot" + + workflow = StateGraph(State) + workflow.add_node("chatbot", chatbot_node) + workflow.add_node("tools", tools_node) + workflow.add_node("human", human_node) + + workflow.add_edge(START, "human") + workflow.add_edge("tools", "chatbot") + + workflow.add_conditional_edges( + "chatbot", tools_condition, {"tools": "tools", "human": "human"} + ) + workflow.add_conditional_edges( + "human", end_condition, {"chatbot": "chatbot", END: END} + ) + + app = workflow.compile() + graph = app.get_graph() + graph_json = graph.to_json() + + assert json.dumps(graph_json, indent=2, sort_keys=True) == snapshot + def test_null_resume_disallowed_with_multiple_interrupts( sync_checkpointer: BaseCheckpointSaver, ) -> None: