boom

libs/langgraph/tests/test_replay.py

@@ -0,0 +1,72 @@
+"""MRE: replay from before an interrupt node uses cached resume values."""
+
+import operator
+from typing import Annotated
+
+from langgraph.checkpoint.memory import MemorySaver
+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 test_replay_uses_cached_resume():
+    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=MemorySaver())
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Run until interrupt
+    result = graph.invoke({"value": []}, config)
+    assert "__interrupt__" in result
+
+    # Resume with answer
+    result = graph.invoke(Command(resume="hello"), config)
+    assert result == {"value": ["a", "human:hello", "b"]}
+
+    # 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 from that checkpoint
+    called.clear()
+    replay_result = graph.invoke(None, before_ask.config)
+
+    # Interrupt is NOT re-triggered — cached resume value used
+    assert replay_result == {"value": ["a", "human:hello", "b"]}
+    assert "__interrupt__" not in replay_result
+    assert "ask_human" in called
+    assert "node_b" in called
+    assert "node_a" not in called
+
+
+if __name__ == "__main__":
+    test_replay_uses_cached_resume()
+    print("PASSED")

libs/langgraph/tests/test_time_travel.py

@@ -859,6 +859,241 @@ def test_subgraph_interrupt_fork_from_subgraph_checkpoint_full_flow_no_sub_check
     assert "post" in final_result["value"]
 
 
