diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index f8031ca6a..ff0f1330a 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -259,8 +259,8 @@ class PregelLoop: _reapplied_task_ids: set[str] # Tasks that wrote during this run, the only ones exit mode saves at exit. _written_task_ids: set[str] - # An update or fork checkpoint addressed by `checkpoint_id` that the thread - # has moved past, so it may already have children. + # A checkpoint addressed by `checkpoint_id` that the thread has moved past, + # so it may already have children. _addressed_checkpoint_superseded: bool = False # The checkpoint_config that points at the parent loaded at `__enter__` @@ -351,7 +351,6 @@ class PregelLoop: self.interrupt_before = interrupt_before self.manager = manager self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) - self.is_replaying = CONFIG_KEY_CHECKPOINT_ID in config[CONF] self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes self.retry_policy = retry_policy @@ -409,6 +408,14 @@ class PregelLoop: else () ) self.prev_checkpoint_config = None + # Whether the run points this loop at a checkpoint: a truthy + # `checkpoint_id`, this graph's `checkpoint_map` entry, or a replaying + # parent. Every task config carries `checkpoint_id: None`, so a key + # check would count every subgraph. + self._addressed = bool( + self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) + or self.config[CONF].get(CONFIG_KEY_REPLAY_STATE) + ) runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME) self.control = runtime.control if isinstance(runtime, Runtime) else None @@ -937,7 +944,7 @@ class PregelLoop: # writes so that interrupt() calls re-fire instead of returning # stale values. But if we're actively resuming, keep them — # multi-interrupt scenarios need previously resolved values preserved. - is_time_traveling = self.is_replaying and ( + is_time_traveling = self._addressed and ( # Time-travel to a subgraph checkpoint: the parent sets # RESUMING=True (it can't distinguish time-travel from resume), # so we check if this subgraph's own ns is in checkpoint_map. @@ -956,6 +963,10 @@ class PregelLoop: or configurable.get(CONFIG_KEY_RESUMING, False) ) ) + # Finished tasks rerun only on time travel, or on a resume at a + # checkpoint the thread has moved past. A resume addressed at the + # thread's latest checkpoint hands their writes back, as without an id. + self.is_replaying = is_time_traveling or self._addressed_checkpoint_superseded if is_time_traveling: self.checkpoint_pending_writes = [ w for w in self.checkpoint_pending_writes if w[1] != RESUME @@ -1171,11 +1182,11 @@ class PregelLoop: # checkpoint ID since the fork was created after the subgraph's # checkpoints from the original execution. # - # Only gate on is_time_traveling (not is_replaying). When the + # Only gate on is_time_traveling (not _addressed). When the # client resumes with an explicit checkpoint_id that happens to # point at the current head (e.g. LangGraph Studio sending # `checkpoint: {checkpoint_id}` alongside Command(resume=...)), - # is_replaying is True but is_time_traveling is False. In that + # the run is addressed but not time traveling. In that # case subgraphs should load their latest checkpoint normally, # not go through ReplayState's before-bound lookup which would # miss subgraph checkpoints created during processing of the @@ -1795,7 +1806,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): # This covers both normal replay and subgraphs resolved via # checkpoint_map during time-travel. saved = self.checkpointer.get_tuple(self.checkpoint_config) - if saved is not None and saved.metadata.get("source") in ("update", "fork"): + if saved is not None: self._addressed_checkpoint_superseded = checkpoint_superseded( self.checkpointer, self.checkpoint_config, saved ) @@ -2051,7 +2062,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): # This covers both normal replay and subgraphs resolved via # checkpoint_map during time-travel. saved = await self.checkpointer.aget_tuple(self.checkpoint_config) - if saved is not None and saved.metadata.get("source") in ("update", "fork"): + if saved is not None: self._addressed_checkpoint_superseded = await acheckpoint_superseded( self.checkpointer, self.checkpoint_config, saved ) diff --git a/libs/langgraph/tests/test_delta_channel_exit_mode.py b/libs/langgraph/tests/test_delta_channel_exit_mode.py index 70e2627f6..b119bfa48 100644 --- a/libs/langgraph/tests/test_delta_channel_exit_mode.py +++ b/libs/langgraph/tests/test_delta_channel_exit_mode.py @@ -521,10 +521,11 @@ def test_addressed_resume_keeps_a_rerun_tasks_write_to_a_new_channel( graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "t"}} graph.invoke(_both("in"), config, durability=durability) - head = graph.get_state(config).config + paused = graph.get_state(config).config + graph.update_state(config, None) live = graph.invoke( - Command(resume="yes", update={"flag": True}), head, durability=durability + Command(resume="yes", update={"flag": True}), paused, durability=durability ) state = graph.get_state(config) diff --git a/libs/langgraph/tests/test_interrupt_state.py b/libs/langgraph/tests/test_interrupt_state.py index 1d7cf1b38..0c0ee5662 100644 --- a/libs/langgraph/tests/test_interrupt_state.py +++ b/libs/langgraph/tests/test_interrupt_state.py @@ -175,6 +175,34 @@ def test_finished_task_output_applies_once( assert sorted(graph.get_state(config).values["log"]) == sorted(result["log"]) +@pytest.mark.parametrize("checkpoint_id", ["none", "head"]) +def test_resume_at_the_latest_checkpoint_does_not_rerun_finished_task( + sync_checkpointer: BaseCheckpointSaver, checkpoint_id: str +) -> None: + calls: Counter[str] = Counter() + graph = _build_parallel(sync_checkpointer, calls) + config = _config() + graph.invoke({"log": []}, config) + + def resume(question: str, answer: str) -> Any: + snapshot = graph.get_state(config) + target = ( + snapshot.config + if checkpoint_id == "head" + else {"configurable": {**config["configurable"], "checkpoint_id": None}} + ) + return graph.invoke( + Command(resume={_interrupt_by_value(snapshot, question).id: answer}), + target, + ) + + resume("A1", "yes") + result = resume("B", "ok") + + assert sorted(result["log"]) == ["a:yes", "b:ok"] + assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again" + + @NEEDS_CONTEXTVARS async def test_finished_task_does_not_report_answered_interrupt_async( async_checkpointer: BaseCheckpointSaver, @@ -497,6 +525,64 @@ def test_functional_task_finished_with_none_is_not_rerun( assert calls == {"a": 2, "b": 3} +def test_subgraph_finished_task_is_not_rerun( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + calls: Counter[str] = Counter() + builder = StateGraph(State) + builder.add_node("child", _build_parallel(None, calls)) + builder.add_edge(START, "child") + graph = builder.compile(checkpointer=sync_checkpointer) + config = _config() + + graph.invoke({"log": []}, config, durability=durability) + snapshot = graph.get_state(config) + graph.invoke( + Command(resume={_interrupt_by_value(snapshot, "A1").id: "yes"}), + config, + durability=durability, + ) + snapshot = graph.get_state(config) + result = graph.invoke( + Command(resume={_interrupt_by_value(snapshot, "B").id: "ok"}), + config, + durability=durability, + ) + + assert sorted(result["log"]) == ["a:yes", "b:ok"] + assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again" + + +def test_nested_entrypoint_finished_task_is_not_rerun( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + calls: Counter[str] = Counter() + + @task + def step_1(query: str) -> str: + calls["step_1"] += 1 + return f"{query} bar" + + @task + def human_feedback(query: str) -> str: + calls["human_feedback"] += 1 + return f"{query} {interrupt('feedback?')}" + + @entrypoint() + def child(query: str) -> str: + return human_feedback(step_1(query).result()).result() + + @entrypoint(checkpointer=sync_checkpointer) + def parent(query: str) -> str: + return child.invoke(query) + + config = _config() + parent.invoke("foo", config) + + assert parent.invoke(Command(resume="baz"), config) == "foo bar baz" + assert calls == {"step_1": 1, "human_feedback": 2} + + # --- Reading task status from recorded writes --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d698da84f..e7ff4bd10 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7737,6 +7737,8 @@ def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None: def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None: # --- CHILD GRAPH --- + answered: Counter[str] = Counter() + class ChildState(BaseModel): prompt: str = Field(..., description="What is going to be asked to the user?") human_input: str | None = Field(None, description="What the human said") @@ -7746,6 +7748,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N def get_human_input(state: ChildState): human_input = interrupt(state.prompt) + answered["human"] += 1 return dict( human_inputs=[human_input], # update parent state @@ -7753,6 +7756,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N def get_dolphin_input(state: ChildState): human_input = interrupt(state.prompt) + answered["dolphin"] += 1 return dict( human_inputs=[human_input], # update parent state @@ -7846,6 +7850,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N assert invokes == 5 assert len(events) == 5 + assert answered == {"human": 2, "dolphin": 2} def test_pregel_loop_refcount(): diff --git a/libs/langgraph/tests/test_time_travel.py b/libs/langgraph/tests/test_time_travel.py index fe544274f..2a2649a46 100644 --- a/libs/langgraph/tests/test_time_travel.py +++ b/libs/langgraph/tests/test_time_travel.py @@ -1177,6 +1177,43 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id( assert result["value"] == ["sub_a", "human:answer", "sub_b"] +def test_subgraph_replay_reasks_answered_parallel_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying a subgraph checkpoint asks its step's questions again, including + one already answered. An update checkpoint replays in place instead of + forking, so its tasks keep their ids and its saved answers could be reused.""" + + subgraph = ( + StateGraph(State) + .add_node("a", lambda s: {"value": [f"a:{interrupt('a')}"]}) + .add_node("b", lambda s: {"value": [f"b:{interrupt('b')}"]}) + .add_edge(START, "a") + .add_edge(START, "b") + .compile(checkpointer=True) + ) + graph = ( + StateGraph(State) + .add_node("subgraph_node", subgraph) + .add_edge(START, "subgraph_node") + .compile(checkpointer=sync_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + + def subgraph_head() -> RunnableConfig: + return graph.get_state(config, subgraphs=True).tasks[0].state.config + + graph.invoke({"value": []}, config) + graph.update_state(subgraph_head(), {"value": ["u"]}, as_node="__start__") + result = graph.invoke(None, config) + question_a = next(i for i in result["__interrupt__"] if i.value == "a") + graph.invoke(Command(resume={question_a.id: "yes"}), config) + + result = graph.invoke(None, subgraph_head()) + + assert sorted(i.value for i in result["__interrupt__"]) == ["a", "b"] + + def test_subgraph_replay_loads_accumulated_state_then_resume( sync_checkpointer: BaseCheckpointSaver, ) -> None: