From 54e18445fc27c2ac63eeb974a281a4acf36743fa Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Mon, 17 Mar 2025 19:53:03 -0700
Subject: [PATCH] Fix concurrency issue in PregelScratchpad.consume_null_resume

- Need to use a single operation to check if present and remove item from list
- This doesn't fix the separate issue that parallel tasks claiming a single interrupt value have somewhat undefined behavior (in the sense that they will race to be the first to take it). That will be fixed in a future PR
---
 libs/langgraph/langgraph/pregel/algo.py | 18 +++++++++++++-----
 libs/langgraph/langgraph/pregel/loop.py |  9 ++++-----
 libs/langgraph/langgraph/types.py       | 15 +++------------
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index 03b2af6f4..e08f0eff4 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -1,4 +1,3 @@
-import functools
 import itertools
 import sys
 from collections import defaultdict, deque
@@ -768,6 +767,18 @@ def _scratchpad(
     null_resume_write = next(
         (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
     )
+
+    def get_null_resume(consume: bool = False) -> Any:
+        if null_resume_write is None:
+            return None
+        if consume:
+            try:
+                pending_writes.remove(null_resume_write)
+                return null_resume_write[2]
+            except ValueError:
+                return None
+        return null_resume_write[2]
+
     # using itertools.count as an atomic counter (+= 1 is not thread-safe)
     return PregelScratchpad(
         # call
@@ -777,10 +788,7 @@ def _scratchpad(
         resume=next(
             (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
         ),
-        null_resume=null_resume_write[2] if null_resume_write is not None else None,
-        _consume_null_resume=functools.partial(pending_writes.remove, null_resume_write)
-        if null_resume_write is not None
-        else lambda: None,
+        get_null_resume=get_null_resume,
         # subgraph
         subgraph_counter=itertools.count(0).__next__,
     )
diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
index dec608b62..6ef009058 100644
--- a/libs/langgraph/langgraph/pregel/loop.py
+++ b/libs/langgraph/langgraph/pregel/loop.py
@@ -591,11 +591,10 @@ class PregelLoop(LoopProtocol):
         if scratchpad := cast(
             Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD)
         ):
-            if (
-                isinstance(scratchpad, PregelScratchpad)
-                and scratchpad.null_resume is not None
-            ):
-                self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad.null_resume)])
+            if isinstance(scratchpad, PregelScratchpad):
+                null_resume = scratchpad.get_null_resume(False)
+                if null_resume is not None:
+                    self.put_writes(NULL_TASK_ID, [(RESUME, null_resume)])
         # map command to writes
         if isinstance(self.input, Command):
             if self.input.resume is not None and not self.checkpointer:
diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py
index c2b878a4b..4c3e84c57 100644
--- a/libs/langgraph/langgraph/types.py
+++ b/libs/langgraph/langgraph/types.py
@@ -351,20 +351,11 @@ class PregelScratchpad:
     call_counter: Callable[[], int]
     # interrupt
     interrupt_counter: Callable[[], int]
+    get_null_resume: Callable[[bool], Any]
     resume: list[Any]
-    null_resume: Optional[Any]
-    _consume_null_resume: Callable[[], None]
     # subgraph
     subgraph_counter: Callable[[], int]
 
-    def consume_null_resume(self) -> Any:
-        if self.null_resume is not None:
-            value = self.null_resume
-            self._consume_null_resume()
-            self.null_resume = None
-            return value
-        raise ValueError("No null resume to consume")
-
 
 def interrupt(value: Any) -> Any:
     """Interrupt the graph with a resumable exception from within a node.
@@ -480,9 +471,9 @@ def interrupt(value: Any) -> Any:
         if idx < len(scratchpad.resume):
             return scratchpad.resume[idx]
     # find current resume value
-    if scratchpad.null_resume is not None:
+    v = scratchpad.get_null_resume(True)
+    if v is not None:
         assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
-        v = scratchpad.consume_null_resume()
         scratchpad.resume.append(v)
         conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
         return v