From 51cbdbd5cd32614d3a73a4dd3f93823f472fd1a4 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Thu, 16 Apr 2026 08:29:48 -0400 Subject: [PATCH] fix: time travel when going back to interrupt node (#7498) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Fix: Create fork checkpoint on subgraph time travel ## Problem When time-traveling to a subgraph checkpoint that has an interrupt, and then resuming, the resume would load the **wrong state** — it would pick up the original execution's latest checkpoint instead of the time-traveled one. This happened because replaying from a subgraph checkpoint never created a new parent checkpoint. If the replay hit an interrupt before `after_tick()` ran, no checkpoint was written at all, so the parent's "latest" checkpoint was still the old one from the original execution. ## Fix When the loop detects a time-travel replay (not an `update_state` fork), it now **eagerly writes a fork checkpoint** at the start of the tick. This ensures: 1. The parent thread's latest checkpoint points to the replayed state 2. Subsequent `Command(resume=...)` calls find the correct checkpoint 3. Stale `INTERRUPT` pending writes from the old checkpoint are cleared (they reference old task IDs) Additionally, the subgraph replay logic now uses the **parent checkpoint ID** (from `prev_checkpoint_config`) when resolving subgraph checkpoints during time-travel, matching the existing behavior for `update_state` forks. ## Checkpoint flow diagrams ### Before fix: time travel leaves no fork ``` Original execution: C0 (start) --> C1 (step_a) --> C2 (ask_1 interrupt) --> C3 (resume) --> C4 (ask_2 interrupt) --> C5 (done) Time travel to C2 (subgraph config): Replay runs... hits interrupt... no new checkpoint written. Parent "latest" is still C5. Command(resume="new_answer"): Loads C5 (wrong!) instead of the replayed C2 state. ``` ### After fix: time travel creates a fork ``` Original execution: C0 --> C1 --> C2 --> C3 --> C4 --> C5 (done) Time travel to C2 (subgraph config): C0 --> C1 --> C2 --> C3 --> C4 --> C5 \ F1 (fork, source="fork") <-- new latest Command(resume="new_answer"): Loads F1 (correct!) --> resumes from the right state. After full resume: C0 --> C1 --> C2 --> C3 --> C4 --> C5 \ F1 --> F2 (ask_1 result) --> F3 (ask_2 interrupt) --> F4 (done) ``` ### Manual fork via `update_state` (unchanged) ``` C0 --> C1 --> C2 --> C3 \ U1 (source="update") <-- created by update_state() This path already worked. The fix skips update/fork sources so existing behavior is preserved. ``` ## Changes - **`libs/langgraph/langgraph/pregel/_loop.py`**: - Extract `is_time_traveling` flag from the existing replay detection logic for reuse - Write a fork checkpoint (`source="fork"`) eagerly at the start of a time-travel tick, before execution begins - Clear stale `INTERRUPT` pending writes when creating the fork (they reference old task IDs that won't match the new checkpoint) - Unify subgraph replay ID resolution: check `source in ("update", "fork")` instead of a separate `is_time_traveling` condition, since the new fork checkpoint now has `source="fork"` - **`libs/langgraph/tests/test_time_travel.py`** and **`test_time_travel_async.py`**: Added 4 new test cases (sync + async): - `test_replay_from_before_interrupt_then_resume` — replays from a checkpoint before an interrupt, resumes with a new answer, and verifies the full checkpoint history (source, next, values) at each stage - `test_subgraph_time_travel_resume_from_first_interrupt` — time-travels to a subgraph's first interrupt, resumes both interrupts with new answers, and verifies the fork creates a new branch while preserving the original - `test_subgraph_time_travel_resume_from_second_interrupt` — time-travels to a subgraph's second interrupt, resumes with a new answer, and verifies the first interrupt's original answer is preserved - `test_subgraph_time_travel_checkpoint_pattern` — verifies the fork checkpoint branches from the correct replay point and that the full checkpoint tree is correct after resume - **`libs/langgraph/tests/test_pregel.py`** / **`test_pregel_async.py`**: Updated existing `test_weather_subgraph_state` to account for the new fork checkpoint appearing in history (history length increases by 1) --- libs/langgraph/langgraph/pregel/_loop.py | 33 +- libs/langgraph/tests/test_pregel.py | 9 +- libs/langgraph/tests/test_pregel_async.py | 9 +- libs/langgraph/tests/test_time_travel.py | 734 +++++++++++++++++- .../langgraph/tests/test_time_travel_async.py | 410 +++++++++- 5 files changed, 1171 insertions(+), 24 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index aec4b74cd..fa26d2d9c 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -692,7 +692,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. - if self.is_replaying and ( + is_time_traveling = self.is_replaying 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. @@ -710,7 +710,8 @@ class PregelLoop: # (subgraph input is a Send arg, not a Command) or configurable.get(CONFIG_KEY_RESUMING, False) ) - ): + ) + if is_time_traveling: self.checkpoint_pending_writes = [ w for w in self.checkpoint_pending_writes if w[1] != RESUME ] @@ -765,6 +766,26 @@ class PregelLoop: if k in self.checkpoint["channel_versions"]: version = self.checkpoint["channel_versions"][k] self.checkpoint["versions_seen"][INTERRUPT][k] = version + # When time-traveling (replaying from a specific checkpoint), + # save a fork checkpoint so the replayed execution creates a + # new branch. Without this, if the execution hits an interrupt + # before after_tick() runs, no new checkpoint is created — + # the parent's latest checkpoint remains the old one and + # subsequent resumes load the wrong state. + # Skip for update_state forks (source=update/fork) since they + # already have their own fork checkpoint. + if is_time_traveling and self.checkpoint_metadata.get("source") not in ( + "update", + "fork", + ): + # Clear old INTERRUPT writes from the loaded checkpoint. + # The fork will have a new checkpoint_id which changes + # task IDs — stale interrupt writes would accumulate and + # confuse the multiple-interrupt check in future resumes. + self.checkpoint_pending_writes = [ + w for w in self.checkpoint_pending_writes if w[1] != INTERRUPT + ] + self._put_checkpoint({"source": "fork"}) # produce values output self._emit( "values", map_output_values, self.output_keys, True, self.channels @@ -807,14 +828,18 @@ class PregelLoop: if not self.is_nested: # Pass the resolved before-bound checkpoint ID so subgraphs can # find their corresponding checkpoint without re-fetching the - # parent. For forks (source=update), use the fork's parent + # parent. For forks (source=update/fork), use the fork's parent # checkpoint ID since the fork was created after the subgraph's # checkpoints from the original execution. replay_state: ReplayState | None = None if self.is_replaying: replay_checkpoint_id = self.checkpoint["id"] if ( - self.checkpoint_metadata.get("source") == "update" + self.checkpoint_metadata.get("source") + in ( + "update", + "fork", + ) and self.prev_checkpoint_config ): replay_checkpoint_id = self.prev_checkpoint_config[CONF].get( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f84f85c37..e4d49c46d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -615,8 +615,11 @@ def test_run_from_checkpoint_id_retains_previous_writes( ) ] - assert len(new_history) == len(history) + 1 - for original, new in zip(history, new_history[1:]): + # +2: one fork checkpoint from time travel, one from the new execution + assert len(new_history) == len(history) + 2 + # new_history[0] is the new execution result, new_history[1] is the fork + assert new_history[1].metadata["source"] == "fork" + for original, new in zip(history, new_history[2:]): assert original.values == new.values assert original.next == new.next assert original.metadata["step"] == new.metadata["step"] @@ -624,7 +627,7 @@ def test_run_from_checkpoint_id_retains_previous_writes( def _get_tasks(hist: list, start: int): return [h.tasks for h in hist[start:]] - assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + assert _get_tasks(new_history, 2) == _get_tasks(history, 0) def test_batch_two_processes_in_out() -> None: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index efb3bbfc1..b6ef6b088 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2086,8 +2086,11 @@ async def test_run_from_checkpoint_id_retains_previous_writes( ) ] - assert len(new_history) == len(history) + 1 - for original, new in zip(history, new_history[1:]): + # +2: one fork checkpoint from time travel, one from the new execution + assert len(new_history) == len(history) + 2 + # new_history[0] is the new execution result, new_history[1] is the fork + assert new_history[1].metadata["source"] == "fork" + for original, new in zip(history, new_history[2:]): assert original.values == new.values assert original.next == new.next assert original.metadata["step"] == new.metadata["step"] @@ -2095,7 +2098,7 @@ async def test_run_from_checkpoint_id_retains_previous_writes( def _get_tasks(hist: list, start: int): return [h.tasks for h in hist[start:]] - assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + assert _get_tasks(new_history, 2) == _get_tasks(history, 0) async def test_cond_edge_after_send() -> None: diff --git a/libs/langgraph/tests/test_time_travel.py b/libs/langgraph/tests/test_time_travel.py index 95097ba7e..4eea8d2aa 100644 --- a/libs/langgraph/tests/test_time_travel.py +++ b/libs/langgraph/tests/test_time_travel.py @@ -37,6 +37,7 @@ def _checkpoint_summary(history: list) -> list[dict]: Returns a list of dicts (newest-first, matching get_state_history order) with: - id: short checkpoint id suffix (last 6 chars) - parent_id: short parent checkpoint id suffix or None + - source: checkpoint metadata source (input, loop, fork, update) - next: tuple of next node names - values: channel values snapshot """ @@ -52,6 +53,7 @@ def _checkpoint_summary(history: list) -> list[dict]: { "id": cid[-6:], "parent_id": pid[-6:] if pid else None, + "source": s.metadata.get("source"), "next": s.next, "values": s.values, } @@ -280,6 +282,116 @@ def test_replay_from_before_interrupt_refires( assert call_count["node_b"] == 1 # NOT re-executed (after interrupt) +def test_replay_from_before_interrupt_then_resume( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Replay from checkpoint before interrupt node, then resume with a new + answer and verify the graph completes with the new value. + + Graph: START --> node_a --> ask_human (interrupt) --> node_b --> END + + Original run: + source=input next=(__start__,) values=[] + source=loop next=(node_a,) values=[] + source=loop next=(ask_human,) values=[a] <-- replay from here + source=loop next=(node_b,) values=[a, human:old_answer] + source=loop next=() values=[a, human:old_answer, b] + + After replay (fork created) + resume with "new_answer": + source=input next=(__start__,) values=[] + source=loop next=(node_a,) values=[] + source=loop next=(ask_human,) values=[a] <-- branch point + source=loop next=(node_b,) values=[a, human:old_answer] + source=loop next=() values=[a, human:old_answer, b] (old branch) + source=fork next=(ask_human,) values=[a] <-- fork from branch point + source=loop next=(node_b,) values=[a, human:new_answer] + source=loop next=() values=[a, human:new_answer, b] (new branch) + """ + + called: list[str] = [] + + def node_a(state: State) -> State: + called.append("node_a") + return {"value": ["a"]} + + def ask_human(state: State) -> State: + called.append("ask_human") + answer = interrupt("What is your input?") + return {"value": [f"human:{answer}"]} + + def node_b(state: State) -> State: + called.append("node_b") + return {"value": ["b"]} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("ask_human", ask_human) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "ask_human") + .add_edge("ask_human", "node_b") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: invoke until interrupt, then resume to complete --- + graph.invoke({"value": []}, config) + graph.invoke(Command(resume="old_answer"), config) + + original_history = list(graph.get_state_history(config)) + original = _checkpoint_summary(original_history) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["a", "human:old_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:old_answer"]}), + ("loop", ("ask_human",), {"value": ["a"]}), + ("loop", ("node_a",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Replay from checkpoint before ask_human --- + before_ask = next(s for s in original_history if s.next == ("ask_human",)) + + called.clear() + replay_result = graph.invoke(None, before_ask.config) + assert replay_result["__interrupt__"][0].value == "What is your input?" + assert "ask_human" in called + assert "node_a" not in called # before the replay point, not re-executed + + # A fork checkpoint is now the latest — it branches from the replay point + post_replay = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"]) for s in post_replay] == [ + ("fork", ("ask_human",)), # <-- new fork (latest) + ("loop", ()), # original done + ("loop", ("node_b",)), + ("loop", ("ask_human",)), # branch point + ("loop", ("node_a",)), + ("input", ("__start__",)), + ] + + # --- Resume with a new answer --- + called.clear() + final_result = graph.invoke(Command(resume="new_answer"), config) + assert final_result["value"] == ["a", "human:new_answer", "b"] + assert "ask_human" in called + assert "node_b" in called + + final = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from fork) + ("loop", (), {"value": ["a", "human:new_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:new_answer"]}), + ("fork", ("ask_human",), {"value": ["a"]}), + # Original branch (preserved) + ("loop", (), {"value": ["a", "human:old_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:old_answer"]}), + ("loop", ("ask_human",), {"value": ["a"]}), + ("loop", ("node_a",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + def test_replay_interrupt_stable_across_replays( sync_checkpointer: BaseCheckpointSaver, ) -> None: @@ -320,8 +432,14 @@ def test_replay_interrupt_stable_across_replays( r = graph.invoke(None, before_ask.config) results.append(r) - assert all(r == results[0] for r in results) - assert "__interrupt__" in results[0] + # Each replay creates a fork with a unique interrupt ID, so we compare + # interrupt values and state values rather than full equality. + assert all("__interrupt__" in r for r in results) + assert all( + r["__interrupt__"][0].value == results[0]["__interrupt__"][0].value + for r in results + ) + assert all(r["value"] == results[0]["value"] for r in results) def test_fork_from_before_interrupt_refires( @@ -854,6 +972,290 @@ def test_subgraph_interrupt_replay_from_interrupt_checkpoint( assert "step_b" not in called +def test_subgraph_interrupt_replay_from_parent_then_resume( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Replay from the parent checkpoint where a subgraph interrupt fired, + then resume with a new answer. Verifies that a fork is created and the + full graph completes. Checks full checkpoint history at each stage.""" + + called: list[str] = [] + + def router(state: State) -> State: + called.append("router") + return {"value": ["routed"]} + + def step_a(state: State) -> State: + called.append("step_a") + return {"value": ["sub_a"]} + + def ask_human(state: State) -> State: + called.append("ask_human") + answer = interrupt("Provide input:") + return {"value": [f"human:{answer}"]} + + def step_b(state: State) -> State: + called.append("step_b") + return {"value": ["sub_b"]} + + subgraph = ( + StateGraph(State) + .add_node("step_a", step_a) + .add_node("ask_human", ask_human) + .add_node("step_b", step_b) + .add_edge(START, "step_a") + .add_edge("step_a", "ask_human") + .add_edge("ask_human", "step_b") + .compile(checkpointer=True) + ) + + def post_process(state: State) -> State: + called.append("post_process") + return {"value": ["post"]} + + graph = ( + StateGraph(State) + .add_node("router", router) + .add_node("subgraph_node", subgraph) + .add_node("post_process", post_process) + .add_edge(START, "router") + .add_edge("router", "subgraph_node") + .add_edge("subgraph_node", "post_process") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # Run until interrupt, then resume to complete + graph.invoke({"value": []}, config) + graph.invoke(Command(resume="old_answer"), config) + + # Original parent history (newest first) + original_history = list(graph.get_state_history(config)) + assert [s.next for s in original_history] == [ + (), # done + ("post_process",), + ("subgraph_node",), # subgraph ran, interrupt fired here + ("router",), + ("__start__",), + ] + + # Find the parent checkpoint where the interrupt fired + interrupt_checkpoint = next( + s for s in original_history if s.next == ("subgraph_node",) + ) + + # Replay from parent checkpoint — subgraph re-executes, interrupt re-fires + called.clear() + replay_result = graph.invoke(None, interrupt_checkpoint.config) + assert "__interrupt__" in replay_result + assert replay_result["__interrupt__"][0].value == "Provide input:" + assert "step_a" in called + assert "ask_human" in called + assert "step_b" not in called + + # Verify fork checkpoint was created + post_replay_history = list(graph.get_state_history(config)) + assert [s.next for s in post_replay_history] == [ + ("subgraph_node",), # fork (interrupt pending) + (), # original done + ("post_process",), + ("subgraph_node",), + ("router",), + ("__start__",), + ] + assert [s.metadata["source"] for s in post_replay_history] == [ + "fork", + "loop", + "loop", + "loop", + "loop", + "input", + ] + fork = post_replay_history[0] + assert ( + fork.parent_config["configurable"]["checkpoint_id"] + == interrupt_checkpoint.config["configurable"]["checkpoint_id"] + ) + + # Resume with a new answer — full graph should complete + called.clear() + final_result = graph.invoke(Command(resume="new_answer"), config) + assert "__interrupt__" not in final_result + assert "human:new_answer" in final_result["value"] + assert "sub_b" in final_result["value"] + assert "post" in final_result["value"] + assert "ask_human" in called + assert "step_b" in called + assert "post_process" in called + + # Final checkpoint history + final_history = list(graph.get_state_history(config)) + assert [s.next for s in final_history] == [ + (), # new branch done + ("post_process",), # new branch post_process + ("subgraph_node",), # fork + (), # original done + ("post_process",), + ("subgraph_node",), + ("router",), + ("__start__",), + ] + assert [s.metadata["source"] for s in final_history] == [ + "loop", + "loop", + "fork", + "loop", + "loop", + "loop", + "loop", + "input", + ] + + +def test_subgraph_replay_loads_accumulated_state_then_resume( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Two parent invocations, then replay from before the subgraph in the + 2nd invocation. The subgraph (checkpointer=True) should load its + accumulated state from the 1st invocation via ReplayState, re-fire + the interrupt, and then resume + complete. + + This tests the ReplayState path: the parent is replaying and the + subgraph uses list(before=parent_checkpoint_id) to find its + corresponding checkpoint from the original execution. + """ + + class SubState(TypedDict): + value: Annotated[list[str], operator.add] + + class ParentState(TypedDict): + results: Annotated[list[str], operator.add] + + started_state: list[dict] = [] + + def step_a(state: SubState) -> SubState: + started_state.append(dict(state)) + answer = interrupt("question_a") + return {"value": [f"a:{answer}"]} + + subgraph = ( + StateGraph(SubState) + .add_node("step_a", step_a) + .add_edge(START, "step_a") + .compile(checkpointer=True) + ) + + def parent_node(state: ParentState) -> ParentState: + return {"results": ["p"]} + + graph = ( + StateGraph(ParentState) + .add_node("parent_node", parent_node) + .add_node("sub_node", subgraph) + .add_edge(START, "parent_node") + .add_edge("parent_node", "sub_node") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # === 1st invocation: complete with answer "a1" === + graph.invoke({"results": []}, config) + graph.invoke(Command(resume="a1"), config) + + # step_a saw empty state (fresh subgraph) + assert started_state[0] == {"value": []} + + # === 2nd invocation: complete with answer "a2" === + started_state.clear() + graph.invoke({"results": []}, config) + graph.invoke(Command(resume="a2"), config) + + # Stateful subgraph retained state from 1st invocation + assert started_state[0] == {"value": ["a:a1"]} + + # Original history (newest first) + original_history = list(graph.get_state_history(config)) + assert [s.next for s in original_history] == [ + (), # 2nd done + ("sub_node",), # 2nd sub_node + ("parent_node",), # 2nd parent_node + ("__start__",), # 2nd input + (), # 1st done + ("sub_node",), # 1st sub_node + ("parent_node",), # 1st parent_node + ("__start__",), # 1st input + ] + + # Replay from before sub_node in 2nd invocation (newest match) + before_sub_2nd = [s for s in original_history if s.next == ("sub_node",)][0] + started_state.clear() + replay = graph.invoke(None, before_sub_2nd.config) + assert "__interrupt__" in replay + + # Subgraph should see accumulated state from END of 1st invocation + assert started_state[0] == {"value": ["a:a1"]} + + # Verify fork was created + post_replay_history = list(graph.get_state_history(config)) + assert [s.next for s in post_replay_history] == [ + ("sub_node",), # fork (interrupt pending) + (), # 2nd done + ("sub_node",), # 2nd sub_node + ("parent_node",), # 2nd parent_node + ("__start__",), # 2nd input + (), # 1st done + ("sub_node",), # 1st sub_node + ("parent_node",), # 1st parent_node + ("__start__",), # 1st input + ] + assert [s.metadata["source"] for s in post_replay_history] == [ + "fork", + "loop", + "loop", + "loop", + "input", + "loop", + "loop", + "loop", + "input", + ] + + # Resume with a new answer + started_state.clear() + final = graph.invoke(Command(resume="a3"), config) + assert "__interrupt__" not in final + assert final["results"] == ["p", "p"] + + # Final history + final_history = list(graph.get_state_history(config)) + assert [s.next for s in final_history] == [ + (), # new branch done + ("sub_node",), # fork + (), # 2nd done + ("sub_node",), # 2nd sub_node + ("parent_node",), # 2nd parent_node + ("__start__",), # 2nd input + (), # 1st done + ("sub_node",), # 1st sub_node + ("parent_node",), # 1st parent_node + ("__start__",), # 1st input + ] + assert [s.metadata["source"] for s in final_history] == [ + "loop", + "fork", + "loop", + "loop", + "loop", + "input", + "loop", + "loop", + "loop", + "input", + ] + + def test_subgraph_interrupt_full_flow( sync_checkpointer: BaseCheckpointSaver, ) -> None: @@ -1290,6 +1692,321 @@ def test_subgraph_time_travel_to_second_interrupt( assert "ask_1" not in called +def test_subgraph_time_travel_resume_from_first_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Time travel to a subgraph checkpoint at the first interrupt, then + resume through both interrupts with new answers. + + This verifies the key bug fix: after time-traveling to a subgraph + checkpoint with an interrupt, a fork checkpoint is created so that + subsequent resumes find the correct state (not the old branch tip). + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END + + Parent history after original run completes: + source=input next=(__start__,) values=[] + source=loop next=(executor,) values=[] + source=loop next=() values=[step_a_done, ask_1:answer_1, ask_2:answer_2] + + After time-traveling to 1st interrupt + resuming with new answers: + source=input next=(__start__,) values=[] + source=loop next=(executor,) values=[] <-- branch point + source=loop next=() values=[..., ask_2:answer_2] (old branch) + source=fork next=(executor,) values=[] <-- fork from time travel + source=loop next=() values=[..., ask_2:new_answer_2] (new branch) + """ + + called: list[str] = [] + + def step_a(state: State) -> State: + called.append("step_a") + return {"value": ["step_a_done"]} + + def ask_1(state: State) -> State: + called.append("ask_1") + answer = interrupt("Question 1?") + return {"value": [f"ask_1:{answer}"]} + + def ask_2(state: State) -> State: + called.append("ask_2") + answer = interrupt("Question 2?") + return {"value": [f"ask_2:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("step_a", step_a) + .add_node("ask_1", ask_1) + .add_node("ask_2", ask_2) + .add_edge(START, "step_a") + .add_edge("step_a", "ask_1") + .add_edge("ask_1", "ask_2") + .add_edge("ask_2", "__end__") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: hit both interrupts and resume --- + graph.invoke({"value": []}, config) + sub_config_at_first = graph.get_state(config, subgraphs=True).tasks[0].state.config + graph.invoke(Command(resume="answer_1"), config) + graph.invoke(Command(resume="answer_2"), config) + + original = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Time travel to first interrupt's subgraph checkpoint --- + called.clear() + replay_result = graph.invoke(None, sub_config_at_first) + assert replay_result["__interrupt__"][0].value == "Question 1?" + assert "step_a" not in called # before interrupt, not re-executed + + # Fork is now the latest parent checkpoint + post_tt = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"]) for s in post_tt] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), # original done + ("loop", ("executor",)), + ("input", ("__start__",)), + ] + + # --- Resume both interrupts with new answers --- + called.clear() + resume_1 = graph.invoke(Command(resume="new_answer_1"), config) + assert resume_1["__interrupt__"][0].value == "Question 2?" + assert "ask_1" in called + + called.clear() + resume_2 = graph.invoke(Command(resume="new_answer_2"), config) + assert resume_2["value"] == [ + "step_a_done", + "ask_1:new_answer_1", + "ask_2:new_answer_2", + ] + + # Verify final history: original branch preserved, new branch appended + final = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from time travel fork) + ( + "loop", + (), + {"value": ["step_a_done", "ask_1:new_answer_1", "ask_2:new_answer_2"]}, + ), + ("fork", ("executor",), {"value": []}), + # Original branch (preserved) + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + +def test_subgraph_time_travel_resume_from_second_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Time travel to a subgraph checkpoint at the second interrupt, then + resume with a new answer. The first interrupt's answer should be preserved. + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END + + Key assertion: after resuming from a time-travel to the 2nd interrupt, + the final state keeps ask_1's original answer but uses the new ask_2 answer. + """ + + called: list[str] = [] + + def step_a(state: State) -> State: + called.append("step_a") + return {"value": ["step_a_done"]} + + def ask_1(state: State) -> State: + called.append("ask_1") + answer = interrupt("Question 1?") + return {"value": [f"ask_1:{answer}"]} + + def ask_2(state: State) -> State: + called.append("ask_2") + answer = interrupt("Question 2?") + return {"value": [f"ask_2:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("step_a", step_a) + .add_node("ask_1", ask_1) + .add_node("ask_2", ask_2) + .add_edge(START, "step_a") + .add_edge("step_a", "ask_1") + .add_edge("ask_1", "ask_2") + .add_edge("ask_2", "__end__") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: hit both interrupts and resume --- + graph.invoke({"value": []}, config) + graph.invoke(Command(resume="answer_1"), config) + sub_config_at_second = graph.get_state(config, subgraphs=True).tasks[0].state.config + graph.invoke(Command(resume="answer_2"), config) + + original = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Time travel to second interrupt --- + called.clear() + replay_result = graph.invoke(None, sub_config_at_second) + assert replay_result["__interrupt__"][0].value == "Question 2?" + assert "step_a" not in called + assert "ask_1" not in called # already resolved, not re-executed + + # Fork is now the latest parent checkpoint + post_tt = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"]) for s in post_tt] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), # original done + ("loop", ("executor",)), + ("input", ("__start__",)), + ] + + # --- Resume with a new answer for ask_2 only --- + called.clear() + resume_result = graph.invoke(Command(resume="new_answer_2"), config) + # ask_1's original answer preserved, ask_2 uses the new answer + assert resume_result["value"] == [ + "step_a_done", + "ask_1:answer_1", + "ask_2:new_answer_2", + ] + + # Verify final history: original branch preserved, new branch appended + final = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from time travel fork) + ( + "loop", + (), + {"value": ["step_a_done", "ask_1:answer_1", "ask_2:new_answer_2"]}, + ), + ("fork", ("executor",), {"value": []}), + # Original branch (preserved) + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + +def test_subgraph_time_travel_checkpoint_pattern( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Verify the checkpoint pattern created by time travel to a subgraph + interrupt. A fork checkpoint should branch from the replay point and + become the latest parent checkpoint. + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> ask (interrupt) --> END + + Original run (after completing): + source=input next=(__start__,) values=[] + source=loop next=(executor,) values=[] <-- replay point + source=loop next=() values=[a:first] + + After time travel to interrupt + resume with "second": + source=input next=(__start__,) values=[] + source=loop next=(executor,) values=[] <-- branch point + source=loop next=() values=[a:first] (old branch) + source=fork next=(executor,) values=[] <-- fork + source=loop next=() values=[a:second] (new branch) + """ + + def ask(state: State) -> State: + answer = interrupt("Q?") + return {"value": [f"a:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("ask", ask) + .add_edge(START, "ask") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # Run until interrupt, then complete + graph.invoke({"value": []}, config) + sub_config = graph.get_state(config, subgraphs=True).tasks[0].state.config + graph.invoke(Command(resume="first"), config) + + original = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["a:first"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # Time travel to the interrupt + graph.invoke(None, sub_config) + + # Fork is now the latest, branching from the original replay point + post_tt = list(graph.get_state_history(config)) + post_tt_summary = _checkpoint_summary(post_tt) + assert [(s["source"], s["next"]) for s in post_tt_summary] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), + ("loop", ("executor",)), # <-- replay point / fork parent + ("input", ("__start__",)), + ] + # Verify the fork's parent is the original replay point + replay_point_id = sub_config["configurable"]["checkpoint_map"][""] + assert post_tt[0].parent_config["configurable"]["checkpoint_id"] == replay_point_id + + # Resume from the fork — graph completes with new answer + result = graph.invoke(Command(resume="second"), config) + assert result["value"] == ["a:second"] + + final = _checkpoint_summary(list(graph.get_state_history(config))) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch + ("loop", (), {"value": ["a:second"]}), + ("fork", ("executor",), {"value": []}), + # Original branch + ("loop", (), {"value": ["a:first"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + def test_subgraph_time_travel_after_completion( sync_checkpointer: BaseCheckpointSaver, ) -> None: @@ -2283,14 +3000,16 @@ def test_replay_creates_branch_preserving_old_checkpoints( # -- Post-replay checkpoint history (newest first) -- post_replay_history = list(graph.get_state_history(config)) post_summary = _checkpoint_summary(post_replay_history) - assert len(post_summary) == 7 # 5 original + 2 new branch checkpoints + # 5 original + 1 fork + 2 new branch checkpoints = 8 + assert len(post_summary) == 8 # Verify the full shape after replay assert [s["next"] for s in post_summary] == [ - (), # new branch tip (C6) - ("node_c",), # new branch (C5) - (), # old branch tip (C4) - ("node_c",), # old (C3) + (), # new branch tip + ("node_c",), # new branch + ("node_b",), # fork from replay point + (), # old branch tip + ("node_c",), # old ("node_b",), # branch point (C2) ("node_a",), # old (C1) ("__start__",), # old (C0) @@ -2298,6 +3017,7 @@ def test_replay_creates_branch_preserving_old_checkpoints( assert [s["values"] for s in post_summary] == [ {"value": ["a", "b2", "c"]}, # new branch tip {"value": ["a", "b2"]}, # new: node_b re-ran with call_count=2 + {"value": ["a"]}, # fork from replay point {"value": ["a", "b1", "c"]}, # old branch tip preserved {"value": ["a", "b1"]}, # old {"value": ["a"]}, # branch point diff --git a/libs/langgraph/tests/test_time_travel_async.py b/libs/langgraph/tests/test_time_travel_async.py index 6f72a6d1d..458b32c6a 100644 --- a/libs/langgraph/tests/test_time_travel_async.py +++ b/libs/langgraph/tests/test_time_travel_async.py @@ -46,6 +46,7 @@ def _checkpoint_summary(history: list) -> list[dict]: Returns a list of dicts (newest-first, matching get_state_history order) with: - id: short checkpoint id suffix (last 6 chars) - parent_id: short parent checkpoint id suffix or None + - source: checkpoint metadata source (input, loop, fork, update) - next: tuple of next node names - values: channel values snapshot """ @@ -61,6 +62,7 @@ def _checkpoint_summary(history: list) -> list[dict]: { "id": cid[-6:], "parent_id": pid[-6:] if pid else None, + "source": s.metadata.get("source"), "next": s.next, "values": s.values, } @@ -335,8 +337,14 @@ async def test_replay_interrupt_stable_across_replays( r = await graph.ainvoke(None, before_ask.config) results.append(r) - assert all(r == results[0] for r in results) - assert "__interrupt__" in results[0] + # Each replay creates a fork with a unique interrupt ID, so we compare + # interrupt values and state values rather than full equality. + assert all("__interrupt__" in r for r in results) + assert all( + r["__interrupt__"][0].value == results[0]["__interrupt__"][0].value + for r in results + ) + assert all(r["value"] == results[0]["value"] for r in results) @NEEDS_CONTEXTVARS @@ -1261,6 +1269,391 @@ async def test_subgraph_time_travel_after_completion_async( assert "ask_2:answer_2" in replay_result["value"] +@NEEDS_CONTEXTVARS +async def test_replay_from_before_interrupt_then_resume_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Replay from checkpoint before interrupt node, then resume with a new + answer and verify the graph completes with the new value. + + Graph: START --> node_a --> ask_human (interrupt) --> node_b --> END + """ + + called: list[str] = [] + + async def node_a(state: State) -> State: + called.append("node_a") + return {"value": ["a"]} + + async def ask_human(state: State) -> State: + called.append("ask_human") + answer = interrupt("What is your input?") + return {"value": [f"human:{answer}"]} + + async def node_b(state: State) -> State: + called.append("node_b") + return {"value": ["b"]} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("ask_human", ask_human) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "ask_human") + .add_edge("ask_human", "node_b") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: invoke until interrupt, then resume to complete --- + await graph.ainvoke({"value": []}, config) + await graph.ainvoke(Command(resume="old_answer"), config) + + original_history = [s async for s in graph.aget_state_history(config)] + original = _checkpoint_summary(original_history) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["a", "human:old_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:old_answer"]}), + ("loop", ("ask_human",), {"value": ["a"]}), + ("loop", ("node_a",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Replay from checkpoint before ask_human --- + before_ask = next(s for s in original_history if s.next == ("ask_human",)) + + called.clear() + replay_result = await graph.ainvoke(None, before_ask.config) + assert replay_result["__interrupt__"][0].value == "What is your input?" + assert "ask_human" in called + assert "node_a" not in called + + # A fork checkpoint is now the latest + post_replay = _checkpoint_summary( + [s async for s in graph.aget_state_history(config)] + ) + assert [(s["source"], s["next"]) for s in post_replay] == [ + ("fork", ("ask_human",)), + ("loop", ()), + ("loop", ("node_b",)), + ("loop", ("ask_human",)), + ("loop", ("node_a",)), + ("input", ("__start__",)), + ] + + # --- Resume with a new answer --- + called.clear() + final_result = await graph.ainvoke(Command(resume="new_answer"), config) + assert final_result["value"] == ["a", "human:new_answer", "b"] + assert "ask_human" in called + assert "node_b" in called + + final = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from fork) + ("loop", (), {"value": ["a", "human:new_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:new_answer"]}), + ("fork", ("ask_human",), {"value": ["a"]}), + # Original branch (preserved) + ("loop", (), {"value": ["a", "human:old_answer", "b"]}), + ("loop", ("node_b",), {"value": ["a", "human:old_answer"]}), + ("loop", ("ask_human",), {"value": ["a"]}), + ("loop", ("node_a",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + +@NEEDS_CONTEXTVARS +async def test_subgraph_time_travel_resume_from_first_interrupt_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Time travel to a subgraph checkpoint at the first interrupt, then + resume through both interrupts with new answers. + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END + """ + + called: list[str] = [] + + async def step_a(state: State) -> State: + called.append("step_a") + return {"value": ["step_a_done"]} + + async def ask_1(state: State) -> State: + called.append("ask_1") + answer = interrupt("Question 1?") + return {"value": [f"ask_1:{answer}"]} + + async def ask_2(state: State) -> State: + called.append("ask_2") + answer = interrupt("Question 2?") + return {"value": [f"ask_2:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("step_a", step_a) + .add_node("ask_1", ask_1) + .add_node("ask_2", ask_2) + .add_edge(START, "step_a") + .add_edge("step_a", "ask_1") + .add_edge("ask_1", "ask_2") + .add_edge("ask_2", "__end__") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: hit both interrupts and resume --- + await graph.ainvoke({"value": []}, config) + sub_config_at_first = ( + (await graph.aget_state(config, subgraphs=True)).tasks[0].state.config + ) + await graph.ainvoke(Command(resume="answer_1"), config) + await graph.ainvoke(Command(resume="answer_2"), config) + + original = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Time travel to first interrupt's subgraph checkpoint --- + called.clear() + replay_result = await graph.ainvoke(None, sub_config_at_first) + assert replay_result["__interrupt__"][0].value == "Question 1?" + assert "step_a" not in called + + # Fork is now the latest parent checkpoint + post_tt = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"]) for s in post_tt] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), # original done + ("loop", ("executor",)), + ("input", ("__start__",)), + ] + + # --- Resume both interrupts with new answers --- + called.clear() + resume_1 = await graph.ainvoke(Command(resume="new_answer_1"), config) + assert resume_1["__interrupt__"][0].value == "Question 2?" + assert "ask_1" in called + + called.clear() + resume_2 = await graph.ainvoke(Command(resume="new_answer_2"), config) + assert resume_2["value"] == [ + "step_a_done", + "ask_1:new_answer_1", + "ask_2:new_answer_2", + ] + + # Verify final history: original branch preserved, new branch appended + final = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from time travel fork) + ( + "loop", + (), + {"value": ["step_a_done", "ask_1:new_answer_1", "ask_2:new_answer_2"]}, + ), + ("fork", ("executor",), {"value": []}), + # Original branch (preserved) + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + +@NEEDS_CONTEXTVARS +async def test_subgraph_time_travel_resume_from_second_interrupt_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Time travel to a subgraph checkpoint at the second interrupt, then + resume with a new answer. The first interrupt's answer should be preserved. + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END + """ + + called: list[str] = [] + + async def step_a(state: State) -> State: + called.append("step_a") + return {"value": ["step_a_done"]} + + async def ask_1(state: State) -> State: + called.append("ask_1") + answer = interrupt("Question 1?") + return {"value": [f"ask_1:{answer}"]} + + async def ask_2(state: State) -> State: + called.append("ask_2") + answer = interrupt("Question 2?") + return {"value": [f"ask_2:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("step_a", step_a) + .add_node("ask_1", ask_1) + .add_node("ask_2", ask_2) + .add_edge(START, "step_a") + .add_edge("step_a", "ask_1") + .add_edge("ask_1", "ask_2") + .add_edge("ask_2", "__end__") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # --- Original run: hit both interrupts and resume --- + await graph.ainvoke({"value": []}, config) + await graph.ainvoke(Command(resume="answer_1"), config) + sub_config_at_second = ( + (await graph.aget_state(config, subgraphs=True)).tasks[0].state.config + ) + await graph.ainvoke(Command(resume="answer_2"), config) + + original = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # --- Time travel to second interrupt --- + called.clear() + replay_result = await graph.ainvoke(None, sub_config_at_second) + assert replay_result["__interrupt__"][0].value == "Question 2?" + assert "step_a" not in called + assert "ask_1" not in called + + # Fork is now the latest parent checkpoint + post_tt = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"]) for s in post_tt] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), # original done + ("loop", ("executor",)), + ("input", ("__start__",)), + ] + + # --- Resume with a new answer for ask_2 only --- + called.clear() + resume_result = await graph.ainvoke(Command(resume="new_answer_2"), config) + assert resume_result["value"] == [ + "step_a_done", + "ask_1:answer_1", + "ask_2:new_answer_2", + ] + + # Verify final history + final = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch (from time travel fork) + ( + "loop", + (), + {"value": ["step_a_done", "ask_1:answer_1", "ask_2:new_answer_2"]}, + ), + ("fork", ("executor",), {"value": []}), + # Original branch (preserved) + ("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + +@NEEDS_CONTEXTVARS +async def test_subgraph_time_travel_checkpoint_pattern_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Verify the checkpoint pattern created by time travel to a subgraph + interrupt. A fork checkpoint should branch from the replay point. + + Parent: START --> executor (subgraph, checkpointer=True) --> END + Executor: START --> ask (interrupt) --> END + """ + + async def ask(state: State) -> State: + answer = interrupt("Q?") + return {"value": [f"a:{answer}"]} + + executor = ( + StateGraph(State) + .add_node("ask", ask) + .add_edge(START, "ask") + .compile(checkpointer=True) + ) + + graph = ( + StateGraph(State) + .add_node("executor", executor) + .add_edge(START, "executor") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # Run until interrupt, then complete + await graph.ainvoke({"value": []}, config) + sub_config = (await graph.aget_state(config, subgraphs=True)).tasks[0].state.config + await graph.ainvoke(Command(resume="first"), config) + + original = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in original] == [ + ("loop", (), {"value": ["a:first"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + # Time travel to the interrupt + await graph.ainvoke(None, sub_config) + + # Fork is now the latest, branching from the original replay point + post_tt = [s async for s in graph.aget_state_history(config)] + post_tt_summary = _checkpoint_summary(post_tt) + assert [(s["source"], s["next"]) for s in post_tt_summary] == [ + ("fork", ("executor",)), # <-- new fork (latest) + ("loop", ()), + ("loop", ("executor",)), # <-- replay point / fork parent + ("input", ("__start__",)), + ] + # Verify the fork's parent is the original replay point + replay_point_id = sub_config["configurable"]["checkpoint_map"][""] + assert post_tt[0].parent_config["configurable"]["checkpoint_id"] == replay_point_id + + # Resume from the fork + result = await graph.ainvoke(Command(resume="second"), config) + assert result["value"] == ["a:second"] + + final = _checkpoint_summary([s async for s in graph.aget_state_history(config)]) + assert [(s["source"], s["next"], s["values"]) for s in final] == [ + # New branch + ("loop", (), {"value": ["a:second"]}), + ("fork", ("executor",), {"value": []}), + # Original branch + ("loop", (), {"value": ["a:first"]}), + ("loop", ("executor",), {"value": []}), + ("input", ("__start__",), {"value": []}), + ] + + @NEEDS_CONTEXTVARS async def test_3_levels_deep_time_travel_to_first_interrupt_async( async_checkpointer: BaseCheckpointSaver, @@ -2088,13 +2481,15 @@ async def test_replay_creates_branch_preserving_old_checkpoints( # -- Post-replay checkpoint history (newest first) -- post_replay_history = [s async for s in graph.aget_state_history(config)] post_summary = _checkpoint_summary(post_replay_history) - assert len(post_summary) == 7 # 5 original + 2 new branch checkpoints + # 5 original + 1 fork + 2 new branch checkpoints = 8 + assert len(post_summary) == 8 assert [s["next"] for s in post_summary] == [ - (), # new branch tip (C6) - ("node_c",), # new branch (C5) - (), # old branch tip (C4) - ("node_c",), # old (C3) + (), # new branch tip + ("node_c",), # new branch + ("node_b",), # fork from replay point + (), # old branch tip + ("node_c",), # old ("node_b",), # branch point (C2) ("node_a",), # old (C1) ("__start__",), # old (C0) @@ -2102,6 +2497,7 @@ async def test_replay_creates_branch_preserving_old_checkpoints( assert [s["values"] for s in post_summary] == [ {"value": ["a", "b2", "c"]}, # new branch tip {"value": ["a", "b2"]}, # new: node_b re-ran with call_count=2 + {"value": ["a"]}, # fork from replay point {"value": ["a", "b1", "c"]}, # old branch tip preserved {"value": ["a", "b1"]}, # old {"value": ["a"]}, # branch point