From a9ea0cd28a328d5f109f2dedcee32910a372fad9 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 7 May 2025 09:12:12 -0700
Subject: [PATCH] Add a limit to Pregel.draw (#4575)

Co-authored-by: vbarda <vadym@langchain.dev>
---
 libs/langgraph/langgraph/pregel/draw.py       |  5 +-
 .../tests/__snapshots__/test_pregel.ambr      | 79 +++++++++++++++++++
 libs/langgraph/tests/test_pregel.py           | 23 ++++++
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py
index 44b7c1995..8df6dd458 100644
--- a/libs/langgraph/langgraph/pregel/draw.py
+++ b/libs/langgraph/langgraph/pregel/draw.py
@@ -34,6 +34,7 @@ def draw_graph(
     trigger_to_nodes: Optional[Mapping[str, Sequence[str]]],
     checkpointer: Checkpointer,
     subgraphs: dict[str, Graph],
+    limit: int = 250,
 ) -> Graph:
     """Get the graph for this Pregel instance.
 
@@ -97,7 +98,9 @@ def draw_graph(
         )
         start_tasks = tasks
         # run the pregel loop
-        while tasks:
+        for _ in range(limit):
+            if not tasks:
+                break
             conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
             # run task writers
             for task in tasks.values():
diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr
index 69303056d..46955d25c 100644
--- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr
+++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr
@@ -317,6 +317,85 @@
   
   '''
 # ---
+# name: test_get_graph_loop
+  '''
+  {
+    "nodes": [
+      {
+        "id": "__start__",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langchain",
+            "schema",
+            "runnable",
+            "RunnablePassthrough"
+          ],
+          "name": "__start__"
+        }
+      },
+      {
+        "id": "human",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langgraph",
+            "utils",
+            "runnable",
+            "RunnableCallable"
+          ],
+          "name": "human"
+        }
+      },
+      {
+        "id": "agent",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langgraph",
+            "utils",
+            "runnable",
+            "RunnableCallable"
+          ],
+          "name": "agent"
+        }
+      },
+      {
+        "id": "__end__"
+      }
+    ],
+    "edges": [
+      {
+        "source": "__start__",
+        "target": "human"
+      },
+      {
+        "source": "agent",
+        "target": "human"
+      },
+      {
+        "source": "human",
+        "target": "agent"
+      },
+      {
+        "source": "human",
+        "target": "__end__",
+        "conditional": true
+      }
+    ]
+  }
+  '''
+# ---
+# name: test_get_graph_loop.1
+  '''
+  graph TD;
+  	__start__ --> human;
+  	agent --> human;
+  	human --> agent;
+  	human -.-> __end__;
+  
+  '''
+# ---
 # name: test_in_one_fan_out_state_graph_waiting_edge[memory]
   '''
   graph TD;
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 97e6786a9..a8fd2b5f3 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -8031,3 +8031,26 @@ def test_migration_graph(snapshot: SnapshotAssertion) -> None:
     app = migration_graph.compile()
 
     assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+
+def test_get_graph_loop(snapshot: SnapshotAssertion) -> None:
+    class State(TypedDict):
+        foo: str
+
+    def human_node(state: State) -> State:
+        value = interrupt()
+        return {"foo": value}
+
+    def agent_node(state: State) -> State:
+        return {"foo": "Hi " + state["foo"]}
+
+    workflow = StateGraph(State)
+    workflow.add_node("human", human_node)
+    workflow.add_node("agent", agent_node)
+    workflow.add_edge(START, "human")
+    workflow.add_edge("human", "agent")
+    workflow.add_edge("agent", "human")
+
+    app = workflow.compile()
+    assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot