feat(langgraph): prevent arbitrary resumes w/ multiple pending interrupts (#6108)

The idea here is that we don't want to allow resuming a graph w/ an
arbitrary resume value if there are multiple interrupts in the queue,
because the order in which interrupts enter the queue is not
deterministic. We want to instead enforce that each resume value is
mapped to an interrupt id.

Instead, when multiple interrupts are present, a user should invoke w/ a
resume map, mapping interrupt id -> resume value.

The logic was more complex than expected because there are 2 copies of
an interrupt in `checkpoint_pending_writes` for the cases w/ the
functional API, because an interrupt in a task interrupts the task and
entrypoint.

This is technically breaking (users resuming multiple hanging interrupts
w/ multiple resume calls can no longer do this... but the behavior for
this case was non-deterministic in the first place so we can sell this
as a fix).

libs/langgraph/langgraph/pregel/_loop.py

@@ -568,6 +568,36 @@ class PregelLoop:
             if task := tasks.get(tid):
                 task.writes.append((k, v))
 
+    def _pending_interrupts(self) -> set[str]:
+        """Return the set of interrupt ids that are pending without corresponding resume values."""
+        # mapping of task ids to interrupt ids
+        pending_interrupts: dict[str, str] = {}
+
+        # set of resume task ids
+        pending_resumes: set[str] = set()
+
+        for task_id, write_type, value in self.checkpoint_pending_writes:
+            if write_type == INTERRUPT:
+                # interrupts is always a list, but there should only be one element
+                pending_interrupts[task_id] = value[0].id
+            elif write_type == RESUME:
+                pending_resumes.add(task_id)
+
+        resumed_interrupt_ids = {
+            pending_interrupts[task_id]
+            for task_id in pending_resumes
+            if task_id in pending_interrupts
+        }
+
+        # Keep only interrupts whose interrupt_id is not resumed
+        hanging_interrupts: set[str] = {
+            interrupt_id
+            for interrupt_id in pending_interrupts.values()
+            if interrupt_id not in resumed_interrupt_ids
+        }
+
+        return hanging_interrupts
+
     def _first(
         self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
     ) -> set[str] | None:
@@ -590,16 +620,24 @@ class PregelLoop:
 
         # map command to writes
         if isinstance(self.input, Command):
-            if resume_is_map := (
-                (resume := self.input.resume) is not None
-                and isinstance(resume, dict)
-                and all(is_xxh3_128_hexdigest(k) for k in resume)
-            ):
-                self.config[CONF][CONFIG_KEY_RESUME_MAP] = self.input.resume
-            if resume is not None and not self.checkpointer:
-                raise RuntimeError(
-                    "Cannot use Command(resume=...) without checkpointer"
-                )
+            if (resume := self.input.resume) is not None:
+                if not self.checkpointer:
+                    raise RuntimeError(
+                        "Cannot use Command(resume=...) without checkpointer"
+                    )
+
+                if resume_is_map := (
+                    isinstance(resume, dict)
+                    and all(is_xxh3_128_hexdigest(k) for k in resume)
+                ):
+                    self.config[CONF][CONFIG_KEY_RESUME_MAP] = resume
+                else:
+                    if len(self._pending_interrupts()) > 1:
+                        raise RuntimeError(
+                            "When there are multiple pending interrupts, you must specify the interrupt id when resuming. "
+                            "Docs: https://docs.langchain.com/oss/python/langgraph/add-human-in-the-loop#resume-multiple-interrupts-with-one-invocation."
+                        )
+
             writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
             # group writes by task ID
             for tid, c, v in map_command(cmd=self.input):

libs/langgraph/tests/test_pregel.py

@@ -7222,7 +7222,11 @@ def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None:
 
         # get human input and resume
         if len(current_interrupts) > 0:
-            current_input = Command(resume=f"Resume #{invokes}")
+            # we resume one at a time to preserve original test behavior,
+            # but we could also resume all at once if we wanted
+            # with a single dict mapping of interrupt ids to resume values
+            resume = {current_interrupts[0].id: f"Resume #{invokes}"}
+            current_input = Command(resume=resume)
 
         # not more human input required, must be completed
         else:
@@ -7383,7 +7387,11 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
 
         # get human input and resume
         if len(current_interrupts) > 0:
-            current_input = Command(resume=f"Resume #{invokes}")
+            # we resume one at a time to preserve original test behavior,
+            # but we could also resume all at once if we wanted
+            # with a single dict mapping of interrupt ids to resume values
+            resume = {current_interrupts[0].id: f"Resume #{invokes}"}
+            current_input = Command(resume=resume)
 
         # not more human input required, must be completed
         else:
@@ -8333,3 +8341,51 @@ def test_subgraph_streaming_sync() -> None:
 
     assert result["last_chunk"].content == "today."
     assert result["num_chunks"] == 9
+
+
+def test_null_resume_disallowed_with_multiple_interrupts(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    class State(TypedDict):
+        text_1: str
+        text_2: str
+
+    def human_node_1(state: State):
+        value = interrupt({"text_to_revise": state["text_1"]})
+        return {"text_1": value}
+
+    def human_node_2(state: State):
+        value = interrupt({"text_to_revise": state["text_2"]})
+        return {"text_2": value}
+
+    graph_builder = StateGraph(State)
+    graph_builder.add_node("human_node_1", human_node_1)
+    graph_builder.add_node("human_node_2", human_node_2)
+
+    # Add both nodes in parallel from START
+    graph_builder.add_edge(START, "human_node_1")
+    graph_builder.add_edge(START, "human_node_2")
+
+    checkpointer = InMemorySaver()
+    graph = graph_builder.compile(checkpointer=checkpointer)
+
+    thread_id = str(uuid.uuid4())
+    config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
+    graph.invoke(
+        {"text_1": "original text 1", "text_2": "original text 2"}, config=config
+    )
+
+    resume_map = {
+        i.id: f"resume for prompt: {i.value['text_to_revise']}"
+        for i in graph.get_state(config).interrupts
+    }
+    with pytest.raises(
+        RuntimeError,
+        match="When there are multiple pending interrupts, you must specify the interrupt id when resuming.",
+    ):
+        graph.invoke(Command(resume="singular resume"), config=config)
+
+    assert graph.invoke(Command(resume=resume_map), config=config) == {
+        "text_1": "resume for prompt: original text 1",
+        "text_2": "resume for prompt: original text 2",
+    }

libs/langgraph/tests/test_pregel_async.py

@@ -9107,3 +9107,52 @@ async def test_subgraph_streaming_async() -> None:
 
     assert result["last_chunk"].content == "today."
     assert result["num_chunks"] == 9
+
+
+@NEEDS_CONTEXTVARS
+async def test_null_resume_disallowed_with_multiple_interrupts(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    class State(TypedDict):
+        text_1: str
+        text_2: str
+
+    async def human_node_1(state: State):
+        value = interrupt(state["text_1"])
+        return {"text_1": value}
+
+    async def human_node_2(state: State):
+        value = interrupt(state["text_2"])
+        return {"text_2": value}
+
+    graph_builder = StateGraph(State)
+    graph_builder.add_node("human_node_1", human_node_1)
+    graph_builder.add_node("human_node_2", human_node_2)
+
+    # Add both nodes in parallel from START
+    graph_builder.add_edge(START, "human_node_1")
+    graph_builder.add_edge(START, "human_node_2")
+
+    checkpointer = InMemorySaver()
+    graph = graph_builder.compile(checkpointer=checkpointer)
+
+    thread_id = str(uuid.uuid4())
+    config: RunnableConfig = {"configurable": {"thread_id": thread_id}}
+    await graph.ainvoke(
+        {"text_1": "original text 1", "text_2": "original text 2"}, config=config
+    )
+
+    resume_map = {
+        i.id: f"resume for prompt: {i.value}"
+        for i in (await graph.aget_state(config)).interrupts
+    }
+    with pytest.raises(
+        RuntimeError,
+        match="When there are multiple pending interrupts, you must specify the interrupt id when resuming.",
+    ):
+        await graph.ainvoke(Command(resume="singular resume"), config=config)
+
+    assert await graph.ainvoke(Command(resume=resume_map), config=config) == {
+        "text_1": "resume for prompt: original text 1",
+        "text_2": "resume for prompt: original text 2",
+    }