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", + }