From e9e9a96a0d1bb2ae3dbe77bd0092c19512003d8f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 31 Mar 2025 14:50:10 -0700
Subject: [PATCH] Reduce perf impact of pregel scratchpad creation

- make scratchpad class frozen now that its members are never reassigned
- replace next(gen expr) with for-loop to avoid allocating generator objects
---
 libs/langgraph/langgraph/pregel/algo.py | 33 +++++++++++++++++++------
 libs/langgraph/langgraph/types.py       |  2 +-
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 355ccdcff..4d9157da1 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -869,11 +869,30 @@ def _scratchpad(
     pending_writes: list[PendingWrite],
     task_id: str,
 ) -> PregelScratchpad:
-    # None cannot be used as a resume value, because it would be difficult to
-    # distinguish from missing when used over http
-    null_resume_write = next(
-        (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
-    )
+    if len(pending_writes) > 0:
+        # find global resume value
+        for w in pending_writes:
+            if w[0] == NULL_TASK_ID and w[1] == RESUME:
+                null_resume_write = w
+                break
+        else:
+            # None cannot be used as a resume value, because it would be difficult to
+            # distinguish from missing when used over http
+            null_resume_write = None
+        # find task-specific resume value
+        for w in pending_writes:
+            if w[0] == task_id and w[1] == RESUME:
+                task_resume_write = w[2]
+                if not isinstance(task_resume_write, list):
+                    task_resume_write = [task_resume_write]
+                break
+        else:
+            task_resume_write = []
+        # clear var
+        del w
+    else:
+        null_resume_write = None
+        task_resume_write = []
 
     def get_null_resume(consume: bool = False) -> Any:
         if null_resume_write is None:
@@ -894,9 +913,7 @@ def _scratchpad(
         call_counter=itertools.count(0).__next__,
         # interrupt
         interrupt_counter=itertools.count(0).__next__,
-        resume=next(
-            (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
-        ),
+        resume=task_resume_write,
         get_null_resume=get_null_resume,
         # subgraph
         subgraph_counter=itertools.count(0).__next__,
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index 6fb0b9b4a..a24ee1ea5 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/libs/langgraph/langgraph/types.py
@@ -369,7 +369,7 @@ class LoopProtocol:
         self.stop = stop
 
 
-@dataclasses.dataclass(**{**_DC_KWARGS, "frozen": False})
+@dataclasses.dataclass(**_DC_KWARGS)
 class PregelScratchpad:
     # call
     call_counter: Callable[[], int]