langgraph: add test for interrupting multiple times from a task (#3148)

libs/langgraph/langgraph/pregel/algo.py

@@ -1,3 +1,4 @@
+import functools
 import sys
 from collections import defaultdict, deque
 from functools import partial
@@ -46,7 +47,6 @@ from langgraph.constants import (
     EMPTY_SEQ,
     ERROR,
     INTERRUPT,
-    MISSING,
     NO_WRITES,
     NS_END,
     NS_SEP,
@@ -324,7 +324,7 @@ def apply_writes(
 @overload
 def prepare_next_tasks(
     checkpoint: Checkpoint,
-    pending_writes: Sequence[PendingWrite],
+    pending_writes: list[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -341,7 +341,7 @@ def prepare_next_tasks(
 @overload
 def prepare_next_tasks(
     checkpoint: Checkpoint,
-    pending_writes: Sequence[PendingWrite],
+    pending_writes: list[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -357,7 +357,7 @@ def prepare_next_tasks(
 
 def prepare_next_tasks(
     checkpoint: Checkpoint,
-    pending_writes: Sequence[PendingWrite],
+    pending_writes: list[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -418,7 +418,7 @@ def prepare_single_task(
     task_id_checksum: Optional[str],
     *,
     checkpoint: Checkpoint,
-    pending_writes: Sequence[PendingWrite],
+    pending_writes: list[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
@@ -761,9 +761,13 @@ def prepare_single_task(
 
 
 def _scratchpad(
-    pending_writes: Sequence[PendingWrite],
+    pending_writes: list[PendingWrite],
     task_id: str,
 ) -> PregelScratchpad:
+    null_resume_write = next(
+        (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
+    )
+
     return PregelScratchpad(
         # call
         call_counter=0,
@@ -772,10 +776,10 @@ def _scratchpad(
         resume=next(
             (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
         ),
-        null_resume=next(
-            (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME),
-            MISSING,
-        ),
+        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,
         # subgraph
         subgraph_counter=0,
     )

libs/langgraph/langgraph/pregel/io.py

@@ -89,7 +89,7 @@ def map_command(
                 raise TypeError(
                     f"In Command.goto, expected Send/str, got {type(send).__name__}"
                 )
-    if cmd.resume:
+    if cmd.resume is not None:
         if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
             for tid, resume in cmd.resume.items():
                 existing: list[Any] = next(

libs/langgraph/langgraph/pregel/loop.py

@@ -54,7 +54,6 @@ from langgraph.constants import (
     ERROR,
     INPUT,
     INTERRUPT,
-    MISSING,
     NS_SEP,
     NULL_TASK_ID,
     PUSH,
@@ -229,20 +228,22 @@ class PregelLoop(LoopProtocol):
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
         scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
-        if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None:
-            if scratchpad["subgraph_counter"]:
+        if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance(
+            scratchpad, PregelScratchpad
+        ):
+            if scratchpad.subgraph_counter:
                 self.config = patch_configurable(
                     self.config,
                     {
                         CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
                             (
                                 config[CONF][CONFIG_KEY_CHECKPOINT_NS],
-                                str(scratchpad["subgraph_counter"]),
+                                str(scratchpad.subgraph_counter),
                             )
                         )
                     },
                 )
-            scratchpad["subgraph_counter"] += 1
+            scratchpad.subgraph_counter += 1
         if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
             self.config = patch_configurable(
                 self.config,
@@ -563,9 +564,14 @@ class PregelLoop(LoopProtocol):
         )
 
         # take resume value from parent
-        if scratchpad := configurable.get(CONFIG_KEY_SCRATCHPAD):
-            if scratchpad["null_resume"] is not MISSING:
-                self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad["null_resume"])])
+        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)])
         # map command to writes
         if isinstance(self.input, Command):
             if self.input.resume is not None and not self.checkpointer:

libs/langgraph/langgraph/pregel/runner.py

@@ -39,7 +39,7 @@ from langgraph.errors import GraphBubbleUp, GraphInterrupt
 from langgraph.pregel.algo import Call
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
-from langgraph.types import PregelExecutableTask, RetryPolicy
+from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy
 from langgraph.utils.future import chain_future
 
 F = TypeVar("F", concurrent.futures.Future, asyncio.Future)
@@ -135,8 +135,7 @@ class PregelRunner:
                 return task.config[CONF][CONFIG_KEY_SEND](writes)
 
             # schedule PUSH tasks, collect futures
-            scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
-            scratchpad.setdefault("call_counter", 0)
+            scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
             rtn: dict[int, Optional[concurrent.futures.Future]] = {}
             for idx, w in enumerate(writes):
                 # bail if not a PUSH write
@@ -144,8 +143,8 @@ class PregelRunner:
                     continue
                 # schedule the next task, if the callback returns one
                 wcall = calls[idx] if calls else None
-                cnt = scratchpad["call_counter"]
-                scratchpad["call_counter"] += 1
+                cnt = scratchpad.call_counter
+                scratchpad.call_counter += 1
                 if next_task := self.schedule_task(task, cnt, wcall):
                     if fut := next(
                         (
@@ -324,8 +323,7 @@ class PregelRunner:
                 return task.config[CONF][CONFIG_KEY_SEND](writes)
 
             # schedule PUSH tasks, collect futures
-            scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
-            scratchpad.setdefault("call_counter", 0)
+            scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
             rtn: dict[int, Optional[asyncio.Future]] = {}
             for idx, w in enumerate(writes):
                 # bail if not a PUSH write
@@ -333,8 +331,8 @@ class PregelRunner:
                     continue
                 # schedule the next task, if the callback returns one
                 wcall = calls[idx] if calls is not None else None
-                cnt = scratchpad["call_counter"]
-                scratchpad["call_counter"] += 1
+                cnt = scratchpad.call_counter
+                scratchpad.call_counter += 1
                 if next_task := self.schedule_task(task, cnt, wcall):
                     # if the parent task was retried,
                     # the next task might already be running

libs/langgraph/langgraph/types.py

@@ -19,7 +19,7 @@ from typing import (
 )
 
 from langchain_core.runnables import Runnable, RunnableConfig
-from typing_extensions import Self, TypedDict
+from typing_extensions import Self
 
 from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
 
@@ -339,16 +339,26 @@ class LoopProtocol:
         self.stop = stop
 
 
-class PregelScratchpad(TypedDict):
+@dataclasses.dataclass(**{**_DC_KWARGS, "frozen": False})
+class PregelScratchpad:
     # call
     call_counter: int
     # interrupt
     interrupt_counter: int
     resume: list[Any]
-    null_resume: Any
+    null_resume: Optional[Any]
+    _consume_null_resume: Callable[[], None]
     # subgraph
     subgraph_counter: 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.
@@ -449,7 +459,6 @@ def interrupt(value: Any) -> Any:
         CONFIG_KEY_CHECKPOINT_NS,
         CONFIG_KEY_SCRATCHPAD,
         CONFIG_KEY_SEND,
-        MISSING,
         NS_SEP,
         RESUME,
     )
@@ -459,19 +468,19 @@ def interrupt(value: Any) -> Any:
     conf = get_config()["configurable"]
     # track interrupt index
     scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
-    scratchpad["interrupt_counter"] += 1
-    idx = scratchpad["interrupt_counter"]
+    scratchpad.interrupt_counter += 1
+    idx = scratchpad.interrupt_counter
     # find previous resume values
-    if scratchpad["resume"]:
-        if idx < len(scratchpad["resume"]):
-            return scratchpad["resume"][idx]
+    if scratchpad.resume:
+        if idx < len(scratchpad.resume):
+            return scratchpad.resume[idx]
     # find current resume value
-    if scratchpad["null_resume"] is not MISSING:
-        assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
-        v = scratchpad["null_resume"]
-        scratchpad["null_resume"] = MISSING
-        scratchpad["resume"].append(v)
-        conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
+    if scratchpad.null_resume is not None:
+        assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
+        print("consume null resume", scratchpad.null_resume)
+        v = scratchpad.consume_null_resume()
+        scratchpad.resume.append(v)
+        conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
         return v
     # no resume value found
     raise GraphInterrupt(

libs/langgraph/tests/test_pregel.py

@@ -5082,11 +5082,27 @@ def test_interrupt_task_functional(
 
     config = {"configurable": {"thread_id": "1"}}
     # First run, interrupted at bar
-    graph.invoke({"a": ""}, config)
+    assert not graph.invoke({"a": ""}, config)
     # Resume with an answer
     res = graph.invoke(Command(resume="bar"), config)
     assert res == {"a": "foobar"}
 
+    # Test that we can interrupt the same task multiple times
+    config = {"configurable": {"thread_id": "2"}}
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(inputs: dict) -> dict:
+        foo_result = foo(inputs).result()
+        bar_result = bar(foo_result).result()
+        baz_result = bar(bar_result).result()
+        return baz_result
+
+    # First run, interrupted at bar
+    assert not graph.invoke({"a": ""}, config)
+    # Provide resumes
+    assert not graph.invoke(Command(resume="bar"), config)
+    assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"}
+
 
 def test_root_mixed_return() -> None:
     def my_node(state: list[str]):

libs/langgraph/tests/test_pregel_async.py

@@ -6408,6 +6408,22 @@ async def test_interrupt_task_functional(checkpointer_name: str) -> None:
         res = await graph.ainvoke(Command(resume="bar"), config)
         assert res == {"a": "foobar"}
 
+        # Test that we can interrupt the same task multiple times
+        config = {"configurable": {"thread_id": "2"}}
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(inputs: dict) -> dict:
+            foo_result = await foo(inputs)
+            bar_result = await bar(foo_result)
+            baz_result = await bar(bar_result)
+            return baz_result
+
+        # First run, interrupted at bar
+        assert not await graph.ainvoke({"a": ""}, config)
+        # Provide resumes
+        assert not await graph.ainvoke(Command(resume="bar"), config)
+        assert await graph.ainvoke(Command(resume="baz"), config) == {"a": "foobarbaz"}
+
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -13,10 +13,8 @@ from typing_extensions import Self
 
 import langgraph.scheduler.kafka.serde as serde
 from langgraph.constants import (
-    CONF,
     CONFIG_KEY_DEDUPE_TASKS,
     CONFIG_KEY_ENSURE_LATEST,
-    CONFIG_KEY_SCRATCHPAD,
     INTERRUPT,
     SCHEDULED,
 )
@@ -178,8 +176,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
                             CONFIG_KEY_ENSURE_LATEST: True,
                         },
                     )
-                    if CONFIG_KEY_SCRATCHPAD in config[CONF]:
-                        config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
                     # send messages to executor
                     futures = await asyncio.gather(
                         *(
@@ -366,8 +362,6 @@ class KafkaOrchestrator(AbstractContextManager):
                             CONFIG_KEY_ENSURE_LATEST: True,
                         },
                     )
-                    if CONFIG_KEY_SCRATCHPAD in config[CONF]:
-                        config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
                     # send messages to executor
                     futures = [
                         self.producer.send(