mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 08:32:24 +02:00
move tests
This commit is contained in:
Sydney Runkle
@@ -5316,262 +5316,6 @@ def test_multiple_interrupt_state_persistence(
|
||||
assert state.values["steps"] == ["step1", "step2"]
|
||||
|
||||
|
||||
def test_fork_before_all_interrupts(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Replaying from a checkpoint before any interrupt node should re-trigger
|
||||
the interrupt rather than reusing cached resume values from the original
|
||||
execution."""
|
||||
|
||||
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=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# 1. Run until interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
|
||||
# 2. Resume with answer — completes the full graph
|
||||
result = graph.invoke(Command(resume="hello"), config)
|
||||
assert result == {"value": ["a", "human:hello", "b"]}
|
||||
|
||||
# 3. Find checkpoint before ask_human (after node_a completed)
|
||||
history = list(graph.get_state_history(config))
|
||||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||||
|
||||
# 4. Replay from that checkpoint — interrupt should re-fire
|
||||
called.clear()
|
||||
replay_result = graph.invoke(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
|
||||
|
||||
|
||||
def test_fork_between_two_interrupt_nodes(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""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 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"}}
|
||||
|
||||
# 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
|
||||
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
|
||||
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
|
||||
interrupt_state2 = graph.get_state(config)
|
||||
result = graph.invoke(Command(resume="ans2"), interrupt_state2.config)
|
||||
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_concurrent_execution_thread_safety():
|
||||
"""Test thread safety during concurrent execution."""
|
||||
|
||||
|
||||
@@ -6530,252 +6530,6 @@ 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_concurrent_execution():
|
||||
"""Test concurrent execution with async nodes."""
|
||||
|
||||
|
||||
@@ -1244,3 +1244,264 @@ def test_copy_fork_creates_sibling_checkpoint(
|
||||
regular_config = graph.update_state(before_b.config, {"value": ["x"]})
|
||||
regular_state = graph.get_state(regular_config)
|
||||
assert regular_state.metadata["source"] == "update"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Section: Replay / fork with interrupts (moved from test_pregel)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fork_before_all_interrupts(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Replaying from a checkpoint before any interrupt node should re-trigger
|
||||
the interrupt rather than reusing cached resume values from the original
|
||||
execution."""
|
||||
|
||||
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=sync_checkpointer)
|
||||
)
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# 1. Run until interrupt
|
||||
result = graph.invoke({"value": []}, config)
|
||||
assert "__interrupt__" in result
|
||||
|
||||
# 2. Resume with answer — completes the full graph
|
||||
result = graph.invoke(Command(resume="hello"), config)
|
||||
assert result == {"value": ["a", "human:hello", "b"]}
|
||||
|
||||
# 3. Find checkpoint before ask_human (after node_a completed)
|
||||
history = list(graph.get_state_history(config))
|
||||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||||
|
||||
# 4. Replay from that checkpoint — interrupt should re-fire
|
||||
called.clear()
|
||||
replay_result = graph.invoke(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
|
||||
|
||||
|
||||
def test_fork_between_two_interrupt_nodes(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""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 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"}}
|
||||
|
||||
# 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
|
||||
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
|
||||
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
|
||||
interrupt_state2 = graph.get_state(config)
|
||||
result = graph.invoke(Command(resume="ans2"), interrupt_state2.config)
|
||||
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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user