diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 7ea448ac0..ba737d4ea 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -316,35 +316,36 @@ class PregelLoop: ] writes_to_save: WritesT = [ w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id - ] + list(writes) + ] + [ + (c, v) for c, v in writes if c != RESUME + ] self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes) else: - # build map of existing interrupts by interrupt id for this task + # build map of existing interrupts for this task for quick lookup existing_interrupts_by_id = { # interrupt id -> list of interrupts v[0].id: v for tid, ch, v in self.checkpoint_pending_writes if tid == task_id and ch == INTERRUPT } - # check if a resume write exists - has_resume = any(ch == RESUME for ch, _ in writes) - writes_to_save = [] for ch, v in writes: if ch == INTERRUPT: - # merge with existing interrupts if same interrupt id + # we merge new interrupt writes with existing interrupts writes if they + # occured within the same task (which means they have the same interrupt id) new_interrupts = v if isinstance(v, list) else list(v) if new_interrupts and ( existing := existing_interrupts_by_id.get(new_interrupts[0].id) ): - # found existing interrupts with same interrupt id - # if a resume write exists, it means this is a new interrupt - # so we are safe to merge it into the existing interrupt writes - v = existing + new_interrupts if has_resume else existing + # if the graph is invoked with None, we will hit the same interrupt + # that was raised before, in this case we don't want to duplicate its write + # so we just keep the existing checkpoint writes + v = existing + new_interrupts if self.input is not None else existing writes_to_save.append((ch, v)) else: + # we add non-interrupt writes as-is writes_to_save.append((ch, v)) - # replace all writes for this task_id in one shot + # replace all writes for this task_id with the merged writes self.checkpoint_pending_writes = [ w for w in self.checkpoint_pending_writes if w[0] != task_id ] + [(task_id, c, v) for c, v in writes_to_save] @@ -477,6 +478,8 @@ class PregelLoop: return False # prepare next tasks + print("checkpoint_pending_writes before: ", self.checkpoint_pending_writes) + self.tasks = prepare_next_tasks( self.checkpoint, self.checkpoint_pending_writes, @@ -496,6 +499,8 @@ class PregelLoop: cache_policy=self.cache_policy, ) + print("checkpoint_pending_writes after: ", self.checkpoint_pending_writes) + resume_map = self.config.get(CONF, {}).get(CONFIG_KEY_RESUME_MAP, {}) if resume_map: skipped_interrupt_ids = self._pending_interrupts() - set(resume_map) diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index 4bd30124d..85de4649b 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -582,3 +582,91 @@ async def test_node_with_multiple_interrupts_requires_full_resume_async( assert "input" in final_result assert final_result["input"] == "human_first-human_second-human_third" assert node_counter == 5 + +def test_invoke_interrupted_graph_with_none( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that invoking an interrupted graph with None does not duplicate interrupt writes""" + + 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") + 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_none_resume"}} + + result = graph.invoke({"input": "start"}, config=config) + first_history = list(graph.get_state_history(config)) + interrupts = result.get("__interrupt__", []) + assert len(interrupts) == 1 + assert node_counter == 1 + + # invoke with None. this should execute the node and the history should + # look the same as the first run + partial = graph.invoke( + None, config=config + ) + second_history = list(graph.get_state_history(config)) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "first" + assert node_counter == 2 + + # history should look the same for tasks and interrupts + print("first_history[0].interrupts: ", first_history[0].interrupts) + print("second_history[0].interrupts: ", second_history[0].interrupts) + print("first_history[0].tasks: ", first_history[0].tasks) + print("second_history[0].tasks: ", second_history[0].tasks) + assert first_history[0].interrupts == second_history[0].interrupts + assert first_history[0].tasks == second_history[0].tasks + + # now resume the first interrupt with some value + partial = graph.invoke( + Command(resume="weet"), config=config + ) + print("partial 3", partial) + third_history = list(graph.get_state_history(config)) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert remaining_interrupts[0].value == "second" + assert node_counter == 3 + + # invoke with None again. the history should look the same as + # the third run + partial = graph.invoke( + None, config=config + ) + print("partial 4", partial) + fourth_history = list(graph.get_state_history(config)) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 1 + assert node_counter == 4 + + print("\nthird_history[0].interrupts: ", third_history[0].interrupts) + print("fourth_history[0].interrupts: ", fourth_history[0].interrupts) + print("third_history[0].tasks: ", third_history[0].tasks) + print("fourth_history[0].tasks: ", fourth_history[0].tasks) + assert third_history[0].interrupts == fourth_history[0].interrupts + assert third_history[0].tasks == fourth_history[0].tasks + + # resume the graph once more with a real value + partial = graph.invoke( + Command(resume="bix"), config=config + ) + remaining_interrupts = partial.get("__interrupt__", []) + assert len(remaining_interrupts) == 0 + assert node_counter == 5