+# ---------------------------------------------------------------------------
+# Section 4b: With subgraph, MULTIPLE interrupts in subgraph (ask_a + ask_b)
+# ---------------------------------------------------------------------------
+
+
+def _build_subgraph_multi_interrupt_graph(
+    checkpointer: BaseCheckpointSaver,
+    called: list[str],
+    subgraph_checkpointer=True,
+):
+    """Build: START -> router -> [subgraph: ask_a [interrupt] -> ask_b [interrupt] -> step_c] -> post_process -> END
+
+    Subgraph has TWO interrupt nodes to test multiple-interrupt scenarios.
+    """
+
+    class SubMultiIntState(TypedDict):
+        value: Annotated[list[str], operator.add]
+
+    class ParentMultiIntState(TypedDict):
+        value: Annotated[list[str], operator.add]
+
+    def router(state: ParentMultiIntState) -> ParentMultiIntState:
+        called.append("router")
+        return {"value": ["routed"]}
+
+    def ask_a(state: SubMultiIntState) -> SubMultiIntState:
+        called.append("ask_a")
+        answer = interrupt("Question A")
+        return {"value": [f"a:{answer}"]}
+
+    def ask_b(state: SubMultiIntState) -> SubMultiIntState:
+        called.append("ask_b")
+        answer = interrupt("Question B")
+        return {"value": [f"b:{answer}"]}
+
+    def step_c(state: SubMultiIntState) -> SubMultiIntState:
+        called.append("step_c")
+        return {"value": ["sub_c"]}
+
+    subgraph = (
+        StateGraph(SubMultiIntState)
+        .add_node("ask_a", ask_a)
+        .add_node("ask_b", ask_b)
+        .add_node("step_c", step_c)
+        .add_edge(START, "ask_a")
+        .add_edge("ask_a", "ask_b")
+        .add_edge("ask_b", "step_c")
+        .compile(checkpointer=subgraph_checkpointer)
+    )
+
+    def post_process(state: ParentMultiIntState) -> ParentMultiIntState:
+        called.append("post_process")
+        return {"value": ["post"]}
+
+    graph = (
+        StateGraph(ParentMultiIntState)
+        .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=checkpointer)
+    )
+
+    return graph
+
+
+def test_subgraph_multi_interrupt_full_flow_checkpointer_true(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Subgraph with two interrupt nodes (ask_a, ask_b), checkpointer=True.
+    Run through both interrupts, then fork from parent before subgraph and
+    verify behavior — with checkpointer=True, cached resumes are used."""
+
+    called: list[str] = []
+    graph = _build_subgraph_multi_interrupt_graph(sync_checkpointer, called)
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Hit first interrupt (ask_a)
+    r1 = graph.invoke({"value": []}, config)
+    assert "__interrupt__" in r1
+    assert r1["__interrupt__"][0].value == "Question A"
+
+    # Resume ask_a -> hits ask_b
+    r2 = graph.invoke(Command(resume="ans_a"), config)
+    assert "__interrupt__" in r2
+    assert r2["__interrupt__"][0].value == "Question B"
+
+    # Resume ask_b -> completes
+    r3 = graph.invoke(Command(resume="ans_b"), config)
+    assert "a:ans_a" in r3["value"]
+    assert "b:ans_b" in r3["value"]
+    assert "sub_c" in r3["value"]
+    assert "post" in r3["value"]
+
+    # Fork from parent checkpoint before subgraph_node
+    history = list(graph.get_state_history(config))
+    before_sub_candidates = [s for s in history if s.next == ("subgraph_node",)]
+    before_sub = before_sub_candidates[-1]
+
+    called.clear()
+    fork_config = graph.update_state(before_sub.config, {"value": ["forked"]})
+    fork_result = graph.invoke(None, fork_config)
+
+    # With checkpointer=True, subgraph has its own persistent state.
+    # Parent fork does NOT clear subgraph checkpoints — cached resumes used.
+    assert "__interrupt__" not in fork_result
+    assert "a:ans_a" in fork_result["value"]
+    assert "b:ans_b" in fork_result["value"]
+
+
+def test_subgraph_multi_interrupt_full_flow_checkpointer_none(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Subgraph with two interrupt nodes (ask_a, ask_b), checkpointer=None.
+    Run through both interrupts, then fork from parent before subgraph.
+    With no sub-checkpointer, forking parent re-triggers the first interrupt."""
+
+    called: list[str] = []
+    graph = _build_subgraph_multi_interrupt_graph(
+        sync_checkpointer, called, subgraph_checkpointer=None
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Hit first interrupt (ask_a)
+    r1 = graph.invoke({"value": []}, config)
+    assert "__interrupt__" in r1
+    assert r1["__interrupt__"][0].value == "Question A"
+
+    # Resume ask_a -> hits ask_b
+    r2 = graph.invoke(Command(resume="ans_a"), config)
+    assert "__interrupt__" in r2
+    assert r2["__interrupt__"][0].value == "Question B"
+
+    # Resume ask_b -> completes
+    r3 = graph.invoke(Command(resume="ans_b"), config)
+    assert "a:ans_a" in r3["value"]
+    assert "b:ans_b" in r3["value"]
+    assert "post" in r3["value"]
+
+    # Fork from parent checkpoint before subgraph_node
+    history = list(graph.get_state_history(config))
+    before_sub_candidates = [s for s in history if s.next == ("subgraph_node",)]
+    before_sub = before_sub_candidates[-1]
+
+    called.clear()
+    fork_config = graph.update_state(before_sub.config, {"value": ["forked"]})
+    fork_result = graph.invoke(None, fork_config)
+
+    # With no sub-checkpointer, forking parent re-triggers the first interrupt
+    assert "__interrupt__" in fork_result
+    assert fork_result["__interrupt__"][0].value == "Question A"
+
+    # Resume ask_a with new answer -> hits ask_b
+    r4 = graph.invoke(Command(resume="new_a"), fork_config)
+    assert "__interrupt__" in r4
+    assert r4["__interrupt__"][0].value == "Question B"
+
+    # Resume ask_b with new answer -> completes
+    final = graph.invoke(Command(resume="new_b"), fork_config)
+    assert "a:new_a" in final["value"]
+    assert "b:new_b" in final["value"]
+    assert "post" in final["value"]
+
+
+def test_subgraph_multi_interrupt_fork_from_subgraph_between_interrupts(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Subgraph with two interrupt nodes (ask_a, ask_b), checkpointer=True.
+    Fork from within the subgraph at the ask_b interrupt point using
+    update_state. With checkpointer=True, the subgraph's cached resume is
+    used so the interrupt does NOT re-fire — the graph completes."""
+
+    called: list[str] = []
+    graph = _build_subgraph_multi_interrupt_graph(sync_checkpointer, called)
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Hit first interrupt (ask_a)
+    r1 = graph.invoke({"value": []}, config)
+    assert "__interrupt__" in r1
+    assert r1["__interrupt__"][0].value == "Question A"
+
+    # Resume ask_a -> hits ask_b
+    r2 = graph.invoke(Command(resume="ans_a"), config)
+    assert "__interrupt__" in r2
+    assert r2["__interrupt__"][0].value == "Question B"
+
+    # Get subgraph state to find the checkpoint at the ask_b interrupt
+    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 the subgraph's checkpoint (at the ask_b interrupt)
+    called.clear()
+    fork_config = graph.update_state(sub_config, {"value": ["sub_forked"]})
+
+    # Invoke — with checkpointer=True, cached resume is used so the
+    # interrupt does NOT re-fire. The graph completes through parent.
+    fork_result = graph.invoke(None, fork_config)
+    assert "__interrupt__" not in fork_result
+    assert "sub_c" in fork_result["value"]
+    assert "post" in fork_result["value"]
+
+
+def test_subgraph_multi_interrupt_resume_from_parent_at_second_interrupt(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    """Subgraph with two interrupts, checkpointer=True. Resume from the parent
+    level when paused at the second interrupt (ask_b). Verifies that resuming
+    at the subgraph_node level correctly forwards to the subgraph."""
+
+    called: list[str] = []
+    graph = _build_subgraph_multi_interrupt_graph(sync_checkpointer, called)
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Hit first interrupt (ask_a)
+    graph.invoke({"value": []}, config)
+
+    # Resume ask_a -> hits ask_b
+    graph.invoke(Command(resume="ans_a"), config)
+
+    # Now resume ask_b from the parent level
+    called.clear()
+    final = graph.invoke(Command(resume="ans_b"), config)
+
+    assert "a:ans_a" in final["value"]
+    assert "b:ans_b" in final["value"]
+    assert "step_c" in called
+    assert "sub_c" in final["value"]
+    assert "post_process" in called
+    assert "post" in final["value"]
+
+
 # ---------------------------------------------------------------------------
 # Section 5: Additional scenarios from customer thread
 # ---------------------------------------------------------------------------