lib: Add support for multiple interrupts per node

libs/langgraph/langgraph/constants.py

@@ -1,4 +1,4 @@
 import sys
 from os import getenv
 from types import MappingProxyType
 from typing import Any, Literal, Mapping, cast
@@ -72,9 +72,11 @@
 CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
 # holds the current checkpoint_ns, "" for root graph
 CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
-# callback to be called when a node is finished
-CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
 # holds the value that "answers" an interrupt() call
+CONFIG_KEY_WRITES = sys.intern("__pregel_writes")
+# read-only list of existing task writes
+CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad")
+# holds a mutable dict for temporary storage scoped to the current task
 
 # --- Other constants ---
 PUSH = sys.intern("__pregel_push")

libs/langgraph/langgraph/pregel/algo.py

@@ -37,13 +37,13 @@
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_CHECKPOINTER,
     CONFIG_KEY_READ,
-    CONFIG_KEY_RESUME_VALUE,
+    CONFIG_KEY_SCRATCHPAD,
     CONFIG_KEY_SEND,
     CONFIG_KEY_STORE,
     CONFIG_KEY_TASK_ID,
+    CONFIG_KEY_WRITES,
     EMPTY_SEQ,
     INTERRUPT,
-    MISSING,
     NO_WRITES,
     NS_END,
     NS_SEP,
@@ -589,14 +589,13 @@ def prepare_single_task(
                             },
                             CONFIG_KEY_CHECKPOINT_ID: None,
                             CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                            CONFIG_KEY_RESUME_VALUE: next(
-                                (
-                                    v
-                                    for tid, c, v in pending_writes
-                                    if tid in (NULL_TASK_ID, task_id) and c == RESUME
-                                ),
-                                configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
-                            ),
+                            CONFIG_KEY_WRITES: [
+                                w
+                                for w in pending_writes
+                                + configurable.get(CONFIG_KEY_WRITES, [])
+                                if w[0] in (NULL_TASK_ID, task_id)
+                            ],
+                            CONFIG_KEY_SCRATCHPAD: {},
                         },
                     ),
                     triggers,
@@ -713,15 +712,13 @@ def prepare_single_task(
                                 },
                                 CONFIG_KEY_CHECKPOINT_ID: None,
                                 CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                                CONFIG_KEY_RESUME_VALUE: next(
-                                    (
-                                        v
-                                        for tid, c, v in pending_writes
-                                        if tid in (NULL_TASK_ID, task_id)
-                                        and c == RESUME
-                                    ),
-                                    configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
-                                ),
+                                CONFIG_KEY_WRITES: [
+                                    w
+                                    for w in pending_writes
+                                    + configurable.get(CONFIG_KEY_WRITES, [])
+                                    if w[0] in (NULL_TASK_ID, task_id)
+                                ],
+                                CONFIG_KEY_SCRATCHPAD: {},
                             },
                         ),
                         triggers,

libs/langgraph/langgraph/pregel/io.py

@@ -4,6 +4,7 @@
 from langchain_core.runnables.utils import AddableDict
 
 from langgraph.channels.base import BaseChannel, EmptyChannelError
