Add refcount test

libs/langgraph/tests/test_pregel.py

@@ -1,5 +1,6 @@
 import enum
 import functools
+import gc
 import json
 import logging
 import operator
@@ -62,7 +63,9 @@ from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.loop import SyncPregelLoop
 from langgraph.pregel.retry import RetryPolicy
+from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     Command,
@@ -7613,3 +7616,28 @@ def test_parallel_interrupts_double(
 
     assert invokes == 5
     assert len(events) == 5
+
+
+def test_pregel_loop_refcount():
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    graph_builder = StateGraph(State)
+
+    def chatbot(state: State):
+        return {"messages": [("ai", "HIYA")]}
+
+    graph_builder.add_node("chatbot", chatbot)
+    graph_builder.set_entry_point("chatbot")
+    graph_builder.set_finish_point("chatbot")
+    graph = graph_builder.compile()
+
+    for _ in range(5):
+        graph.invoke({"messages": [{"role": "user", "content": "hi"}]})
+        assert (
+            len([obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)])
+            == 0
+        )
+        assert (
+            len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0
+        )

libs/langgraph/tests/test_pregel_async.py

@@ -1,5 +1,6 @@
 import asyncio
 import functools
+import gc
 import logging
 import operator
 import random
@@ -52,7 +53,9 @@ from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessagesState, add_messages
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.loop import AsyncPregelLoop
 from langgraph.pregel.retry import RetryPolicy
+from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     Command,
@@ -7838,3 +7841,28 @@ async def test_handles_multiple_interrupts_from_tasks() -> None:
     assert len(result) == 2
     assert result[0] == "Added James!"
     assert result[1] == "Added Will!"
+
+
+async def test_pregel_loop_refcount():
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    graph_builder = StateGraph(State)
+
+    async def chatbot(state: State):
+        return {"messages": [("ai", "HIYA")]}
+
+    graph_builder.add_node("chatbot", chatbot)
+    graph_builder.set_entry_point("chatbot")
+    graph_builder.set_finish_point("chatbot")
+    graph = graph_builder.compile()
+
+    for _ in range(5):
+        await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]})
+        assert (
+            len([obj for obj in gc.get_objects() if isinstance(obj, AsyncPregelLoop)])
+            == 0
+        )
+        assert (
+            len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0
+        )