From 63528f25af89d47e2f1a323a1cb8b8b854db1a4e Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 4 Mar 2026 17:36:15 -0800 Subject: [PATCH] more tests --- libs/langgraph/tests/test_pregel.py | 409 ++++++++++++++++++-- libs/langgraph/tests/test_pregel_async.py | 436 ++++++++++++++++++++++ 2 files changed, 810 insertions(+), 35 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index db1baea3c..b3a3c11fa 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5316,13 +5316,12 @@ def test_multiple_interrupt_state_persistence( assert state.values["steps"] == ["step1", "step2"] -def test_fork_from_resolved_interrupt_retriggers( +def test_fork_before_all_interrupts( sync_checkpointer: BaseCheckpointSaver, ) -> None: - """Replaying from a checkpoint before an interrupt node should re-trigger + """Replaying from a checkpoint before any interrupt node should re-trigger the interrupt rather than reusing cached resume values from the original - execution. This covers both the checkpoint that directly resolved the - interrupt and an earlier checkpoint (before the interrupt node).""" + execution.""" called: list[str] = [] @@ -5374,44 +5373,51 @@ def test_fork_from_resolved_interrupt_retriggers( assert "__interrupt__" in replay_result assert replay_result["value"] == ["a"] assert replay_result["__interrupt__"][0].value == "What is your input?" - # ask_human was called but hit interrupt before returning assert "ask_human" in called - # node_a should NOT run (it's before our checkpoint) assert "node_a" not in called - # node_b should NOT run (interrupt halted execution) assert "node_b" not in called - # 5. Resume the re-triggered interrupt on a fresh thread to verify - # the interrupt is functional (the fork's checkpoints are not the - # latest on the original thread, so we use a new thread). - called.clear() - fresh_config = {"configurable": {"thread_id": "2"}} - result = graph.invoke({"value": []}, fresh_config) - assert "__interrupt__" in result - result = graph.invoke(Command(resume="world"), fresh_config) - assert result == {"value": ["a", "human:world", "b"]} - -def test_fork_multiple_interrupts_resume_with_checkpoint_id( +def test_fork_between_two_interrupt_nodes( sync_checkpointer: BaseCheckpointSaver, ) -> None: - """When a node has multiple interrupts and we resume them one at a time - using Command(resume=...) with a specific checkpoint_id, previously - resolved RESUME values must be preserved so that later interrupts can - also be resolved.""" + """Replaying from a checkpoint between two interrupt nodes (after the first + interrupt was resolved, before the second) should re-trigger the second + interrupt.""" + + called: list[str] = [] class State(TypedDict): value: Annotated[list[str], operator.add] - def multi_interrupt_node(state: State) -> State: - answer1 = interrupt("First question?") - answer2 = interrupt("Second question?") - return {"value": [f"a1:{answer1}", f"a2:{answer2}"]} + 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("ask", multi_interrupt_node) - .add_edge(START, "ask") + .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) ) @@ -5422,21 +5428,354 @@ def test_fork_multiple_interrupts_resume_with_checkpoint_id( assert "__interrupt__" in result assert result["__interrupt__"][0].value == "First question?" - # Grab the checkpoint where the interrupt fired - interrupt_state = graph.get_state(config) - interrupt_config = interrupt_state.config + # 2. Resume first interrupt + result = graph.invoke(Command(resume="ans1"), config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Second question?" + + # 3. Resume second interrupt — completes the full graph + result = graph.invoke(Command(resume="ans2"), config) + assert result == {"value": ["a", "i1:ans1", "i2:ans2", "b"]} + + # 4. Find checkpoint between the two interrupts (after interrupt_1, before + # interrupt_2) + history = list(graph.get_state_history(config)) + between = [s for s in history if s.next == ("interrupt_2",)][-1] + + # 5. Replay from that checkpoint — second interrupt should re-fire + called.clear() + replay_result = graph.invoke(None, between.config) + assert "__interrupt__" in replay_result + assert replay_result["__interrupt__"][0].value == "Second question?" + assert "interrupt_2" in called + assert "interrupt_1" not in called + assert "node_a" not in called + assert "node_b" not in called + + # 6. Also replay from before interrupt_1 — first interrupt should re-fire + 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_fork_multiple_interrupts_in_one_node( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from a checkpoint before a node with multiple interrupts + should re-trigger the first interrupt. Resuming with checkpoint_id should + preserve previously resolved RESUME values.""" + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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"}} + + # 1. Run until first interrupt + result = graph.invoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "First question?" # 2. Resume first interrupt with checkpoint_id — should hit second interrupt - result = graph.invoke(Command(resume="ans1"), interrupt_config) + interrupt_state = graph.get_state(config) + result = graph.invoke(Command(resume="ans1"), interrupt_state.config) assert "__interrupt__" in result assert result["__interrupt__"][0].value == "Second question?" # 3. Resume second interrupt with checkpoint_id — should complete - # This is the critical test: the first RESUME value ("ans1") must still - # be present in pending writes, otherwise the first interrupt re-fires. interrupt_state2 = graph.get_state(config) result = graph.invoke(Command(resume="ans2"), interrupt_state2.config) - assert result == {"value": ["a1:ans1", "a2:ans2"]} + assert result == {"value": ["a1:ans1", "a2:ans2", "done"]} + + # 4. Replay from before the multi-interrupt 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 "__interrupt__" in replay_result + assert replay_result["__interrupt__"][0].value == "First question?" + + +def test_fork_after_all_interrupts( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from the final checkpoint (after all interrupts resolved and + graph completed) should not re-trigger any interrupts.""" + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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("Question?") + 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"}} + + # 1. Run until interrupt + result = graph.invoke({"value": []}, config) + assert "__interrupt__" in result + + # 2. Resume — completes the full graph + result = graph.invoke(Command(resume="hello"), config) + assert result == {"value": ["a", "human:hello", "b"]} + + # 3. Get the final checkpoint (graph completed, no next nodes) + history = list(graph.get_state_history(config)) + final = [s for s in history if not s.next][0] + + # 4. Replay from final checkpoint — nothing should run + called.clear() + replay_result = graph.invoke(None, final.config) + assert "__interrupt__" not in replay_result + assert "ask_human" not in called + assert "node_b" not in called + + +def test_fork_subgraph_interrupt_no_checkpointer( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with a subgraph that has no checkpointer (checkpointer=False/None). + The subgraph inherits the parent's checkpointer via config and its saved + checkpoint retains RESUME writes (CONFIG_KEY_RESUMING is propagated from + the parent). So the subgraph does NOT re-fire the interrupt on replay — + it uses the cached resume value and completes.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_interrupt(state: State) -> State: + called.append("sub_interrupt") + answer = interrupt("Sub question?") + return {"value": [f"sub:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_interrupt", sub_interrupt) + .add_edge(START, "sub_interrupt") + .compile() # no checkpointer + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + def after(state: State) -> State: + called.append("after") + return {"value": ["after"]} + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_node("after", after) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "after") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until interrupt + result = graph.invoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub question?" + + # 2. Resume — completes + result = graph.invoke(Command(resume="answer"), config) + assert result == {"value": ["sub:answer", "after"]} + + # 3. Find checkpoint before subgraph node + history = list(graph.get_state_history(config)) + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + # 4. Replay — subgraph uses cached resume value, does NOT re-fire interrupt + called.clear() + replay_result = graph.invoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["sub:answer", "after"]} + assert "call_subgraph" in called + assert "after" in called + + +def test_fork_subgraph_interrupt_checkpointer_true( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with a subgraph that has checkpointer=True. + Same behavior as no checkpointer — the subgraph's checkpoint retains + RESUME writes and the interrupt does NOT re-fire on replay.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_node(state: State) -> State: + called.append("sub_node") + return {"value": ["sub_node"]} + + def sub_interrupt(state: State) -> State: + called.append("sub_interrupt") + answer = interrupt("Sub question?") + return {"value": [f"sub:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_node", sub_node) + .add_node("sub_interrupt", sub_interrupt) + .add_edge(START, "sub_node") + .add_edge("sub_node", "sub_interrupt") + .compile(checkpointer=True) + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + def after(state: State) -> State: + called.append("after") + return {"value": ["after"]} + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_node("after", after) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "after") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until interrupt + result = graph.invoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub question?" + + # 2. Resume — completes + result = graph.invoke(Command(resume="answer"), config) + assert result == {"value": ["sub_node", "sub:answer", "after"]} + + # 3. Find checkpoint before subgraph node + history = list(graph.get_state_history(config)) + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + # 4. Replay — subgraph uses cached resume value, does NOT re-fire interrupt + called.clear() + replay_result = graph.invoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["sub_node", "sub:answer", "after"]} + assert "call_subgraph" in called + assert "after" in called + + +def test_fork_subgraph_two_interrupts_no_checkpointer( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with a subgraph (no checkpointer) containing two interrupt + nodes. Same as single interrupt — the subgraph uses cached resume values + and completes without re-firing interrupts.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_int_1(state: State) -> State: + called.append("sub_int_1") + answer = interrupt("Sub Q1?") + return {"value": [f"s1:{answer}"]} + + def sub_int_2(state: State) -> State: + called.append("sub_int_2") + answer = interrupt("Sub Q2?") + return {"value": [f"s2:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_int_1", sub_int_1) + .add_node("sub_int_2", sub_int_2) + .add_edge(START, "sub_int_1") + .add_edge("sub_int_1", "sub_int_2") + .compile() # no checkpointer + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_edge(START, "call_subgraph") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until first sub-interrupt + result = graph.invoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub Q1?" + + # 2. Resume first + result = graph.invoke(Command(resume="a1"), config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub Q2?" + + # 3. Resume second — completes + result = graph.invoke(Command(resume="a2"), config) + assert result == {"value": ["s1:a1", "s2:a2"]} + + # 4. Replay from before subgraph — uses cached resume values + history = list(graph.get_state_history(config)) + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + called.clear() + replay_result = graph.invoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["s1:a1", "s2:a2"]} def test_concurrent_execution_thread_safety(): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3956a60b6..a00534124 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6530,6 +6530,442 @@ async def test_multiple_interrupt_state_persistence( assert state.values["steps"] == ["step1", "step2"] +async def test_fork_before_all_interrupts( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from a checkpoint before any interrupt node should re-trigger + the interrupt rather than reusing cached resume values.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until interrupt + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + + # 2. Resume — completes the full graph + result = await graph.ainvoke(Command(resume="hello"), config) + assert result == {"value": ["a", "human:hello", "b"]} + + # 3. Find checkpoint before ask_human + history = [s async for s in graph.aget_state_history(config)] + before_ask = [s for s in history if s.next == ("ask_human",)][-1] + + # 4. Replay — interrupt should re-fire + called.clear() + replay_result = await graph.ainvoke(None, before_ask.config) + + assert "__interrupt__" in replay_result + assert replay_result["value"] == ["a"] + assert replay_result["__interrupt__"][0].value == "What is your input?" + assert "ask_human" in called + assert "node_a" not in called + assert "node_b" not in called + + +async def test_fork_between_two_interrupt_nodes( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from a checkpoint between two interrupt nodes should + re-trigger the second interrupt.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until first interrupt + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "First question?" + + # 2. Resume first interrupt + result = await graph.ainvoke(Command(resume="ans1"), config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Second question?" + + # 3. Resume second — completes + result = await graph.ainvoke(Command(resume="ans2"), config) + assert result == {"value": ["a", "i1:ans1", "i2:ans2", "b"]} + + # 4. Find checkpoint between the two interrupts + history = [s async for s in graph.aget_state_history(config)] + between = [s for s in history if s.next == ("interrupt_2",)][-1] + + # 5. Replay — second interrupt re-fires + called.clear() + replay_result = await graph.ainvoke(None, between.config) + assert "__interrupt__" in replay_result + assert replay_result["__interrupt__"][0].value == "Second question?" + assert "interrupt_2" in called + assert "interrupt_1" not in called + assert "node_a" not in called + + # 6. Replay from before interrupt_1 + before_i1 = [s for s in history if s.next == ("interrupt_1",)][-1] + called.clear() + replay_result = await graph.ainvoke(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 + + +async def test_fork_multiple_interrupts_in_one_node( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from before a node with multiple interrupts should re-trigger + the first. Resuming with checkpoint_id preserves resolved values.""" + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # 1. Run until first interrupt + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "First question?" + + # 2. Resume first with checkpoint_id + interrupt_state = await graph.aget_state(config) + result = await graph.ainvoke(Command(resume="ans1"), interrupt_state.config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Second question?" + + # 3. Resume second with checkpoint_id — completes + interrupt_state2 = await graph.aget_state(config) + result = await graph.ainvoke(Command(resume="ans2"), interrupt_state2.config) + assert result == {"value": ["a1:ans1", "a2:ans2", "done"]} + + # 4. Replay from before multi-interrupt node + history = [s async for s in graph.aget_state_history(config)] + before_ask = [s for s in history if s.next == ("ask",)][-1] + replay_result = await graph.ainvoke(None, before_ask.config) + assert "__interrupt__" in replay_result + assert replay_result["__interrupt__"][0].value == "First question?" + + +async def test_fork_after_all_interrupts( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Replaying from the final checkpoint should not re-trigger interrupts.""" + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + 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("Question?") + 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=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + + result = await graph.ainvoke(Command(resume="hello"), config) + assert result == {"value": ["a", "human:hello", "b"]} + + history = [s async for s in graph.aget_state_history(config)] + final = [s for s in history if not s.next][0] + + called.clear() + replay_result = await graph.ainvoke(None, final.config) + assert "__interrupt__" not in replay_result + assert "ask_human" not in called + assert "node_b" not in called + + +async def test_fork_subgraph_interrupt_no_checkpointer( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with subgraph (no checkpointer). The subgraph inherits the + parent's checkpointer and CONFIG_KEY_RESUMING, so it uses cached resume + values and does NOT re-fire the interrupt.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_interrupt(state: State) -> State: + called.append("sub_interrupt") + answer = interrupt("Sub question?") + return {"value": [f"sub:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_interrupt", sub_interrupt) + .add_edge(START, "sub_interrupt") + .compile() + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + def after(state: State) -> State: + called.append("after") + return {"value": ["after"]} + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_node("after", after) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "after") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub question?" + + result = await graph.ainvoke(Command(resume="answer"), config) + assert result == {"value": ["sub:answer", "after"]} + + history = [s async for s in graph.aget_state_history(config)] + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + called.clear() + replay_result = await graph.ainvoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["sub:answer", "after"]} + assert "call_subgraph" in called + assert "after" in called + + +async def test_fork_subgraph_interrupt_checkpointer_true( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with subgraph (checkpointer=True). Same as no checkpointer — + uses cached resume values, does NOT re-fire the interrupt.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_node(state: State) -> State: + called.append("sub_node") + return {"value": ["sub_node"]} + + def sub_interrupt(state: State) -> State: + called.append("sub_interrupt") + answer = interrupt("Sub question?") + return {"value": [f"sub:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_node", sub_node) + .add_node("sub_interrupt", sub_interrupt) + .add_edge(START, "sub_node") + .add_edge("sub_node", "sub_interrupt") + .compile(checkpointer=True) + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + def after(state: State) -> State: + called.append("after") + return {"value": ["after"]} + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_node("after", after) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "after") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub question?" + + result = await graph.ainvoke(Command(resume="answer"), config) + assert result == {"value": ["sub_node", "sub:answer", "after"]} + + history = [s async for s in graph.aget_state_history(config)] + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + called.clear() + replay_result = await graph.ainvoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["sub_node", "sub:answer", "after"]} + assert "call_subgraph" in called + assert "after" in called + + +async def test_fork_subgraph_two_interrupts_no_checkpointer( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Fork/replay with subgraph (no checkpointer) with two interrupt nodes. + Same behavior — uses cached resume values, does NOT re-fire interrupts.""" + + called: list[str] = [] + + class State(TypedDict): + value: Annotated[list[str], operator.add] + + def sub_int_1(state: State) -> State: + called.append("sub_int_1") + answer = interrupt("Sub Q1?") + return {"value": [f"s1:{answer}"]} + + def sub_int_2(state: State) -> State: + called.append("sub_int_2") + answer = interrupt("Sub Q2?") + return {"value": [f"s2:{answer}"]} + + subgraph = ( + StateGraph(State) + .add_node("sub_int_1", sub_int_1) + .add_node("sub_int_2", sub_int_2) + .add_edge(START, "sub_int_1") + .add_edge("sub_int_1", "sub_int_2") + .compile() + ) + + def call_subgraph(state: State) -> State: + called.append("call_subgraph") + return subgraph.invoke(state) + + graph = ( + StateGraph(State) + .add_node("call_subgraph", call_subgraph) + .add_edge(START, "call_subgraph") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + result = await graph.ainvoke({"value": []}, config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub Q1?" + + result = await graph.ainvoke(Command(resume="a1"), config) + assert "__interrupt__" in result + assert result["__interrupt__"][0].value == "Sub Q2?" + + result = await graph.ainvoke(Command(resume="a2"), config) + assert result == {"value": ["s1:a1", "s2:a2"]} + + history = [s async for s in graph.aget_state_history(config)] + before_sub = [s for s in history if s.next == ("call_subgraph",)][-1] + + called.clear() + replay_result = await graph.ainvoke(None, before_sub.config) + assert "__interrupt__" not in replay_result + assert replay_result == {"value": ["s1:a1", "s2:a2"]} + + async def test_concurrent_execution(): """Test concurrent execution with async nodes."""