From 6a2d20fd5b2f1c8d5c47b6445864da30548bbf5a Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 19 Mar 2025 13:02:07 -0700
Subject: [PATCH] Reference to PregelExecutableTask

---
 libs/langgraph/langgraph/pregel/loop.py   |  5 ++++-
 libs/langgraph/langgraph/pregel/runner.py | 26 +++++++++++------------
 libs/langgraph/langgraph/types.py         |  9 +++++++-
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index edd69db01..807be5f4d 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -1,6 +1,7 @@
 import asyncio
 import binascii
 import concurrent.futures
+import dataclasses
 from collections import defaultdict, deque
 from contextlib import AsyncExitStack, ExitStack
 from inspect import signature
@@ -571,7 +572,9 @@ class PregelLoop(LoopProtocol):
                         self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(),
                         default=None,
                     ):
-                        self.tasks[tid] = task._replace(scheduled=True)
+                        self.tasks[tid] = PregelExecutableTask(
+                            **dataclasses.asdict(task) | {"scheduled": True}
+                        )
                 else:
                     task.writes.append((k, v))
 
diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py
index 67f9c7b3e..0c312e145 100644
--- a/libs/langgraph/langgraph/pregel/runner.py
+++ b/libs/langgraph/langgraph/pregel/runner.py
@@ -149,7 +149,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _call,
-                            t,
+                            weakref.ref(t),
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
@@ -185,7 +185,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _call,
-                            t,
+                            weakref.ref(t),
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
@@ -263,7 +263,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
-                            t,
+                            weakref.ref(t),
                             stream=self.use_astream,
                             retry=retry_policy,
                             futures=weakref.ref(futures),
@@ -304,7 +304,7 @@ class PregelRunner:
                         configurable={
                             CONFIG_KEY_CALL: partial(
                                 _acall,
-                                t,
+                                weakref.ref(t),
                                 retry=retry_policy,
                                 stream=self.use_astream,
                                 futures=weakref.ref(futures),
@@ -469,7 +469,7 @@ def _panic_or_proceed(
 
 
 def _call(
-    task: PregelExecutableTask,
+    task: weakref.ref[PregelExecutableTask],
     func: Callable[[Any], Union[Awaitable[Any], Any]],
     input: Any,
     *,
@@ -489,10 +489,10 @@ def _call(
 
     fut: Optional[concurrent.futures.Future] = None
     # schedule PUSH tasks, collect futures
-    scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
+    scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]
     # schedule the next task, if the callback returns one
     if next_task := schedule_task()(  # type: ignore[misc]
-        task,
+        task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
         Call(func, input, retry=retry, callbacks=callbacks),
     ):
@@ -528,7 +528,7 @@ def _call(
                 configurable={
                     CONFIG_KEY_CALL: partial(
                         _call,
-                        next_task,
+                        weakref.ref(next_task),
                         futures=futures,
                         retry=retry,
                         callbacks=callbacks,
@@ -550,7 +550,7 @@ def _call(
 
 
 def _acall(
-    task: PregelExecutableTask,
+    task: weakref.ref[PregelExecutableTask],
     func: Callable[[Any], Union[Awaitable[Any], Any]],
     input: Any,
     *,
@@ -570,10 +570,10 @@ def _acall(
 ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
     fut: Optional[asyncio.Future] = None
     # schedule PUSH tasks, collect futures
-    scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
+    scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]
     # schedule the next task, if the callback returns one
     if next_task := schedule_task()(  # type: ignore[misc]
-        task,
+        task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
         Call(func, input, retry=retry, callbacks=callbacks),
     ):
@@ -614,7 +614,7 @@ def _acall(
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
-                            next_task,
+                            weakref.ref(next_task),
                             stream=stream,
                             futures=futures,
                             schedule_task=schedule_task,
@@ -623,7 +623,7 @@ def _acall(
                             reraise=reraise,
                         ),
                     },
-                    __name__=task.name,
+                    __name__=task().name,  # type: ignore[union-attr]
                     __cancel_on_exit__=True,
                     __reraise_on_exit__=reraise,
                     # starting a new task in the next tick ensures
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index 00339119d..0a3e9f205 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/libs/langgraph/langgraph/types.py
@@ -148,7 +148,14 @@ class PregelTask(NamedTuple):
     result: Optional[Any] = None
 
 
-class PregelExecutableTask(NamedTuple):
+if sys.version_info > (3, 11):
+    _T_DC_KWARGS = {"weakref_slot": True, "slots": True, "frozen": True}
+else:
+    _T_DC_KWARGS = {"frozen": True}
+
+
+@dataclasses.dataclass(**_T_DC_KWARGS)
+class PregelExecutableTask:
     name: str
     input: Any
     proc: Runnable