"""Tests for time travel (replay and fork) behavior. Covers the intersection of replay vs fork across graph structures: - Replay & fork basics (no interrupt, no subgraph) - Replay & fork with interrupts (no subgraph) - Multiple / sequential interrupts - Subgraph without interrupt - Subgraph with interrupt - __copy__ / update_state(None) - Observability (get_state, config access) Key concepts: - Replay (invoke with checkpoint_id): Re-executes nodes after the checkpoint. Interrupts re-fire on replay. - Fork (update_state then invoke): Creates a new checkpoint without cached pending writes. Nodes re-execute and interrupts DO re-fire. """ import operator from typing import Annotated from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import BaseCheckpointSaver from typing_extensions import TypedDict from langgraph.graph import START, StateGraph from langgraph.types import Command, interrupt class State(TypedDict): value: Annotated[list[str], operator.add] def _checkpoint_summary(history: list) -> list[dict]: """Summarize checkpoint history into a readable format for assertions. 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 """ summaries = [] for s in history: cid = s.config["configurable"]["checkpoint_id"] pid = ( s.parent_config["configurable"]["checkpoint_id"] if s.parent_config else None ) summaries.append( { "id": cid[-6:], "parent_id": pid[-6:] if pid else None, "source": s.metadata.get("source"), "next": s.next, "values": s.values, } ) return summaries # --------------------------------------------------------------------------- # Section 1: Replay & fork basics (no interrupt, no subgraph) # --------------------------------------------------------------------------- def test_replay_reruns_nodes_after_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from checkpoint before node_b. node_b re-executes (it's after the checkpoint), node_a does not.""" called: list[str] = [] def node_a(state: State) -> State: called.append("node_a") return {"value": ["a"]} def node_b(state: State) -> State: called.append("node_b") return {"value": ["b"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": []}, config) assert result == {"value": ["a", "b"]} assert called == ["node_a", "node_b"] # Find checkpoint before node_b (next=(node_b,)) history = list(graph.get_state_history(config)) before_b = next(s for s in history if s.next == ("node_b",)) # Replay from checkpoint before node_b called.clear() replay_result = graph.invoke(None, before_b.config) assert replay_result == {"value": ["a", "b"]} assert "node_b" in called assert "node_a" not in called def test_replay_from_final_checkpoint_is_noop( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from completed checkpoint (no next nodes) is a no-op.""" called: list[str] = [] def node_a(state: State) -> State: called.append("node_a") return {"value": ["a"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_edge(START, "node_a") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": []}, config) assert result == {"value": ["a"]} state = graph.get_state(config) assert state.next == () called.clear() replay_result = graph.invoke(None, state.config) assert replay_result == {"value": ["a"]} assert called == [] def test_fork_reruns_with_modified_state( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Fork from checkpoint before node_b with modified state. node_b re-executes with the new state.""" called: list[str] = [] def node_a(state: State) -> State: called.append("node_a") return {"value": ["a"]} def node_b(state: State) -> State: called.append("node_b") return {"value": ["b"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"value": []}, config) history = list(graph.get_state_history(config)) before_b = next(s for s in history if s.next == ("node_b",)) called.clear() fork_config = graph.update_state(before_b.config, {"value": ["x"]}) fork_result = graph.invoke(None, fork_config) assert "node_b" in called assert fork_result == {"value": ["a", "x", "b"]} def test_multiple_forks_from_same_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Two independent forks from the same checkpoint create independent branches that don't affect each other.""" def node_a(state: State) -> State: return {"value": ["a"]} def node_b(state: State) -> State: return {"value": ["b"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"value": []}, config) history = list(graph.get_state_history(config)) before_b = next(s for s in history if s.next == ("node_b",)) fork1_config = graph.update_state(before_b.config, {"value": ["fork1"]}) result1 = graph.invoke(None, fork1_config) fork2_config = graph.update_state(before_b.config, {"value": ["fork2"]}) result2 = graph.invoke(None, fork2_config) assert "fork1" in result1["value"] assert "fork2" not in result1["value"] assert "fork2" in result2["value"] assert "fork1" not in result2["value"] # --------------------------------------------------------------------------- # Section 2: Replay & fork with interrupts (no subgraph) # --------------------------------------------------------------------------- def test_replay_from_before_interrupt_refires( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from checkpoint before interrupt node. The node re-executes and interrupt re-fires.""" call_count: dict[str, int] = {"node_a": 0, "ask_human": 0, "node_b": 0} def node_a(state: State) -> State: call_count["node_a"] += 1 return {"value": ["a"]} def ask_human(state: State) -> State: call_count["ask_human"] += 1 answer = interrupt("What is your input?") return {"value": [f"human:{answer}"]} def node_b(state: State) -> State: call_count["node_b"] += 1 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"}} # Run until interrupt result = graph.invoke({"value": []}, config) assert "__interrupt__" in result assert call_count["ask_human"] == 1 # Resume result = graph.invoke(Command(resume="hello"), config) assert result == {"value": ["a", "human:hello", "b"]} assert call_count["ask_human"] == 2 # re-executes on resume # Find checkpoint before ask_human history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask_human",)][-1] # Replay — interrupt re-fires, node re-executes replay_result = graph.invoke(None, before_ask.config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "What is your input?" assert replay_result["value"] == ["a"] assert call_count["ask_human"] == 3 # re-executed again assert call_count["node_a"] == 1 # NOT re-executed (before checkpoint) 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: """Replaying the same checkpoint multiple times consistently produces identical results (interrupt re-fires each time).""" def node_a(state: State) -> State: return {"value": ["a"]} def ask_human(state: State) -> State: answer = interrupt("What is your input?") return {"value": [f"human:{answer}"]} def node_b(state: State) -> State: 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"}} graph.invoke({"value": []}, config) graph.invoke(Command(resume="cached_answer"), config) history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask_human",)][-1] results = [] for _ in range(3): r = graph.invoke(None, before_ask.config) results.append(r) # 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( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Fork from checkpoint before interrupt node. Interrupt IS re-triggered because fork has no cached resume values. Resume with new answer.""" 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"}} graph.invoke({"value": []}, config) graph.invoke(Command(resume="hello"), config) history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask_human",)][-1] called.clear() fork_config = graph.update_state(before_ask.config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "What is your input?" # Resume the forked interrupt with a different answer final = graph.invoke(Command(resume="world"), fork_config) assert final == {"value": ["a", "forked", "human:world", "b"]} def test_fork_from_interrupt_checkpoint_refires( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Fork from the checkpoint where interrupt fired. Interrupt re-triggered because fork clears cached data. Resume with different answer.""" 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"}} graph.invoke({"value": []}, config) graph.invoke(Command(resume="hello"), config) history = list(graph.get_state_history(config)) interrupt_checkpoint = next( s for s in history if s.next == ("ask_human",) and s.tasks and any(t.interrupts for t in s.tasks) ) fork_config = graph.update_state(interrupt_checkpoint.config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result final = graph.invoke(Command(resume="different"), fork_config) assert "human:different" in final["value"] # --------------------------------------------------------------------------- # Section 3: Multiple / sequential interrupts # --------------------------------------------------------------------------- def test_sequential_interrupts_fork_from_middle( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Graph with two sequential interrupt nodes. Fork from between them: only the second re-fires, the first's result is preserved. Also verify replaying from before the first re-fires only the first.""" called: list[str] = [] def node_a(state: State) -> State: called.append("node_a") return {"value": ["a"]} def interrupt_1(state: State) -> State: called.append("interrupt_1") answer = interrupt("First question?") return {"value": [f"i1:{answer}"]} def interrupt_2(state: State) -> State: called.append("interrupt_2") answer = interrupt("Second question?") return {"value": [f"i2:{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("interrupt_1", interrupt_1) .add_node("interrupt_2", interrupt_2) .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "interrupt_1") .add_edge("interrupt_1", "interrupt_2") .add_edge("interrupt_2", "node_b") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Hit first interrupt r1 = graph.invoke({"value": []}, config) assert r1["__interrupt__"][0].value == "First question?" # Resume first → hit second r2 = graph.invoke(Command(resume="ans1"), config) assert r2["__interrupt__"][0].value == "Second question?" # Resume second → complete r3 = graph.invoke(Command(resume="ans2"), config) assert r3 == {"value": ["a", "i1:ans1", "i2:ans2", "b"]} history = list(graph.get_state_history(config)) # Fork from between the two interrupts — only second re-fires between = [s for s in history if s.next == ("interrupt_2",)][-1] fork_config = graph.update_state(between.config, {"value": ["mid_fork"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Second question?" # Resume with new answer, first answer preserved final_result = graph.invoke(Command(resume="new_b"), fork_config) assert "i2:new_b" in final_result["value"] assert "i1:ans1" in final_result["value"] # Replay from before first interrupt — first re-fires, second does not before_i1 = [s for s in history if s.next == ("interrupt_1",)][-1] called.clear() replay_result = graph.invoke(None, before_i1.config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "First question?" assert "interrupt_1" in called assert "interrupt_2" not in called def test_multiple_interrupts_in_one_node( sync_checkpointer: BaseCheckpointSaver, ) -> None: """A single node with two sequential interrupt() calls. Resuming resolves them one at a time. Replaying from before the node re-fires the first.""" def multi_interrupt_node(state: State) -> State: answer1 = interrupt("First question?") answer2 = interrupt("Second question?") return {"value": [f"a1:{answer1}", f"a2:{answer2}"]} def after(state: State) -> State: return {"value": ["done"]} graph = ( StateGraph(State) .add_node("ask", multi_interrupt_node) .add_node("after", after) .add_edge(START, "ask") .add_edge("ask", "after") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Hit first interrupt result = graph.invoke({"value": []}, config) assert result["__interrupt__"][0].value == "First question?" # Resume first → hit second interrupt_state = graph.get_state(config) result = graph.invoke(Command(resume="ans1"), interrupt_state.config) assert result["__interrupt__"][0].value == "Second question?" # Resume second → complete interrupt_state2 = graph.get_state(config) result = graph.invoke(Command(resume="ans2"), interrupt_state2.config) assert result == {"value": ["a1:ans1", "a2:ans2", "done"]} # Replay from before the node — first interrupt re-fires history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask",)][-1] replay_result = graph.invoke(None, before_ask.config) assert replay_result["__interrupt__"][0].value == "First question?" # --------------------------------------------------------------------------- # Section 4: Subgraph without interrupt # --------------------------------------------------------------------------- def test_subgraph_replay_from_before( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from parent checkpoint before subgraph node. Subgraph and post_process re-execute, parent_node does not.""" called: list[str] = [] def parent_node(state: State) -> State: called.append("parent_node") return {"value": ["parent"]} def step_a(state: State) -> State: called.append("step_a") return {"value": ["sub_a"]} 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("step_b", step_b) .add_edge(START, "step_a") .add_edge("step_a", "step_b") .compile() ) def post_process(state: State) -> State: called.append("post_process") return {"value": ["post"]} graph = ( StateGraph(State) .add_node("parent_node", parent_node) .add_node("subgraph", subgraph) .add_node("post_process", post_process) .add_edge(START, "parent_node") .add_edge("parent_node", "subgraph") .add_edge("subgraph", "post_process") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": []}, config) assert "sub_a" in result["value"] assert "sub_b" in result["value"] assert "post" in result["value"] history = list(graph.get_state_history(config)) before_sub = next(s for s in history if s.next == ("subgraph",)) called.clear() replay_result = graph.invoke(None, before_sub.config) assert "sub_a" in replay_result["value"] assert "sub_b" in replay_result["value"] assert "post" in replay_result["value"] assert "parent_node" not in called def test_subgraph_fork_from_before( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Fork from parent checkpoint before subgraph with modified state. Subgraph re-executes with forked state.""" called: list[str] = [] def parent_node(state: State) -> State: called.append("parent_node") return {"value": ["parent"]} def step_a(state: State) -> State: called.append("step_a") return {"value": ["sub_a"]} 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("step_b", step_b) .add_edge(START, "step_a") .add_edge("step_a", "step_b") .compile() ) def post_process(state: State) -> State: called.append("post_process") return {"value": ["post"]} graph = ( StateGraph(State) .add_node("parent_node", parent_node) .add_node("subgraph", subgraph) .add_node("post_process", post_process) .add_edge(START, "parent_node") .add_edge("parent_node", "subgraph") .add_edge("subgraph", "post_process") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"value": []}, config) history = list(graph.get_state_history(config)) before_sub = next(s for s in history if s.next == ("subgraph",)) called.clear() fork_config = graph.update_state(before_sub.config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "step_a" in called assert "step_b" in called assert "post_process" in called assert "forked" in fork_result["value"] assert "sub_a" in fork_result["value"] assert "sub_b" in fork_result["value"] # --------------------------------------------------------------------------- # Section 5: Subgraph with interrupt # --------------------------------------------------------------------------- def test_subgraph_interrupt_replay_from_parent( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from parent checkpoint before subgraph. Subgraph re-executes and interrupt re-fires.""" 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 result = graph.invoke({"value": []}, config) assert "__interrupt__" in result completed_result = graph.invoke(Command(resume="answer"), config) assert "human:answer" in completed_result["value"] assert "sub_b" in completed_result["value"] assert "post" in completed_result["value"] # Find parent checkpoint before subgraph_node history = list(graph.get_state_history(config)) before_sub = [s for s in history if s.next == ("subgraph_node",)][-1] # Replay from before subgraph — subgraph starts fresh, interrupt re-fires called.clear() replay_result = graph.invoke(None, before_sub.config) assert "__interrupt__" in replay_result # Subgraph ran from scratch (step_a and ask_human called) assert "step_a" in called assert "ask_human" in called # step_b should NOT be called (interrupt stops execution) assert "step_b" not in called def test_subgraph_interrupt_replay_from_interrupt_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay from the parent checkpoint where subgraph interrupt fired. Interrupt re-fires.""" 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 graph.invoke({"value": []}, config) # Verify subgraph state is accessible parent_state = graph.get_state(config, subgraphs=True) assert len(parent_state.tasks) > 0 assert parent_state.tasks[0].state is not None # Resume graph.invoke(Command(resume="answer"), config) # Find the parent checkpoint where the interrupt fired history = list(graph.get_state_history(config)) interrupt_checkpoint = next( s for s in history if s.next == ("subgraph_node",) and s.tasks and any(t.interrupts for t in s.tasks) ) called.clear() replay_result = graph.invoke(None, interrupt_checkpoint.config) assert "__interrupt__" in replay_result # Subgraph starts fresh during replay — all nodes re-run from scratch. # step_a re-runs, ask_human re-fires interrupt, step_b not reached. assert "step_a" in called assert "ask_human" in called 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_interrupt_resume_with_explicit_head_checkpoint_id( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Resume with Command(resume=...) plus the current head checkpoint_id in config. The subgraph must continue from the interrupted node, not restart from scratch. Explicit checkpoint_id triggers is_replaying but this is a resume, not a time-travel, so ReplayState should not apply.""" called: list[str] = [] 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) ) graph = ( StateGraph(State) .add_node("subgraph_node", subgraph) .add_edge(START, "subgraph_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run until interrupt fires in subgraph graph.invoke({"value": []}, config) assert called == ["step_a", "ask_human"] # Resume with explicit head checkpoint_id in config head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"] called.clear() resume_config = { "configurable": { "thread_id": "1", "checkpoint_id": head_checkpoint_id, "checkpoint_ns": "", } } result = graph.invoke(Command(resume="answer"), resume_config) assert called == ["ask_human", "step_b"] assert "__interrupt__" not in result assert result["value"] == ["sub_a", "human:answer", "sub_b"] 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: """Fork from the subgraph's own checkpoint, resume the interrupt with a new answer, and verify the FULL flow: subgraph completes (step_b runs) AND execution continues back to the parent graph (post_process runs). This is the key test for time-traveling to a subgraph checkpoint, re-triggering the interrupt, providing a new answer, and having the entire graph complete normally including parent nodes after the subgraph.""" 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 fires in subgraph result = graph.invoke({"value": []}, config) assert "__interrupt__" in result assert "router" in called assert "step_a" in called assert "ask_human" in called # Get subgraph's own checkpoint config parent_state = graph.get_state(config, subgraphs=True) sub_task = parent_state.tasks[0] assert sub_task.state is not None sub_config = sub_task.state.config # Fork from subgraph checkpoint called.clear() fork_config = graph.update_state(sub_config, {"value": ["sub_forked"]}) # Invoke from fork — interrupt should re-fire fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Provide input:" # Resume the re-triggered interrupt with a NEW answer called.clear() final_result = graph.invoke(Command(resume="new_answer"), fork_config) # Verify full completion assert "ask_human" in called assert "human:new_answer" in final_result["value"] assert "step_b" in called assert "sub_b" in final_result["value"] assert "post_process" in called assert "post" in final_result["value"] def test_subgraph_interrupt_full_flow_no_sub_checkpointer( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Same as test_subgraph_interrupt_full_flow but with no sub-checkpointer (checkpointer=None). Fork from the parent checkpoint before the subgraph, re-trigger interrupt, resume, and verify full parent completion.""" 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=None) ) 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 result = graph.invoke({"value": []}, config) assert "__interrupt__" in result original_result = graph.invoke(Command(resume="original"), config) assert "human:original" in original_result["value"] assert "post" in original_result["value"] # Find parent checkpoint before subgraph_node history = list(graph.get_state_history(config)) before_sub = [s for s in history if s.next == ("subgraph_node",)][-1] # Fork from parent checkpoint called.clear() fork_config = graph.update_state(before_sub.config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) # Interrupt IS re-triggered assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Provide input:" # Resume with new answer called.clear() final_result = graph.invoke(Command(resume="new_answer"), fork_config) # Verify full completion back through parent assert "human:new_answer" in final_result["value"] assert "step_b" in called assert "sub_b" in final_result["value"] assert "post_process" in called assert "post" in final_result["value"] def test_subgraph_replay_from_subgraph_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replay directly from a subgraph's own checkpoint (via get_state with subgraphs=True). The subgraph resumes from its checkpoint and the parent graph completes normally afterwards.""" 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 fires in subgraph result = graph.invoke({"value": []}, config) assert "__interrupt__" in result # Get subgraph checkpoint config (without forking) parent_state = graph.get_state(config, subgraphs=True) sub_task = parent_state.tasks[0] assert sub_task.state is not None sub_config = sub_task.state.config # Replay directly from subgraph checkpoint (no update_state / no fork) called.clear() replay_result = graph.invoke(None, sub_config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Provide input:" # Resume from the replayed checkpoint called.clear() final_result = graph.invoke(Command(resume="replayed_answer"), sub_config) assert "ask_human" in called assert "human:replayed_answer" in final_result["value"] assert "step_b" in called assert "sub_b" in final_result["value"] assert "post_process" in called assert "post" in final_result["value"] def test_subgraph_time_travel_to_first_interrupt( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel to a subgraph checkpoint at the FIRST interrupt. Architecture: Parent: START --> executor (subgraph, checkpointer=True) --> END Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END Flow: run through both interrupts, then time travel back to the subgraph checkpoint captured at the first interrupt. ask_1 should re-fire, step_a should NOT re-run. Then resume through both interrupts with new answers. """ 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"}} # Run until first interrupt (ask_1) result = graph.invoke({"value": []}, config) assert "__interrupt__" in result assert result["__interrupt__"][0].value == "Question 1?" # Capture subgraph state at the first interrupt parent_state = graph.get_state(config, subgraphs=True) sub_config_at_first = parent_state.tasks[0].state.config # Resume first interrupt result = graph.invoke(Command(resume="answer_1"), config) assert result["__interrupt__"][0].value == "Question 2?" # Resume second interrupt to complete result = graph.invoke(Command(resume="answer_2"), config) assert "__interrupt__" not in result # --- Scenario 1: Replay from subgraph checkpoint at 1st interrupt --- called.clear() replay_result = graph.invoke(None, sub_config_at_first) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Question 1?" # step_a should NOT re-run — it was before this checkpoint assert "step_a" not in called # ask_1 re-fires because the interrupt replays assert "ask_1" in called # --- Scenario 2: Fork from subgraph checkpoint at 1st interrupt --- called.clear() fork_config = graph.update_state(sub_config_at_first, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Question 1?" assert "step_a" not in called assert "ask_1" in called def test_subgraph_time_travel_to_second_interrupt( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel to a subgraph checkpoint at the SECOND interrupt. Architecture: Parent: START --> executor (subgraph, checkpointer=True) --> END Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END Flow: run through both interrupts resuming each, then time travel back to the subgraph checkpoint at the second interrupt. Only ask_2 should re-fire. Then resume with a new answer and verify state. """ 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"}} # Run until first interrupt (ask_1) result = graph.invoke({"value": []}, config) assert result["__interrupt__"][0].value == "Question 1?" # Resume first interrupt result = graph.invoke(Command(resume="answer_1"), config) assert result["__interrupt__"][0].value == "Question 2?" # Capture subgraph state at the second interrupt parent_state = graph.get_state(config, subgraphs=True) sub_config = parent_state.tasks[0].state.config # Resume second interrupt to complete the graph result = graph.invoke(Command(resume="answer_2"), config) assert "__interrupt__" not in result # --- Scenario 1: Replay from subgraph checkpoint at 2nd interrupt --- called.clear() replay_result = graph.invoke(None, sub_config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Question 2?" # step_a and ask_1 should NOT re-run — they were before this checkpoint assert "step_a" not in called assert "ask_1" not in called # --- Scenario 2: Fork from subgraph checkpoint at 2nd interrupt --- called.clear() fork_config = graph.update_state(sub_config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Question 2?" assert "step_a" not in called 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: """Time travel to a subgraph checkpoint AFTER both interrupts are resolved. Architecture: Parent: START --> executor (subgraph, checkpointer=True) --> END Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END After completing the full flow, capture the subgraph's final state checkpoint and replay from it — should be a no-op (no nodes re-run). """ 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"}} # Run through both interrupts graph.invoke({"value": []}, config) graph.invoke(Command(resume="answer_1"), config) # Before resuming 2nd interrupt, get state history to find the # subgraph checkpoint that will exist after ask_2 completes graph.invoke(Command(resume="answer_2"), config) # Get the final parent state — no pending tasks final_state = graph.get_state(config) assert len(final_state.tasks) == 0 # Replay from the final parent checkpoint — should be a no-op called.clear() replay_result = graph.invoke(None, final_state.config) assert "__interrupt__" not in replay_result assert "step_a" not in called assert "ask_1" not in called assert "ask_2" not in called # All values should be present assert "step_a_done" in replay_result["value"] assert "ask_1:answer_1" in replay_result["value"] assert "ask_2:answer_2" in replay_result["value"] def test_3_levels_deep_time_travel_to_first_interrupt( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel to the innermost subgraph checkpoint at the FIRST interrupt. Architecture: Parent: START --> outer (subgraph, checkpointer=True) --> END Outer: START --> inner (subgraph, checkpointer=True) --> END Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END """ 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}"]} inner = ( 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) ) middle = ( StateGraph(State) .add_node("inner", inner) .add_edge(START, "inner") .compile(checkpointer=True) ) graph = ( StateGraph(State) .add_node("outer", middle) .add_edge(START, "outer") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run until first interrupt result = graph.invoke({"value": []}, config) assert result["__interrupt__"][0].value == "Question 1?" # Capture innermost subgraph state at the first interrupt parent_state = graph.get_state(config, subgraphs=True) mid_state = parent_state.tasks[0].state inner_config = mid_state.tasks[0].state.config # Resume through both interrupts to complete graph.invoke(Command(resume="answer_1"), config) graph.invoke(Command(resume="answer_2"), config) # --- Scenario 1: Replay from innermost checkpoint at 1st interrupt --- called.clear() replay_result = graph.invoke(None, inner_config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Question 1?" assert "step_a" not in called assert "ask_1" in called # --- Scenario 2: Fork from innermost checkpoint at 1st interrupt --- called.clear() fork_config = graph.update_state(inner_config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Question 1?" assert "step_a" not in called assert "ask_1" in called def test_3_levels_deep_time_travel_to_second_interrupt( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel to the innermost subgraph checkpoint at the SECOND interrupt. Architecture: Parent: START --> outer (subgraph, checkpointer=True) --> END Outer: START --> inner (subgraph, checkpointer=True) --> END Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END """ 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}"]} inner = ( 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) ) middle = ( StateGraph(State) .add_node("inner", inner) .add_edge(START, "inner") .compile(checkpointer=True) ) graph = ( StateGraph(State) .add_node("outer", middle) .add_edge(START, "outer") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run until first interrupt graph.invoke({"value": []}, config) # Resume first interrupt result = graph.invoke(Command(resume="answer_1"), config) assert result["__interrupt__"][0].value == "Question 2?" # Capture innermost subgraph state at the second interrupt parent_state = graph.get_state(config, subgraphs=True) mid_state = parent_state.tasks[0].state inner_config = mid_state.tasks[0].state.config # Resume second interrupt to complete graph.invoke(Command(resume="answer_2"), config) # --- Scenario 1: Replay from innermost checkpoint at 2nd interrupt --- called.clear() replay_result = graph.invoke(None, inner_config) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Question 2?" assert "step_a" not in called assert "ask_1" not in called # --- Scenario 2: Fork from innermost checkpoint at 2nd interrupt --- called.clear() fork_config = graph.update_state(inner_config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Question 2?" assert "step_a" not in called assert "ask_1" not in called def test_3_levels_deep_time_travel_to_middle_subgraph( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel to the MIDDLE-level subgraph checkpoint (not innermost). Architecture: Parent: START --> outer (subgraph, checkpointer=True) --> END Outer: START --> inner (subgraph, checkpointer=True) --> END Inner: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END After completing the full flow, time travel back to the middle subgraph's checkpoint at the second interrupt. The middle subgraph should replay the inner subgraph from the correct point. """ 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}"]} inner = ( 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) ) middle = ( StateGraph(State) .add_node("inner", inner) .add_edge(START, "inner") .compile(checkpointer=True) ) graph = ( StateGraph(State) .add_node("outer", middle) .add_edge(START, "outer") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run until first interrupt graph.invoke({"value": []}, config) # Resume first, capture middle config at second interrupt graph.invoke(Command(resume="answer_1"), config) parent_state = graph.get_state(config, subgraphs=True) mid_config = parent_state.tasks[0].state.config # Resume second to complete graph.invoke(Command(resume="answer_2"), config) # --- Scenario 1: Replay from middle-level subgraph checkpoint --- # The middle subgraph's checkpoint knows about the inner subgraph's state # via checkpoint_map, so the inner replays from the correct point. called.clear() replay_result = graph.invoke(None, mid_config) assert "__interrupt__" in replay_result # --- Scenario 2: Fork from middle-level subgraph checkpoint --- called.clear() fork_config = graph.update_state(mid_config, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result def test_3_levels_deep_middle_has_interrupts( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Time travel when the MIDDLE subgraph itself has interrupts. Architecture: Parent: START --> outer (subgraph, checkpointer=True) --> END Outer: START --> pre (interrupt) --> inner (subgraph, checkpointer=True) --> END Inner: START --> step_a --> ask_1 (interrupt) --> END Flow: run through both interrupts (pre then ask_1), then time travel back to the middle subgraph checkpoint at each interrupt point. """ called: list[str] = [] def pre(state: State) -> State: called.append("pre") answer = interrupt("Pre-question?") return {"value": [f"pre:{answer}"]} 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}"]} inner = ( StateGraph(State) .add_node("step_a", step_a) .add_node("ask_1", ask_1) .add_edge(START, "step_a") .add_edge("step_a", "ask_1") .add_edge("ask_1", "__end__") .compile(checkpointer=True) ) middle = ( StateGraph(State) .add_node("pre", pre) .add_node("inner", inner) .add_edge(START, "pre") .add_edge("pre", "inner") .add_edge("inner", "__end__") .compile(checkpointer=True) ) graph = ( StateGraph(State) .add_node("outer", middle) .add_edge(START, "outer") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run until first interrupt (pre in middle subgraph) result = graph.invoke({"value": []}, config) assert result["__interrupt__"][0].value == "Pre-question?" # Capture middle subgraph config at the pre interrupt parent_state = graph.get_state(config, subgraphs=True) mid_config_at_pre = parent_state.tasks[0].state.config # Resume pre, hits ask_1 in inner subgraph result = graph.invoke(Command(resume="pre_answer"), config) assert result["__interrupt__"][0].value == "Question 1?" # Capture middle subgraph config at the ask_1 interrupt parent_state = graph.get_state(config, subgraphs=True) mid_config_at_ask1 = parent_state.tasks[0].state.config # Resume ask_1 to complete result = graph.invoke(Command(resume="answer_1"), config) assert "__interrupt__" not in result # --- Time travel to middle checkpoint at pre interrupt --- called.clear() replay_result = graph.invoke(None, mid_config_at_pre) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Pre-question?" # pre should re-fire (interrupt replays), but nothing else should run assert "pre" in called assert "step_a" not in called assert "ask_1" not in called # Fork from middle checkpoint at pre interrupt called.clear() fork_config = graph.update_state(mid_config_at_pre, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Pre-question?" assert "pre" in called assert "step_a" not in called # --- Time travel to middle checkpoint at ask_1 interrupt --- called.clear() replay_result = graph.invoke(None, mid_config_at_ask1) assert "__interrupt__" in replay_result assert replay_result["__interrupt__"][0].value == "Question 1?" # pre should NOT re-run (it completed before this checkpoint) assert "pre" not in called # ask_1 re-fires assert "ask_1" in called # Fork from middle checkpoint at ask_1 interrupt called.clear() fork_config = graph.update_state(mid_config_at_ask1, {"value": ["forked"]}) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "Question 1?" assert "pre" not in called assert "ask_1" in called # --------------------------------------------------------------------------- # Section 6: __copy__ / update_state(None) # --------------------------------------------------------------------------- def test_copy_fork_retriggers_interrupt( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Fork using __copy__ (no state changes) from checkpoint before interrupt. The interrupt is re-triggered because __copy__ creates a new checkpoint without cached resume values. Resume with new answer to verify.""" 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"}} graph.invoke({"value": []}, config) graph.invoke(Command(resume="hello"), config) history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask_human",)][-1] called.clear() fork_config = graph.update_state(before_ask.config, None, as_node="__copy__") fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "What is your input?" final = graph.invoke(Command(resume="new_answer"), fork_config) assert final == {"value": ["a", "human:new_answer", "b"]} def test_copy_fork_creates_sibling_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """__copy__ creates a checkpoint with source="fork", while regular update_state creates one with source="update".""" def node_a(state: State) -> State: return {"value": ["a"]} def node_b(state: State) -> State: return {"value": ["b"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"value": []}, config) history = list(graph.get_state_history(config)) before_b = next(s for s in history if s.next == ("node_b",)) # __copy__ fork → source="fork" copy_config = graph.update_state(before_b.config, None, as_node="__copy__") copy_state = graph.get_state(copy_config) assert copy_state.metadata["source"] == "fork" # Regular update → source="update" regular_config = graph.update_state(before_b.config, {"value": ["x"]}) regular_state = graph.get_state(regular_config) assert regular_state.metadata["source"] == "update" def test_update_state_with_none_values( sync_checkpointer: BaseCheckpointSaver, ) -> None: """update_state with None values (not __copy__) goes through the normal update path, creating a new checkpoint that re-triggers interrupts.""" def node_a(state: State) -> State: return {"value": ["a"]} def ask_human(state: State) -> State: answer = interrupt("What is your input?") return {"value": [f"human:{answer}"]} def node_b(state: State) -> State: 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"}} graph.invoke({"value": []}, config) graph.invoke(Command(resume="hello"), config) history = list(graph.get_state_history(config)) before_ask = [s for s in history if s.next == ("ask_human",)][-1] fork_config = graph.update_state(before_ask.config, None) fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result assert fork_result["__interrupt__"][0].value == "What is your input?" fork_state = graph.get_state(fork_config) assert fork_state.metadata["source"] == "update" # --------------------------------------------------------------------------- # Section 7: Observability (get_state, config access) # --------------------------------------------------------------------------- def test_get_state_with_subgraphs_returns_subgraph_state( sync_checkpointer: BaseCheckpointSaver, ) -> None: """get_state(config, subgraphs=True) returns subgraph state and checkpoint config when paused at interrupt.""" class SubState(TypedDict): data: str def sub_node(state: SubState) -> SubState: interrupt("Continue?") return {"data": "processed"} subgraph = ( StateGraph(SubState) .add_node("process", sub_node) .add_edge(START, "process") .compile() ) class ParentState(TypedDict): data: str graph = ( StateGraph(ParentState) .add_node("sub", subgraph) .add_edge(START, "sub") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"data": "input"}, config) state = graph.get_state(config, subgraphs=True) assert len(state.tasks) > 0 sub_task = state.tasks[0] assert sub_task.state is not None sub_config = sub_task.state.config assert "configurable" in sub_config assert "thread_id" in sub_config["configurable"] def test_checkpoint_ns_accessible_in_subgraph( sync_checkpointer: BaseCheckpointSaver, ) -> None: """RunnableConfig exposes checkpoint_ns and thread_id inside subgraph nodes.""" captured_config: dict = {} class SubState(TypedDict): data: str def sub_node(state: SubState, config: RunnableConfig) -> SubState: captured_config["checkpoint_ns"] = config["configurable"].get("checkpoint_ns") captured_config["thread_id"] = config["configurable"].get("thread_id") return {"data": "done"} subgraph = ( StateGraph(SubState) .add_node("inner", sub_node) .add_edge(START, "inner") .compile() ) class ParentState(TypedDict): data: str graph = ( StateGraph(ParentState) .add_node("outer", subgraph) .add_edge(START, "outer") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} graph.invoke({"data": "test"}, config) assert captured_config["checkpoint_ns"] is not None assert captured_config["checkpoint_ns"] != "" assert captured_config["thread_id"] == "1" # --------------------------------------------------------------------------- # Section 8: Stateful vs stateless subgraph state retention on replay # --------------------------------------------------------------------------- def test_stateful_subgraph_retains_state_on_parent_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Stateful subgraph (checkpointer=True) remembers accumulated state from prior invocations when the parent replays.""" started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_node(state: ParentState) -> ParentState: return {"results": ["p"]} def step_a(state: SubState) -> SubState: started.append(("step_a", dict(state))) answer = interrupt("question_a") observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: started.append(("step_b", dict(state))) answer = interrupt("question_b") observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( StateGraph(SubState) .add_node("step_a", step_a) .add_node("step_b", step_b) .add_edge(START, "step_a") .add_edge("step_a", "step_b") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_node", parent_node) .add_node("sub_node", sub) .add_edge(START, "parent_node") .add_edge("parent_node", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a1"), config) # hits step_b interrupt graph.invoke(Command(resume="b1"), config) # completes # step_a saw empty state (fresh subgraph) assert observed[0] == ("step_a", {"value": []}) # step_b saw step_a's answer assert observed[1] == ("step_b", {"value": ["a:a1"]}) # === 2nd invocation: answer "a2" and "b2" === observed.clear() graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a2"), config) # hits step_b interrupt graph.invoke(Command(resume="b2"), config) # completes # Stateful subgraph retained state from 1st invocation assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) assert observed[1] == ("step_b", {"value": ["a:a1", "b:b1", "a:a2"]}) # === Replay from checkpoint before sub_node in 2nd invocation === history = list(graph.get_state_history(config)) # History is newest-first, so first match = 2nd invocation before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] started.clear() replay = graph.invoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Replay sees 1st invocation's final state, NOT 2nd invocation's assert started[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) def test_stateful_subgraph_retains_state_on_parent_fork( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Stateful subgraph (checkpointer=True) remembers accumulated state from prior invocations when the parent forks.""" observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_node(state: ParentState) -> ParentState: return {"results": ["p"]} def step_a(state: SubState) -> SubState: observed.append(("step_a", dict(state))) answer = interrupt("question_a") return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: observed.append(("step_b", dict(state))) answer = interrupt("question_b") return {"value": [f"b:{answer}"]} sub = ( StateGraph(SubState) .add_node("step_a", step_a) .add_node("step_b", step_b) .add_edge(START, "step_a") .add_edge("step_a", "step_b") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_node", parent_node) .add_node("sub_node", sub) .add_edge(START, "parent_node") .add_edge("parent_node", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a1"), config) # hits step_b interrupt graph.invoke(Command(resume="b1"), config) # completes # === 2nd invocation: answer "a2" and "b2" === graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a2"), config) # hits step_b interrupt graph.invoke(Command(resume="b2"), config) # completes # === Fork from checkpoint before sub_node in 2nd invocation === history = list(graph.get_state_history(config)) before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] fork_config = graph.update_state(before_sub_2nd.config, {"results": ["forked"]}) observed.clear() fork_result = graph.invoke(None, fork_config) assert "__interrupt__" in fork_result # Fork sees 1st invocation's final state, NOT 2nd invocation's assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]}) def test_stateful_subgraph_loads_state_across_ticks( sync_checkpointer: BaseCheckpointSaver, ) -> None: """When replaying from a checkpoint before a non-subgraph node, a stateful subgraph that runs in a later tick should still load its accumulated state from the previous execution. Sequence: node_a -> node_b -> sub_node -> node_a -> node_b -> sub_node Replay from the 2nd node_a: sub_node in the later tick should see state from the end of the 1st sub_node execution. """ observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def node_a(state: ParentState) -> ParentState: return {"results": ["a"]} def node_b(state: ParentState) -> ParentState: return {"results": ["b"]} def sub_step(state: SubState) -> SubState: observed.append(("sub_step", dict(state))) return {"value": ["s"]} sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_node("sub_node", sub) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .add_edge("node_b", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation graph.invoke({"results": []}, config) assert observed[0] == ("sub_step", {"value": []}) # 2nd invocation observed.clear() graph.invoke({"results": []}, config) assert observed[0] == ("sub_step", {"value": ["s"]}) # Replay from checkpoint before node_a in 2nd invocation history = list(graph.get_state_history(config)) before_a_2nd = [s for s in history if s.next == ("node_a",)][0] observed.clear() graph.invoke(None, before_a_2nd.config) # sub_node runs in a later tick (after node_a and node_b replay), # and should see state from end of 1st sub_node execution assert observed[0] == ("sub_step", {"value": ["s"]}) # --------------------------------------------------------------------------- # Section 8: Append-only checkpoint history (branching / forking) # --------------------------------------------------------------------------- def test_replay_creates_branch_preserving_old_checkpoints( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replaying from a mid-run checkpoint creates a new branch of checkpoints while the original checkpoint sequence is preserved (append-only). Original run (newest first): C4 next=() values=[a, b1, c] parent=C3 C3 next=(node_c,) values=[a, b1] parent=C2 C2 next=(node_b,) values=[a] parent=C1 C1 next=(node_a,) values=[] parent=C0 C0 next=(__start__,) values={} parent=None After replay from C2 (newest first): C6 next=() values=[a, b2, c] parent=C5 <- new branch tip C5 next=(node_c,) values=[a, b2] parent=C2 <- branches from C2 C4 next=() values=[a, b1, c] parent=C3 <- old branch preserved C3 next=(node_c,) values=[a, b1] parent=C2 C2 next=(node_b,) values=[a] parent=C1 C1 next=(node_a,) values=[] parent=C0 C0 next=(__start__,) values={} parent=None """ call_count = 0 def node_a(state: State) -> State: return {"value": ["a"]} def node_b(state: State) -> State: nonlocal call_count call_count += 1 return {"value": [f"b{call_count}"]} def node_c(state: State) -> State: return {"value": ["c"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_node("node_c", node_c) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .add_edge("node_b", "node_c") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": []}, config) assert result == {"value": ["a", "b1", "c"]} # -- Original checkpoint history (newest first) -- original_history = list(graph.get_state_history(config)) original_summary = _checkpoint_summary(original_history) assert len(original_summary) == 5 # Verify the shape: next tuples newest->oldest assert [s["next"] for s in original_summary] == [ (), ("node_c",), ("node_b",), ("node_a",), ("__start__",), ] # Verify values at each checkpoint assert [s["values"] for s in original_summary] == [ {"value": ["a", "b1", "c"]}, {"value": ["a", "b1"]}, {"value": ["a"]}, {"value": []}, {"value": []}, ] original_ids = {s.config["configurable"]["checkpoint_id"] for s in original_history} # Find checkpoint before node_b and replay from it before_b = next(s for s in original_history if s.next == ("node_b",)) before_b_id = before_b.config["configurable"]["checkpoint_id"] replay_result = graph.invoke(None, before_b.config) assert replay_result == {"value": ["a", "b2", "c"]} # -- Post-replay checkpoint history (newest first) -- post_replay_history = list(graph.get_state_history(config)) post_summary = _checkpoint_summary(post_replay_history) # 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 ("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) ] 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 {"value": []}, # old {"value": []}, # old ] # All original checkpoint IDs still exist (append-only) post_ids = {s.config["configurable"]["checkpoint_id"] for s in post_replay_history} assert original_ids.issubset(post_ids) # New branch's oldest checkpoint parent is the branch point new_checkpoints = [ s for s in post_replay_history if s.config["configurable"]["checkpoint_id"] not in original_ids ] oldest_new = sorted(new_checkpoints, key=lambda s: s.created_at)[0] assert oldest_new.parent_config is not None assert oldest_new.parent_config["configurable"]["checkpoint_id"] == before_b_id # get_state returns the new branch tip latest = graph.get_state(config) assert latest.values == {"value": ["a", "b2", "c"]} assert latest.config["configurable"]["checkpoint_id"] not in original_ids def test_replay_creates_branch_in_subgraph( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replaying a graph with a subgraph from a mid-run checkpoint creates a new branch while preserving the original checkpoint sequence. The subgraph re-executes on the new branch and the old checkpoints (including sub-checkpoints) remain in the history. """ sub_call_count = 0 class SubState(TypedDict): sub_value: Annotated[list[str], operator.add] class ParentState(TypedDict): value: Annotated[list[str], operator.add] sub_value: Annotated[list[str], operator.add] def parent_start(state: ParentState) -> ParentState: return {"value": ["p_start"]} def sub_step(state: SubState) -> SubState: nonlocal sub_call_count sub_call_count += 1 return {"sub_value": [f"sub{sub_call_count}"]} def parent_end(state: ParentState) -> ParentState: return {"value": ["p_end"]} sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile() ) graph = ( StateGraph(ParentState) .add_node("parent_start", parent_start) .add_node("sub_graph", sub) .add_node("parent_end", parent_end) .add_edge(START, "parent_start") .add_edge("parent_start", "sub_graph") .add_edge("sub_graph", "parent_end") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": [], "sub_value": []}, config) assert result == {"value": ["p_start", "p_end"], "sub_value": ["sub1"]} # Capture original checkpoint IDs original_history = list(graph.get_state_history(config)) original_ids = {s.config["configurable"]["checkpoint_id"] for s in original_history} # Find checkpoint before sub_graph before_sub = next(s for s in original_history if s.next == ("sub_graph",)) before_sub_id = before_sub.config["configurable"]["checkpoint_id"] # Replay from before sub_graph replay_result = graph.invoke(None, before_sub.config) assert replay_result == {"value": ["p_start", "p_end"], "sub_value": ["sub2"]} # Get full history after replay post_replay_history = list(graph.get_state_history(config)) post_replay_ids = { s.config["configurable"]["checkpoint_id"] for s in post_replay_history } # All original checkpoint IDs still exist (append-only) assert original_ids.issubset(post_replay_ids) # New checkpoints were added (the branch) new_ids = post_replay_ids - original_ids assert len(new_ids) >= 2 # sub_graph + parent_end at minimum # The oldest new checkpoint's parent is the checkpoint we replayed from new_checkpoints = [ s for s in post_replay_history if s.config["configurable"]["checkpoint_id"] in new_ids ] oldest_new = sorted(new_checkpoints, key=lambda s: s.created_at)[0] assert oldest_new.parent_config is not None assert oldest_new.parent_config["configurable"]["checkpoint_id"] == before_sub_id # get_state returns the new branch tip latest = graph.get_state(config) assert latest.config["configurable"]["checkpoint_id"] in new_ids assert latest.values == {"value": ["p_start", "p_end"], "sub_value": ["sub2"]} def test_fork_creates_branch_preserving_old_checkpoints( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Forking (update_state + invoke) from a mid-run checkpoint creates a new branch of checkpoints while the original sequence is preserved. Original run (newest first): C4 next=() values=[a, b1, c] parent=C3 C3 next=(node_c,) values=[a, b1] parent=C2 C2 next=(node_b,) values=[a] parent=C1 C1 next=(node_a,) values=[] parent=C0 C0 next=(__start__,) values={} parent=None After fork from C2 with update {"value": ["x"]} (newest first): C7 next=() values=[a, x, b2, c] parent=C6 C6 next=(node_c,) values=[a, x, b2] parent=C5 C5 next=(node_b,) values=[a, x] parent=C2 <- fork checkpoint C4 next=() values=[a, b1, c] parent=C3 <- old branch preserved C3 next=(node_c,) values=[a, b1] parent=C2 C2 next=(node_b,) values=[a] parent=C1 <- fork point C1 next=(node_a,) values=[] parent=C0 C0 next=(__start__,) values={} parent=None """ call_count = 0 def node_a(state: State) -> State: return {"value": ["a"]} def node_b(state: State) -> State: nonlocal call_count call_count += 1 return {"value": [f"b{call_count}"]} def node_c(state: State) -> State: return {"value": ["c"]} graph = ( StateGraph(State) .add_node("node_a", node_a) .add_node("node_b", node_b) .add_node("node_c", node_c) .add_edge(START, "node_a") .add_edge("node_a", "node_b") .add_edge("node_b", "node_c") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": []}, config) assert result == {"value": ["a", "b1", "c"]} # -- Original checkpoint history (newest first) -- original_history = list(graph.get_state_history(config)) original_summary = _checkpoint_summary(original_history) assert len(original_summary) == 5 assert [s["next"] for s in original_summary] == [ (), ("node_c",), ("node_b",), ("node_a",), ("__start__",), ] assert [s["values"] for s in original_summary] == [ {"value": ["a", "b1", "c"]}, {"value": ["a", "b1"]}, {"value": ["a"]}, {"value": []}, {"value": []}, ] original_ids = {s.config["configurable"]["checkpoint_id"] for s in original_history} # Fork from before node_b with modified state before_b = next(s for s in original_history if s.next == ("node_b",)) before_b_id = before_b.config["configurable"]["checkpoint_id"] fork_config = graph.update_state(before_b.config, {"value": ["x"]}) fork_result = graph.invoke(None, fork_config) assert fork_result == {"value": ["a", "x", "b2", "c"]} # -- Post-fork checkpoint history (newest first) -- post_fork_history = list(graph.get_state_history(config)) post_summary = _checkpoint_summary(post_fork_history) # 5 original + 1 fork checkpoint (update_state) + 2 new nodes (node_b, node_c) assert len(post_summary) == 8 assert [s["next"] for s in post_summary] == [ (), # new branch tip (C7) ("node_c",), # new branch (C6) ("node_b",), # fork checkpoint from update_state (C5) (), # old branch tip (C4) ("node_c",), # old (C3) ("node_b",), # fork point (C2) ("node_a",), # old (C1) ("__start__",), # old (C0) ] assert [s["values"] for s in post_summary] == [ {"value": ["a", "x", "b2", "c"]}, # new branch tip {"value": ["a", "x", "b2"]}, # new: node_b re-ran {"value": ["a", "x"]}, # fork: state updated with "x" {"value": ["a", "b1", "c"]}, # old branch tip preserved {"value": ["a", "b1"]}, # old {"value": ["a"]}, # fork point {"value": []}, # old {"value": []}, # old ] # All original checkpoint IDs still exist (append-only) post_ids = {s.config["configurable"]["checkpoint_id"] for s in post_fork_history} assert original_ids.issubset(post_ids) # Fork checkpoint's parent is the branch point new_checkpoints = [ s for s in post_fork_history if s.config["configurable"]["checkpoint_id"] not in original_ids ] oldest_new = sorted(new_checkpoints, key=lambda s: s.created_at)[0] assert oldest_new.parent_config is not None assert oldest_new.parent_config["configurable"]["checkpoint_id"] == before_b_id # get_state returns the new branch tip latest = graph.get_state(config) assert latest.values == {"value": ["a", "x", "b2", "c"]} assert latest.config["configurable"]["checkpoint_id"] not in original_ids def test_fork_creates_branch_in_subgraph( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Forking a graph with a subgraph from a mid-run checkpoint creates a new branch while preserving the original checkpoint sequence. The subgraph re-executes on the new branch and the old checkpoints remain. """ sub_call_count = 0 class SubState(TypedDict): sub_value: Annotated[list[str], operator.add] class ParentState(TypedDict): value: Annotated[list[str], operator.add] sub_value: Annotated[list[str], operator.add] def parent_start(state: ParentState) -> ParentState: return {"value": ["p_start"]} def sub_step(state: SubState) -> SubState: nonlocal sub_call_count sub_call_count += 1 return {"sub_value": [f"sub{sub_call_count}"]} def parent_end(state: ParentState) -> ParentState: return {"value": ["p_end"]} sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile() ) graph = ( StateGraph(ParentState) .add_node("parent_start", parent_start) .add_node("sub_graph", sub) .add_node("parent_end", parent_end) .add_edge(START, "parent_start") .add_edge("parent_start", "sub_graph") .add_edge("sub_graph", "parent_end") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} result = graph.invoke({"value": [], "sub_value": []}, config) assert result == {"value": ["p_start", "p_end"], "sub_value": ["sub1"]} # Capture original checkpoint IDs original_history = list(graph.get_state_history(config)) original_ids = {s.config["configurable"]["checkpoint_id"] for s in original_history} # Find checkpoint before sub_graph and fork with modified state before_sub = next(s for s in original_history if s.next == ("sub_graph",)) before_sub_id = before_sub.config["configurable"]["checkpoint_id"] fork_config = graph.update_state(before_sub.config, {"value": ["extra"]}) fork_result = graph.invoke(None, fork_config) assert fork_result == { "value": ["p_start", "extra", "p_end"], "sub_value": ["sub2"], } # Get full history after fork post_fork_history = list(graph.get_state_history(config)) post_fork_ids = { s.config["configurable"]["checkpoint_id"] for s in post_fork_history } # All original checkpoint IDs still exist (append-only) assert original_ids.issubset(post_fork_ids) # New checkpoints were added (the branch) new_ids = post_fork_ids - original_ids assert len(new_ids) >= 3 # fork checkpoint + sub_graph + parent_end # The oldest new checkpoint's parent is the checkpoint we forked from new_checkpoints = [ s for s in post_fork_history if s.config["configurable"]["checkpoint_id"] in new_ids ] oldest_new = sorted(new_checkpoints, key=lambda s: s.created_at)[0] assert oldest_new.parent_config is not None assert oldest_new.parent_config["configurable"]["checkpoint_id"] == before_sub_id # get_state returns the new branch tip latest = graph.get_state(config) assert latest.config["configurable"]["checkpoint_id"] in new_ids assert latest.values == { "value": ["p_start", "extra", "p_end"], "sub_value": ["sub2"], } def test_stateless_subgraph_starts_fresh_on_parent_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Stateless subgraph (no checkpointer) always starts with empty state, even after prior invocations have completed.""" started: list[tuple[str, dict]] = [] observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_node(state: ParentState) -> ParentState: return {"results": ["p"]} def step_a(state: SubState) -> SubState: started.append(("step_a", dict(state))) answer = interrupt("question_a") observed.append(("step_a", dict(state))) return {"value": [f"a:{answer}"]} def step_b(state: SubState) -> SubState: started.append(("step_b", dict(state))) answer = interrupt("question_b") observed.append(("step_b", dict(state))) return {"value": [f"b:{answer}"]} sub = ( StateGraph(SubState) .add_node("step_a", step_a) .add_node("step_b", step_b) .add_edge(START, "step_a") .add_edge("step_a", "step_b") .compile() # no checkpointer — stateless ) graph = ( StateGraph(ParentState) .add_node("parent_node", parent_node) .add_node("sub_node", sub) .add_edge(START, "parent_node") .add_edge("parent_node", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # === 1st invocation: answer "a1" and "b1" === graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a1"), config) # hits step_b interrupt graph.invoke(Command(resume="b1"), config) # completes # step_a saw empty state, step_b saw only step_a's answer assert observed[0] == ("step_a", {"value": []}) assert observed[1] == ("step_b", {"value": ["a:a1"]}) # === 2nd invocation: answer "a2" and "b2" === observed.clear() graph.invoke({"results": []}, config) # hits step_a interrupt graph.invoke(Command(resume="a2"), config) # hits step_b interrupt graph.invoke(Command(resume="b2"), config) # completes # Stateless subgraph starts fresh — no memory of 1st invocation assert observed[0] == ("step_a", {"value": []}) assert observed[1] == ("step_b", {"value": ["a:a2"]}) # === Replay from checkpoint before sub_node in 2nd invocation === history = list(graph.get_state_history(config)) before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0] started.clear() replay = graph.invoke(None, before_sub_2nd.config) assert "__interrupt__" in replay # Stateless subgraph starts completely fresh on replay assert started[0] == ("step_a", {"value": []}) def test_stateful_subgraph_loads_latest_after_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """After replaying a parent checkpoint, a subsequent (3rd) invocation should load the subgraph state created by the replay — not the state from the checkpoint we replayed from.""" observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_node(state: ParentState) -> ParentState: return {"results": ["p"]} def sub_step(state: SubState) -> SubState: observed.append(("sub_step", dict(state))) return {"value": ["s"]} sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_node", parent_node) .add_node("sub_node", sub) .add_edge(START, "parent_node") .add_edge("parent_node", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation — subgraph starts fresh graph.invoke({"results": []}, config) assert observed[-1] == ("sub_step", {"value": []}) # 2nd invocation — subgraph sees state from 1st graph.invoke({"results": []}, config) assert observed[-1] == ("sub_step", {"value": ["s"]}) # Replay from checkpoint before parent_node in 2nd invocation history = list(graph.get_state_history(config)) before_parent_2nd = [s for s in history if s.next == ("parent_node",)][0] observed.clear() graph.invoke(None, before_parent_2nd.config) # Replay should load subgraph state from end of 1st invocation assert observed[0] == ("sub_step", {"value": ["s"]}) # 3rd invocation — should see state from the replay (2 × "s"), not from # the checkpoint we replayed from (1 × "s") observed.clear() graph.invoke({"results": []}, config) assert observed[0] == ("sub_step", {"value": ["s", "s"]}) def test_three_level_nested_subgraph_loads_state_on_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Three levels of nesting: parent -> mid -> inner. Replaying from the parent should load correct state at all levels.""" observed: list[tuple[str, dict]] = [] class InnerState(TypedDict): inner_trail: Annotated[list[str], operator.add] class MidState(TypedDict): mid_trail: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def inner_step(state: InnerState) -> InnerState: observed.append(("inner_step", dict(state))) return {"inner_trail": ["inner"]} def mid_step(state: MidState) -> MidState: observed.append(("mid_step", dict(state))) return {"mid_trail": ["mid"]} def parent_step(state: ParentState) -> ParentState: return {"results": ["p"]} inner = ( StateGraph(InnerState) .add_node("inner_step", inner_step) .add_edge(START, "inner_step") .compile(checkpointer=True) ) mid = ( StateGraph(MidState) .add_node("mid_step", mid_step) .add_node("inner_node", inner) .add_edge(START, "mid_step") .add_edge("mid_step", "inner_node") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_step", parent_step) .add_node("mid_node", mid) .add_edge(START, "parent_step") .add_edge("parent_step", "mid_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation — everything starts fresh graph.invoke({"results": []}, config) assert observed == [ ("mid_step", {"mid_trail": []}), ("inner_step", {"inner_trail": []}), ] # 2nd invocation — both levels see accumulated state observed.clear() graph.invoke({"results": []}, config) assert observed == [ ("mid_step", {"mid_trail": ["mid"]}), ("inner_step", {"inner_trail": ["inner"]}), ] # Replay from checkpoint before parent_step in 2nd invocation history = list(graph.get_state_history(config)) before_parent_2nd = [s for s in history if s.next == ("parent_step",)][0] observed.clear() graph.invoke(None, before_parent_2nd.config) # Both mid and inner should load state from end of 1st invocation assert observed == [ ("mid_step", {"mid_trail": ["mid"]}), ("inner_step", {"inner_trail": ["inner"]}), ] # 3rd invocation — sees state from replay, not from the replayed checkpoint observed.clear() graph.invoke({"results": []}, config) assert observed == [ ("mid_step", {"mid_trail": ["mid", "mid"]}), ("inner_step", {"inner_trail": ["inner", "inner"]}), ] def test_three_level_nested_subgraph_loads_state_on_fork( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Three levels of nesting with fork instead of replay.""" observed: list[tuple[str, dict]] = [] class InnerState(TypedDict): inner_trail: Annotated[list[str], operator.add] class MidState(TypedDict): mid_trail: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def inner_step(state: InnerState) -> InnerState: observed.append(("inner_step", dict(state))) return {"inner_trail": ["inner"]} def mid_step(state: MidState) -> MidState: observed.append(("mid_step", dict(state))) return {"mid_trail": ["mid"]} def parent_step(state: ParentState) -> ParentState: return {"results": ["p"]} inner = ( StateGraph(InnerState) .add_node("inner_step", inner_step) .add_edge(START, "inner_step") .compile(checkpointer=True) ) mid = ( StateGraph(MidState) .add_node("mid_step", mid_step) .add_node("inner_node", inner) .add_edge(START, "mid_step") .add_edge("mid_step", "inner_node") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_step", parent_step) .add_node("mid_node", mid) .add_edge(START, "parent_step") .add_edge("parent_step", "mid_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation graph.invoke({"results": []}, config) # 2nd invocation graph.invoke({"results": []}, config) # Fork from checkpoint before parent_step in 2nd invocation history = list(graph.get_state_history(config)) before_parent_2nd = [s for s in history if s.next == ("parent_step",)][0] fork_config = graph.update_state(before_parent_2nd.config, {"results": ["forked"]}) observed.clear() graph.invoke(None, fork_config) # Both mid and inner should load state from end of 1st invocation assert observed == [ ("mid_step", {"mid_trail": ["mid"]}), ("inner_step", {"inner_trail": ["inner"]}), ] def test_replay_from_first_invocation_checkpoint( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Replaying from the 1st invocation's checkpoint should load the subgraph state from before that invocation (i.e. empty).""" observed: list[tuple[str, dict]] = [] class SubState(TypedDict): value: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_node(state: ParentState) -> ParentState: return {"results": ["p"]} def sub_step(state: SubState) -> SubState: observed.append(("sub_step", dict(state))) return {"value": ["s"]} sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_node", parent_node) .add_node("sub_node", sub) .add_edge(START, "parent_node") .add_edge("parent_node", "sub_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # Run twice so subgraph accumulates state graph.invoke({"results": []}, config) graph.invoke({"results": []}, config) # Replay from before sub_node in 1st invocation (furthest back) history = list(graph.get_state_history(config)) before_sub_1st = [s for s in history if s.next == ("sub_node",)][-1] observed.clear() graph.invoke(None, before_sub_1st.config) # Should see empty state — no prior subgraph checkpoints exist assert observed[0] == ("sub_step", {"value": []}) def test_parallel_subgraph_nodes_load_correct_state_on_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Two sibling subgraph nodes that fan out from a common predecessor. Each should independently load its own historical checkpoint on replay.""" observed: list[tuple[str, dict]] = [] class SubStateA(TypedDict): a_trail: Annotated[list[str], operator.add] class SubStateB(TypedDict): b_trail: Annotated[list[str], operator.add] class ParentState(TypedDict): results: Annotated[list[str], operator.add] def parent_step(state: ParentState) -> ParentState: return {"results": ["p"]} def sub_a_step(state: SubStateA) -> SubStateA: observed.append(("sub_a", dict(state))) return {"a_trail": ["a"]} def sub_b_step(state: SubStateB) -> SubStateB: observed.append(("sub_b", dict(state))) return {"b_trail": ["b"]} sub_a = ( StateGraph(SubStateA) .add_node("sub_a_step", sub_a_step) .add_edge(START, "sub_a_step") .compile(checkpointer=True) ) sub_b = ( StateGraph(SubStateB) .add_node("sub_b_step", sub_b_step) .add_edge(START, "sub_b_step") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("parent_step", parent_step) .add_node("sub_a_node", sub_a) .add_node("sub_b_node", sub_b) .add_edge(START, "parent_step") # Fan out: both subgraphs run in parallel after parent_step .add_edge("parent_step", "sub_a_node") .add_edge("parent_step", "sub_b_node") .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation — both subgraphs start fresh graph.invoke({"results": []}, config) a_obs = [o for o in observed if o[0] == "sub_a"] b_obs = [o for o in observed if o[0] == "sub_b"] assert a_obs[0] == ("sub_a", {"a_trail": []}) assert b_obs[0] == ("sub_b", {"b_trail": []}) # 2nd invocation — both see accumulated state observed.clear() graph.invoke({"results": []}, config) a_obs = [o for o in observed if o[0] == "sub_a"] b_obs = [o for o in observed if o[0] == "sub_b"] assert a_obs[0] == ("sub_a", {"a_trail": ["a"]}) assert b_obs[0] == ("sub_b", {"b_trail": ["b"]}) # Replay from checkpoint before parent_step in 2nd invocation history = list(graph.get_state_history(config)) before_parent_2nd = [s for s in history if s.next == ("parent_step",)][0] observed.clear() graph.invoke(None, before_parent_2nd.config) # Each subgraph should independently load state from end of 1st invocation a_obs = [o for o in observed if o[0] == "sub_a"] b_obs = [o for o in observed if o[0] == "sub_b"] assert a_obs[0] == ("sub_a", {"a_trail": ["a"]}) assert b_obs[0] == ("sub_b", {"b_trail": ["b"]}) # 3rd invocation — sees state from replay observed.clear() graph.invoke({"results": []}, config) a_obs = [o for o in observed if o[0] == "sub_a"] b_obs = [o for o in observed if o[0] == "sub_b"] assert a_obs[0] == ("sub_a", {"a_trail": ["a", "a"]}) assert b_obs[0] == ("sub_b", {"b_trail": ["b", "b"]}) def test_subgraph_called_in_loop_loads_state_on_replay( sync_checkpointer: BaseCheckpointSaver, ) -> None: """Parent calls a subgraph node multiple times per invocation via a conditional loop. Replay should restore the subgraph's accumulated state from the correct point in the parent's timeline.""" observed: list[tuple[str, dict]] = [] class SubState(TypedDict): sub_trail: Annotated[list[str], operator.add] class ParentState(TypedDict): counter: int results: Annotated[list[str], operator.add] def inc(state: ParentState) -> ParentState: return {"counter": state["counter"] + 1, "results": [f"inc:{state['counter']}"]} def sub_step(state: SubState) -> SubState: observed.append(("sub_step", dict(state))) return {"sub_trail": ["s"]} def should_loop(state: ParentState) -> str: return "inc" if state["counter"] < 2 else "__end__" sub = ( StateGraph(SubState) .add_node("sub_step", sub_step) .add_edge(START, "sub_step") .compile(checkpointer=True) ) graph = ( StateGraph(ParentState) .add_node("inc", inc) .add_node("sub_node", sub) .add_edge(START, "inc") .add_edge("inc", "sub_node") .add_conditional_edges("sub_node", should_loop, ["inc", "__end__"]) .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} # 1st invocation: loop runs inc->sub->inc->sub->end # counter goes 0->1->2, subgraph called twice graph.invoke({"counter": 0, "results": []}, config) assert len(observed) == 2 assert observed[0] == ("sub_step", {"sub_trail": []}) assert observed[1] == ("sub_step", {"sub_trail": ["s"]}) # 2nd invocation: loop runs again, subgraph sees accumulated state observed.clear() graph.invoke({"counter": 0, "results": []}, config) assert len(observed) == 2 assert observed[0] == ("sub_step", {"sub_trail": ["s", "s"]}) assert observed[1] == ("sub_step", {"sub_trail": ["s", "s", "s"]}) # Replay from the START of the 2nd invocation's loop (counter=0). # History has two inc checkpoints with counter=0: 2nd invocation's (newer) # and 1st invocation's (older). Pick the newer one. history = list(graph.get_state_history(config)) start_of_loop_2nd = [ s for s in history if s.next == ("inc",) and s.values["counter"] == 0 ][0] observed.clear() graph.invoke(None, start_of_loop_2nd.config) # Full loop re-runs (2 sub calls). Subgraph loads state from end of # 1st invocation (2 × "s"), same as the original 2nd invocation. assert len(observed) == 2 assert observed[0] == ("sub_step", {"sub_trail": ["s", "s"]}) assert observed[1] == ("sub_step", {"sub_trail": ["s", "s", "s"]}) # Also test replay from MID-loop (counter=1) in the 2nd invocation. # Only one loop iteration remains, and the subgraph should load state # that includes the first loop iteration of the 2nd invocation (3 × "s"). mid_loop_2nd = [ s for s in history if s.next == ("inc",) and s.values["counter"] == 1 ][0] observed.clear() graph.invoke(None, mid_loop_2nd.config) assert len(observed) == 1 assert observed[0] == ("sub_step", {"sub_trail": ["s", "s", "s"]})