From a43acc33bd913095b6180082469bedd38323a182 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Wed, 10 Sep 2025 08:31:11 -0400 Subject: [PATCH] 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 | 58 ++++++++++++++++++---- libs/langgraph/tests/test_pregel.py | 60 ++++++++++++++++++++++- libs/langgraph/tests/test_pregel_async.py | 49 ++++++++++++++++++ 3 files changed, 155 insertions(+), 12 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index c9300f3d4..b0c98a5e3 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/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): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8c9dda4fa..cbf6d03b6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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", + } diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 6d42a0244..b5bd272af 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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", + }