diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 465527bbd..e880617cf 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -317,14 +317,38 @@ class PregelLoop: writes_to_save: WritesT = [ w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id ] + list(writes) + self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) else: + writes_to_save: WritesT = [] + for channel, value in writes: + # + if channel == INTERRUPT: + new_interrupts = list(value) if isinstance(value, (list, tuple)) else [value] + # aggregate existing interrupts for this task + existing = next( + (v for tid, ch, v in self.checkpoint_pending_writes + if tid == task_id and ch == INTERRUPT), + None + ) + if existing is not None: + # backwards compat: support resuming tasks where saved interrupt value is not a list + existing_interrupts = list(existing) if isinstance(existing, (list, tuple)) else [existing] + # append if same interrupt id (multiple interrupt() calls in same task execution), + # otherwise replace (different PUSH tasks, each with unique interrupt id - see tests/test_pregel.py::test_interrupt_task_functional) + value = ( + existing_interrupts + new_interrupts + if existing_interrupts[0].id == new_interrupts[0].id + else new_interrupts + ) + else: + value = new_interrupts + writes_to_save.append((channel, value)) # remove existing writes for this task self.checkpoint_pending_writes = [ w for w in self.checkpoint_pending_writes if w[0] != task_id ] - writes_to_save = writes - # save writes - self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) + self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes_to_save) + if self.durability != "exit" and self.checkpointer_put_writes is not None: config = patch_configurable( self.checkpoint_config, @@ -477,8 +501,11 @@ class PregelLoop: skipped_interrupt_ids = self._pending_interrupts() - set(resume_map) self.skipped_task_ids = { task_id - for task_id, write_type, value in self.checkpoint_pending_writes - if write_type == INTERRUPT and value[0].id in skipped_interrupt_ids + for task_id, channel, value in self.checkpoint_pending_writes + if channel == INTERRUPT + # interrupts within a task are uncovered sequentially as resumes are provided, + # so we only need to check the last interrupt id + and (list(value) if isinstance(value, (list, tuple)) else [value])[-1].id in skipped_interrupt_ids } else: self.skipped_task_ids = set() @@ -535,25 +562,35 @@ class PregelLoop: for task_id in self.skipped_task_ids if not self.tasks[task_id].writes } - # output writes for blocked tasts so they are still visible in the stream - for task_id, write_type, value in self.checkpoint_pending_writes: - if task_id in self.skipped_task_ids: - self.output_writes(task_id, [(write_type, value)]) + # output interrupt writes for blocked tasks so they are still visible in the stream + for task_id, channel, value in self.checkpoint_pending_writes: + if task_id in self.skipped_task_ids and channel == INTERRUPT: + # find resume count for this task + resumes = next( + (v for tid, ch, v in self.checkpoint_pending_writes + if tid == task_id and ch == RESUME), + None + ) + resume_count = len(resumes) if resumes is not None else 0 + interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value] + # only output unresumed interrupts + if resume_count < len(interrupt_list): + self.output_writes(task_id, [(INTERRUPT, interrupt_list[resume_count:])]) return True def after_tick(self) -> None: if self.skipped_task_ids: - # raise an early interrupt for skipped tasks. - # this would ordinarily be raised after the PUSH task is executed, - # but since we know there are no resumes for these tasks, we can - # prevent unecessary node re-execution by raising in this tick. - interrupts = tuple( - value[0] - for _, write_type, value in self.checkpoint_pending_writes - if write_type == INTERRUPT - ) - raise GraphInterrupt(interrupts) + # raise early GraphInterrupt for skipped tasks. + # since we know len(resumes) != len(interrupts) for these tasks, we + # can prevent unecessary node re-execution by raising preemptively + interrupts = [] + for task_id, channel, value in self.checkpoint_pending_writes: + if channel == INTERRUPT and task_id in self.skipped_task_ids: + interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value] + interrupts.extend(interrupt_list) + if interrupts: + raise GraphInterrupt(interrupts) self.skipped_task_ids.clear() # finish superstep @@ -607,30 +644,25 @@ class PregelLoop: 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] = {} + # mapping of task ids to (interrupt_id, interrupt_count) + pending_interrupts: dict[str, tuple[str, int]] = {} + # mapping of task ids to resume count + pending_resumes: dict[str, int] = {} - # set of resume task ids - pending_resumes: set[str] = set() + for task_id, channel, value in self.checkpoint_pending_writes: + if channel == INTERRUPT: + interrupt_list = list(value) if isinstance(value, (list, tuple)) else [value] + pending_interrupts[task_id] = (interrupt_list[0].id, len(interrupt_list)) + elif channel == RESUME: + # count resume values for this task + resume_list = value if isinstance(value, list) else [value] + pending_resumes[task_id] = len(resume_list) - 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 + # keep only interrupt ids where resume_count < interrupt_count hanging_interrupts: set[str] = { interrupt_id - for interrupt_id in pending_interrupts.values() - if interrupt_id not in resumed_interrupt_ids + for task_id, (interrupt_id, interrupt_count) in pending_interrupts.items() + if pending_resumes.get(task_id, 0) < interrupt_count } return hanging_interrupts @@ -1063,6 +1095,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" + super().put_writes(task_id, writes) if not writes or self.cache is None or not hasattr(self, "tasks"): return diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index a2a364358..ccc2b9ab6 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -328,57 +328,6 @@ def test_interrupt_with_send_payloads_sequential_resume( assert node_counter["map_node"] == 5 -@pytest.mark.xfail(reason="Node resumes after partial interrupt resume", strict=False) -def test_node_with_multiple_interrupts_requires_full_resume( - sync_checkpointer: BaseCheckpointSaver, -) -> None: - node_counter = {"double_interrupt": 0} - - class State(TypedDict): - input: str - - def double_interrupt_node(state: State): - node_counter["double_interrupt"] += 1 - first = interrupt({"step": "first"}) - second = interrupt({"step": "second"}) - return {"input": f"{first}-{second}"} - - builder = StateGraph(State) - builder.add_node("double_interrupt", double_interrupt_node) - builder.add_edge(START, "double_interrupt") - builder.add_edge("double_interrupt", END) - - graph = builder.compile(checkpointer=sync_checkpointer) - - config = {"configurable": {"thread_id": "test_double_interrupt_sync"}} - - result = graph.invoke({"input": "start"}, config=config) - - interrupts = result.get("__interrupt__", []) - assert len(interrupts) == 1 - first_interrupt = interrupts[0] - assert node_counter["double_interrupt"] == 1 - - partial = graph.invoke( - Command(resume={first_interrupt.id: "human_first"}), config=config - ) - - # Expected behavior: node should not execute again until all resume values are provided - assert node_counter["double_interrupt"] == 1 - - remaining_interrupts = partial.get("__interrupt__", []) - assert len(remaining_interrupts) == 1 - second_interrupt = remaining_interrupts[0] - - final_result = graph.invoke( - Command(resume={second_interrupt.id: "human_second"}), config=config - ) - - assert node_counter["double_interrupt"] == 2 - assert "input" in final_result - assert final_result["input"] == "human_first-human_second" - - async def test_interrupt_with_send_payloads_sequential_resume_async( async_checkpointer: BaseCheckpointSaver, ) -> None: @@ -466,3 +415,175 @@ async def test_interrupt_with_send_payloads_sequential_resume_async( # Map node runs 3 times initially (item1 completes, 2 dangerous_items interrupt), # then 1 time on first resume, then 1 time on second resume assert node_counter["map_node"] == 5 + + +def test_node_with_multiple_interrupts_requires_full_resume( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test a number of different resume patterns for a node with multiple interrupts, + + Ensures that a node is not re-executed until valid resume values have been provided to all + discovered interrupts""" + + node_counter = 0 + + class State(TypedDict): + input: str + + def double_interrupt_node(state: State): + nonlocal node_counter + node_counter += 1 + first = interrupt("first") + second = interrupt("second") + third = interrupt("third") + return {"input": f"{first}-{second}-{third}"} + + builder = StateGraph(State) + builder.add_node("double_interrupt", double_interrupt_node) + builder.add_edge(START, "double_interrupt") + builder.add_edge("double_interrupt", END) + + graph = builder.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "test_double_interrupt"}} + + result = graph.invoke({"input": "start"}, config=config) + + interrupts = result.get("__interrupt__", []) + assert len(interrupts) == 1 + first_interrupt = interrupts[0] + assert node_counter == 1 + + # invoke with an interrupt map that matches double_interrupt_node. + # this should execute the node + partial = graph.invoke( + Command(resume={first_interrupt.id: "human_first"}), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 2 + + # invoke with an interrupt map that DOES NOT match double_interrupt_node. + # this should not execute the node because the optimization kicks in + partial = graph.invoke( + Command(resume={"00000000000000000000000000000000": "nothing_burger"}), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 2 + + # invoke with None resume. this should execute the node + partial = graph.invoke( + None, config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 3 + + # invoke with nonspecific resume. this should execute the node + partial = graph.invoke( + Command(resume="human_second"), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + print("REMAINING INTERRUPTS: ", remaining_interrupts) + assert remaining_interrupts[0].value == "third" + assert node_counter == 4 + + # finally, invoke with an interrupt map that matches double_interrupt_node. + # this should execute the node and all interrupts should be resolved + final_result = graph.invoke( + Command(resume="human_third"), config=config + ) + assert "input" in final_result + assert final_result["input"] == "human_first-human_second-human_third" + assert node_counter == 5 + + +async def test_node_with_multiple_interrupts_requires_full_resume_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test a number of different resume patterns for a node with multiple interrupts, + + Ensures that a node is not re-executed until valid resume values have been provided to all + discovered interrupts""" + + node_counter = 0 + + class State(TypedDict): + input: str + + def double_interrupt_node(state: State): + nonlocal node_counter + node_counter += 1 + first = interrupt("first") + second = interrupt("second") + third = interrupt("third") + return {"input": f"{first}-{second}-{third}"} + + builder = StateGraph(State) + builder.add_node("double_interrupt", double_interrupt_node) + builder.add_edge(START, "double_interrupt") + builder.add_edge("double_interrupt", END) + + graph = builder.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "test_double_interrupt"}} + + result = await graph.ainvoke({"input": "start"}, config=config) + + interrupts = result.get("__interrupt__", []) + assert len(interrupts) == 1 + first_interrupt = interrupts[0] + assert node_counter == 1 + + # invoke with an interrupt map that matches double_interrupt_node. + # this should execute the node + partial = await graph.ainvoke( + Command(resume={first_interrupt.id: "human_first"}), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 2 + + # invoke with an interrupt map that DOES NOT match double_interrupt_node. + # this should not execute the node because the optimization kicks in + partial = await graph.ainvoke( + Command(resume={"00000000000000000000000000000000": "nothing_burger"}), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 2 + + # invoke with None resume. this should execute the node + partial = await graph.ainvoke( + None, config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 3 + + # invoke with nonspecific resume. this should execute the node + partial = await graph.ainvoke( + Command(resume="human_second"), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + print("REMAINING INTERRUPTS: ", remaining_interrupts) + assert remaining_interrupts[0].value == "third" + assert node_counter == 4 + + # finally, invoke with an interrupt map that matches double_interrupt_node. + # this should execute the node and all interrupts should be resolved + final_result = await graph.ainvoke( + Command(resume="human_third"), config=config + ) + assert "input" in final_result + assert final_result["input"] == "human_first-human_second-human_third" + assert node_counter == 5