+from langgraph.checkpoint.base import PendingWrite
 from langgraph.constants import (
     EMPTY_SEQ,
     ERROR,
@@ -66,7 +67,7 @@ def read_channels(
 
 
 def map_command(
-    cmd: Command,
+    cmd: Command, pending_writes: list[PendingWrite]
 ) -> Iterator[tuple[str, str, Any]]:
     """Map input chunk to a sequence of pending writes in the form (channel, value)."""
     if cmd.graph == Command.PARENT:
@@ -85,7 +86,11 @@ def map_command(
     if cmd.resume:
         if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
             for tid, resume in cmd.resume.items():
-                yield (tid, RESUME, resume)
+                existing: list[Any] = next(
+                    (w[2] for w in pending_writes if w[0] == tid and w[1] == RESUME), []
+                )
+                existing.append(resume)
+                yield (tid, RESUME, existing)
         else:
             yield (NULL_TASK_ID, RESUME, cmd.resume)
     if cmd.update:

libs/langgraph/langgraph/pregel/loop.py

@@ -26,6 +26,7 @@
 
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -263,8 +264,28 @@ def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
         """Put writes for a task, to be read by the next tick."""
         if not writes:
             return
+        # deduplicate writes to special channels, last write wins
+        if all(w[0] in WRITES_IDX_MAP for w in writes):
+            writes = list({w[0]: w for w in writes}.values())
         # save writes
-        self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
+        for c, v in writes:
+            if (
+                c in WRITES_IDX_MAP
+                and (
+                    idx := next(
+                        (
+                            i
+                            for i, w in enumerate(self.checkpoint_pending_writes)
+                            if w[0] == task_id and w[1] == c
+                        ),
+                        None,
+                    )
+                )
+                is not None
+            ):
+                self.checkpoint_pending_writes[idx] = (task_id, c, v)
+            else:
+                self.checkpoint_pending_writes.append((task_id, c, v))
         if self.checkpointer_put_writes is not None:
             self.submit(
                 self.checkpointer_put_writes,
@@ -536,7 +557,7 @@ def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
         elif isinstance(self.input, Command):
             writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
             # group writes by task ID
-            for tid, c, v in map_command(self.input):
+            for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
                 writes[tid].append((c, v))
             if not writes:
                 raise EmptyInputError("Received empty Command input")

libs/langgraph/langgraph/pregel/runner.py

@@ -21,6 +21,7 @@
     INTERRUPT,
     NO_WRITES,
     PUSH,
+    RESUME,
     TAG_HIDDEN,
 )
 from langgraph.errors import GraphBubbleUp, GraphInterrupt
@@ -297,6 +298,8 @@ def commit(
             if isinstance(exception, GraphInterrupt):
                 # save interrupt to checkpointer
                 if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
+                    if resumes := [w for w in task.writes if w[0] == RESUME]:
+                        interrupts.extend(resumes)
                     self.put_writes(task.id, interrupts)
             elif isinstance(exception, GraphBubbleUp):
                 raise exception

libs/langgraph/langgraph/types.py

@@ -13,6 +13,7 @@
     Optional,
     Sequence,
     Type,
+    TypedDict,
     TypeVar,
     Union,
     cast,
@@ -21,11 +22,16 @@
 from langchain_core.runnables import Runnable, RunnableConfig
 from typing_extensions import Self
 
-from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    CheckpointMetadata,
+    PendingWrite,
+)
 
 if TYPE_CHECKING:
     from langgraph.store.base import BaseStore
 
+
 All = Literal["*"]
 """Special value to indicate that graph should interrupt on all nodes."""
 
@@ -300,26 +306,60 @@ def __init__(
         self.stop = stop
 
 
+class PregelScratchpad(TypedDict, total=False):
+    interrupt_counter: int
+    used_null_resume: bool
+    resume: list[Any]
+
+
 def interrupt(value: Any) -> Any:
     from langgraph.constants import (
         CONFIG_KEY_CHECKPOINT_NS,
-        CONFIG_KEY_RESUME_VALUE,
-        MISSING,
+        CONFIG_KEY_SCRATCHPAD,
+        CONFIG_KEY_SEND,
+        CONFIG_KEY_TASK_ID,
+        CONFIG_KEY_WRITES,
         NS_SEP,
+        NULL_TASK_ID,
+        RESUME,
     )
     from langgraph.errors import GraphInterrupt
     from langgraph.utils.config import get_configurable
 
     conf = get_configurable()
-    if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING:
-        return resume
+    # track interrupt index
+    scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
+    if "interrupt_counter" not in scratchpad:
+        scratchpad["interrupt_counter"] = 0
     else:
-        raise GraphInterrupt(
-            (
-                Interrupt(
-                    value=value,
-                    resumable=True,
-                    ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
-                ),
-            )
+        scratchpad["interrupt_counter"] += 1
+    idx = scratchpad["interrupt_counter"]
+    # find previous resume values
+    task_id = conf[CONFIG_KEY_TASK_ID]
+    writes: list[PendingWrite] = conf[CONFIG_KEY_WRITES]
+    scratchpad.setdefault(
+        "resume", next((w[2] for w in writes if w[0] == task_id and w[1] == RESUME), [])
+    )
+    if scratchpad["resume"]:
+        if idx < len(scratchpad["resume"]):
+            return scratchpad["resume"][idx]
+    # find current resume value
+    if not scratchpad.get("used_null_resume"):
+        scratchpad["used_null_resume"] = True
+        for tid, c, v in sorted(writes, key=lambda x: x[0], reverse=True):
+            if tid == NULL_TASK_ID and c == RESUME:
+                assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
+                scratchpad["resume"].append(v)
+                print("saving:", scratchpad["resume"])
+                conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
+                return v
+    # no resume value found
+    raise GraphInterrupt(
+        (
+            Interrupt(
+                value=value,
+                resumable=True,
+                ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP),
+            ),
         )
+    )