wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
3212 lines
107 KiB
Python
3212 lines
107 KiB
Python
"""Async 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
|
||
import sys
|
||
from typing import Annotated
|
||
|
||
import pytest
|
||
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
|
||
|
||
pytestmark = pytest.mark.anyio
|
||
|
||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="Python 3.11+ is required for async contextvars support",
|
||
)
|
||
|
||
|
||
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)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_reruns_nodes_after_checkpoint(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result == {"value": ["a", "b"]}
|
||
assert called == ["node_a", "node_b"]
|
||
|
||
# Find checkpoint before node_b (next=(node_b,))
|
||
history = [s async for s in graph.aget_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 = await graph.ainvoke(None, before_b.config)
|
||
|
||
assert replay_result == {"value": ["a", "b"]}
|
||
assert "node_b" in called
|
||
assert "node_a" not in called
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_from_final_checkpoint_is_noop(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result == {"value": ["a"]}
|
||
|
||
state = await graph.aget_state(config)
|
||
assert state.next == ()
|
||
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, state.config)
|
||
assert replay_result == {"value": ["a"]}
|
||
assert called == []
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_fork_reruns_with_modified_state(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_b = next(s for s in history if s.next == ("node_b",))
|
||
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(before_b.config, {"value": ["x"]})
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
|
||
assert "node_b" in called
|
||
assert fork_result == {"value": ["a", "x", "b"]}
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_multiple_forks_from_same_checkpoint(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_b = next(s for s in history if s.next == ("node_b",))
|
||
|
||
fork1_config = await graph.aupdate_state(before_b.config, {"value": ["fork1"]})
|
||
result1 = await graph.ainvoke(None, fork1_config)
|
||
|
||
fork2_config = await graph.aupdate_state(before_b.config, {"value": ["fork2"]})
|
||
result2 = await graph.ainvoke(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)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_from_before_interrupt_refires(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert "__interrupt__" in result
|
||
assert call_count["ask_human"] == 1
|
||
|
||
# Resume
|
||
result = await graph.ainvoke(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 = [s async for s in graph.aget_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 = await graph.ainvoke(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)
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_interrupt_stable_across_replays(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="cached_answer"), config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||
|
||
results = []
|
||
for _ in range(3):
|
||
r = await graph.ainvoke(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)
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_fork_from_before_interrupt_refires(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="hello"), config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(before_ask.config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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 = await graph.ainvoke(Command(resume="world"), fork_config)
|
||
assert final == {"value": ["a", "forked", "human:world", "b"]}
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_fork_from_interrupt_checkpoint_refires(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="hello"), config)
|
||
|
||
history = [s async for s in graph.aget_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 = await graph.aupdate_state(
|
||
interrupt_checkpoint.config, {"value": ["forked"]}
|
||
)
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
|
||
assert "__interrupt__" in fork_result
|
||
|
||
final = await graph.ainvoke(Command(resume="different"), fork_config)
|
||
assert "human:different" in final["value"]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Section 3: Multiple / sequential interrupts
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_sequential_interrupts_fork_from_middle(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Hit first interrupt
|
||
r1 = await graph.ainvoke({"value": []}, config)
|
||
assert r1["__interrupt__"][0].value == "First question?"
|
||
|
||
# Resume first → hit second
|
||
r2 = await graph.ainvoke(Command(resume="ans1"), config)
|
||
assert r2["__interrupt__"][0].value == "Second question?"
|
||
|
||
# Resume second → complete
|
||
r3 = await graph.ainvoke(Command(resume="ans2"), config)
|
||
assert r3 == {"value": ["a", "i1:ans1", "i2:ans2", "b"]}
|
||
|
||
history = [s async for s in graph.aget_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 = await graph.aupdate_state(between.config, {"value": ["mid_fork"]})
|
||
fork_result = await graph.ainvoke(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 = await graph.ainvoke(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 = 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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_multiple_interrupts_in_one_node(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Hit first interrupt
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result["__interrupt__"][0].value == "First question?"
|
||
|
||
# Resume first → hit second
|
||
interrupt_state = await graph.aget_state(config)
|
||
result = await graph.ainvoke(Command(resume="ans1"), interrupt_state.config)
|
||
assert result["__interrupt__"][0].value == "Second question?"
|
||
|
||
# Resume second → complete
|
||
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"]}
|
||
|
||
# Replay from before the node — first interrupt re-fires
|
||
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 replay_result["__interrupt__"][0].value == "First question?"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Section 4: Subgraph without interrupt
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_replay_from_before(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert "sub_a" in result["value"]
|
||
assert "sub_b" in result["value"]
|
||
assert "post" in result["value"]
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_sub = next(s for s in history if s.next == ("subgraph",))
|
||
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_fork_from_before(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_sub = next(s for s in history if s.next == ("subgraph",))
|
||
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(before_sub.config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_interrupt_replay_from_parent(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt, then resume
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert "__interrupt__" in result
|
||
|
||
completed_result = await graph.ainvoke(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 = [s async for s in graph.aget_state_history(config)]
|
||
before_sub = [s for s in history if s.next == ("subgraph_node",)][-1]
|
||
|
||
# Replay — interrupt re-fires
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, before_sub.config)
|
||
assert "__interrupt__" in replay_result
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_interrupt_replay_from_interrupt_checkpoint(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
# Verify subgraph state is accessible
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
assert len(parent_state.tasks) > 0
|
||
assert parent_state.tasks[0].state is not None
|
||
|
||
# Resume
|
||
await graph.ainvoke(Command(resume="answer"), config)
|
||
|
||
# Find the parent checkpoint where the interrupt fired
|
||
history = [s async for s in graph.aget_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 = await graph.ainvoke(None, interrupt_checkpoint.config)
|
||
assert "__interrupt__" in replay_result
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_interrupt_full_flow(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt fires in subgraph
|
||
result = await graph.ainvoke({"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 = await graph.aget_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 = await graph.aupdate_state(sub_config, {"value": ["sub_forked"]})
|
||
|
||
# Invoke from fork — interrupt should re-fire
|
||
fork_result = await graph.ainvoke(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 = await graph.ainvoke(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"]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_interrupt_full_flow_no_sub_checkpointer(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt, then resume to complete
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert "__interrupt__" in result
|
||
|
||
original_result = await graph.ainvoke(Command(resume="original"), config)
|
||
assert "human:original" in original_result["value"]
|
||
assert "post" in original_result["value"]
|
||
|
||
# Find parent checkpoint before subgraph_node
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_sub = [s for s in history if s.next == ("subgraph_node",)][-1]
|
||
|
||
# Fork from parent checkpoint
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(before_sub.config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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 = await graph.ainvoke(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"]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_to_first_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to a subgraph checkpoint at the FIRST interrupt (async)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("step_a", step_a)
|
||
.add_node("ask_1", ask_1)
|
||
.add_node("ask_2", ask_2)
|
||
.add_edge(START, "step_a")
|
||
.add_edge("step_a", "ask_1")
|
||
.add_edge("ask_1", "ask_2")
|
||
.add_edge("ask_2", "__end__")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until first interrupt (ask_1)
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result["__interrupt__"][0].value == "Question 1?"
|
||
|
||
# Capture subgraph state at the first interrupt
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
sub_config_at_first = parent_state.tasks[0].state.config
|
||
|
||
# Resume through both interrupts to complete
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
# --- Scenario 1: Replay from subgraph checkpoint at 1st interrupt ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, sub_config_at_first)
|
||
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 subgraph checkpoint at 1st interrupt ---
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(sub_config_at_first, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_to_second_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to a subgraph checkpoint at the SECOND interrupt (async)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("step_a", step_a)
|
||
.add_node("ask_1", ask_1)
|
||
.add_node("ask_2", ask_2)
|
||
.add_edge(START, "step_a")
|
||
.add_edge("step_a", "ask_1")
|
||
.add_edge("ask_1", "ask_2")
|
||
.add_edge("ask_2", "__end__")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until first interrupt
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
# Resume first interrupt
|
||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||
assert result["__interrupt__"][0].value == "Question 2?"
|
||
|
||
# Capture subgraph state at the second interrupt
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
sub_config = parent_state.tasks[0].state.config
|
||
|
||
# Resume second interrupt to complete
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
# --- Scenario 1: Replay from subgraph checkpoint at 2nd interrupt ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, sub_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 subgraph checkpoint at 2nd interrupt ---
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(sub_config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_after_completion_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to a subgraph checkpoint AFTER both interrupts resolved (async)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("step_a", step_a)
|
||
.add_node("ask_1", ask_1)
|
||
.add_node("ask_2", ask_2)
|
||
.add_edge(START, "step_a")
|
||
.add_edge("step_a", "ask_1")
|
||
.add_edge("ask_1", "ask_2")
|
||
.add_edge("ask_2", "__end__")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
final_state = await graph.aget_state(config)
|
||
assert len(final_state.tasks) == 0
|
||
|
||
# Replay from the final parent checkpoint — should be a no-op
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(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
|
||
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"]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_from_before_interrupt_then_resume_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Replay from checkpoint before interrupt node, then resume with a new
|
||
answer and verify the graph completes with the new value.
|
||
|
||
Graph: START --> node_a --> ask_human (interrupt) --> node_b --> END
|
||
"""
|
||
|
||
called: list[str] = []
|
||
|
||
async def node_a(state: State) -> State:
|
||
called.append("node_a")
|
||
return {"value": ["a"]}
|
||
|
||
async def ask_human(state: State) -> State:
|
||
called.append("ask_human")
|
||
answer = interrupt("What is your input?")
|
||
return {"value": [f"human:{answer}"]}
|
||
|
||
async def node_b(state: State) -> State:
|
||
called.append("node_b")
|
||
return {"value": ["b"]}
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("node_a", node_a)
|
||
.add_node("ask_human", ask_human)
|
||
.add_node("node_b", node_b)
|
||
.add_edge(START, "node_a")
|
||
.add_edge("node_a", "ask_human")
|
||
.add_edge("ask_human", "node_b")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# --- Original run: invoke until interrupt, then resume to complete ---
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="old_answer"), config)
|
||
|
||
original_history = [s async for s in graph.aget_state_history(config)]
|
||
original = _checkpoint_summary(original_history)
|
||
assert [(s["source"], s["next"], s["values"]) for s in original] == [
|
||
("loop", (), {"value": ["a", "human:old_answer", "b"]}),
|
||
("loop", ("node_b",), {"value": ["a", "human:old_answer"]}),
|
||
("loop", ("ask_human",), {"value": ["a"]}),
|
||
("loop", ("node_a",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
# --- Replay from checkpoint before ask_human ---
|
||
before_ask = next(s for s in original_history if s.next == ("ask_human",))
|
||
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, before_ask.config)
|
||
assert replay_result["__interrupt__"][0].value == "What is your input?"
|
||
assert "ask_human" in called
|
||
assert "node_a" not in called
|
||
|
||
# A fork checkpoint is now the latest
|
||
post_replay = _checkpoint_summary(
|
||
[s async for s in graph.aget_state_history(config)]
|
||
)
|
||
assert [(s["source"], s["next"]) for s in post_replay] == [
|
||
("fork", ("ask_human",)),
|
||
("loop", ()),
|
||
("loop", ("node_b",)),
|
||
("loop", ("ask_human",)),
|
||
("loop", ("node_a",)),
|
||
("input", ("__start__",)),
|
||
]
|
||
|
||
# --- Resume with a new answer ---
|
||
called.clear()
|
||
final_result = await graph.ainvoke(Command(resume="new_answer"), config)
|
||
assert final_result["value"] == ["a", "human:new_answer", "b"]
|
||
assert "ask_human" in called
|
||
assert "node_b" in called
|
||
|
||
final = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in final] == [
|
||
# New branch (from fork)
|
||
("loop", (), {"value": ["a", "human:new_answer", "b"]}),
|
||
("loop", ("node_b",), {"value": ["a", "human:new_answer"]}),
|
||
("fork", ("ask_human",), {"value": ["a"]}),
|
||
# Original branch (preserved)
|
||
("loop", (), {"value": ["a", "human:old_answer", "b"]}),
|
||
("loop", ("node_b",), {"value": ["a", "human:old_answer"]}),
|
||
("loop", ("ask_human",), {"value": ["a"]}),
|
||
("loop", ("node_a",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_resume_from_first_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to a subgraph checkpoint at the first interrupt, then
|
||
resume through both interrupts with new answers.
|
||
|
||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||
"""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("step_a", step_a)
|
||
.add_node("ask_1", ask_1)
|
||
.add_node("ask_2", ask_2)
|
||
.add_edge(START, "step_a")
|
||
.add_edge("step_a", "ask_1")
|
||
.add_edge("ask_1", "ask_2")
|
||
.add_edge("ask_2", "__end__")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# --- Original run: hit both interrupts and resume ---
|
||
await graph.ainvoke({"value": []}, config)
|
||
sub_config_at_first = (
|
||
(await graph.aget_state(config, subgraphs=True)).tasks[0].state.config
|
||
)
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
original = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in original] == [
|
||
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
# --- Time travel to first interrupt's subgraph checkpoint ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, sub_config_at_first)
|
||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||
assert "step_a" not in called
|
||
|
||
# Fork is now the latest parent checkpoint
|
||
post_tt = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"]) for s in post_tt] == [
|
||
("fork", ("executor",)), # <-- new fork (latest)
|
||
("loop", ()), # original done
|
||
("loop", ("executor",)),
|
||
("input", ("__start__",)),
|
||
]
|
||
|
||
# --- Resume both interrupts with new answers ---
|
||
called.clear()
|
||
resume_1 = await graph.ainvoke(Command(resume="new_answer_1"), config)
|
||
assert resume_1["__interrupt__"][0].value == "Question 2?"
|
||
assert "ask_1" in called
|
||
|
||
called.clear()
|
||
resume_2 = await graph.ainvoke(Command(resume="new_answer_2"), config)
|
||
assert resume_2["value"] == [
|
||
"step_a_done",
|
||
"ask_1:new_answer_1",
|
||
"ask_2:new_answer_2",
|
||
]
|
||
|
||
# Verify final history: original branch preserved, new branch appended
|
||
final = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in final] == [
|
||
# New branch (from time travel fork)
|
||
(
|
||
"loop",
|
||
(),
|
||
{"value": ["step_a_done", "ask_1:new_answer_1", "ask_2:new_answer_2"]},
|
||
),
|
||
("fork", ("executor",), {"value": []}),
|
||
# Original branch (preserved)
|
||
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_resume_from_second_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to a subgraph checkpoint at the second interrupt, then
|
||
resume with a new answer. The first interrupt's answer should be preserved.
|
||
|
||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
|
||
"""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("step_a", step_a)
|
||
.add_node("ask_1", ask_1)
|
||
.add_node("ask_2", ask_2)
|
||
.add_edge(START, "step_a")
|
||
.add_edge("step_a", "ask_1")
|
||
.add_edge("ask_1", "ask_2")
|
||
.add_edge("ask_2", "__end__")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# --- Original run: hit both interrupts and resume ---
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
sub_config_at_second = (
|
||
(await graph.aget_state(config, subgraphs=True)).tasks[0].state.config
|
||
)
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
original = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in original] == [
|
||
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
# --- Time travel to second interrupt ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, sub_config_at_second)
|
||
assert replay_result["__interrupt__"][0].value == "Question 2?"
|
||
assert "step_a" not in called
|
||
assert "ask_1" not in called
|
||
|
||
# Fork is now the latest parent checkpoint
|
||
post_tt = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"]) for s in post_tt] == [
|
||
("fork", ("executor",)), # <-- new fork (latest)
|
||
("loop", ()), # original done
|
||
("loop", ("executor",)),
|
||
("input", ("__start__",)),
|
||
]
|
||
|
||
# --- Resume with a new answer for ask_2 only ---
|
||
called.clear()
|
||
resume_result = await graph.ainvoke(Command(resume="new_answer_2"), config)
|
||
assert resume_result["value"] == [
|
||
"step_a_done",
|
||
"ask_1:answer_1",
|
||
"ask_2:new_answer_2",
|
||
]
|
||
|
||
# Verify final history
|
||
final = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in final] == [
|
||
# New branch (from time travel fork)
|
||
(
|
||
"loop",
|
||
(),
|
||
{"value": ["step_a_done", "ask_1:answer_1", "ask_2:new_answer_2"]},
|
||
),
|
||
("fork", ("executor",), {"value": []}),
|
||
# Original branch (preserved)
|
||
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_subgraph_time_travel_checkpoint_pattern_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Verify the checkpoint pattern created by time travel to a subgraph
|
||
interrupt. A fork checkpoint should branch from the replay point.
|
||
|
||
Parent: START --> executor (subgraph, checkpointer=True) --> END
|
||
Executor: START --> ask (interrupt) --> END
|
||
"""
|
||
|
||
async def ask(state: State) -> State:
|
||
answer = interrupt("Q?")
|
||
return {"value": [f"a:{answer}"]}
|
||
|
||
executor = (
|
||
StateGraph(State)
|
||
.add_node("ask", ask)
|
||
.add_edge(START, "ask")
|
||
.compile(checkpointer=True)
|
||
)
|
||
|
||
graph = (
|
||
StateGraph(State)
|
||
.add_node("executor", executor)
|
||
.add_edge(START, "executor")
|
||
.compile(checkpointer=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until interrupt, then complete
|
||
await graph.ainvoke({"value": []}, config)
|
||
sub_config = (await graph.aget_state(config, subgraphs=True)).tasks[0].state.config
|
||
await graph.ainvoke(Command(resume="first"), config)
|
||
|
||
original = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in original] == [
|
||
("loop", (), {"value": ["a:first"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
# Time travel to the interrupt
|
||
await graph.ainvoke(None, sub_config)
|
||
|
||
# Fork is now the latest, branching from the original replay point
|
||
post_tt = [s async for s in graph.aget_state_history(config)]
|
||
post_tt_summary = _checkpoint_summary(post_tt)
|
||
assert [(s["source"], s["next"]) for s in post_tt_summary] == [
|
||
("fork", ("executor",)), # <-- new fork (latest)
|
||
("loop", ()),
|
||
("loop", ("executor",)), # <-- replay point / fork parent
|
||
("input", ("__start__",)),
|
||
]
|
||
# Verify the fork's parent is the original replay point
|
||
replay_point_id = sub_config["configurable"]["checkpoint_map"][""]
|
||
assert post_tt[0].parent_config["configurable"]["checkpoint_id"] == replay_point_id
|
||
|
||
# Resume from the fork
|
||
result = await graph.ainvoke(Command(resume="second"), config)
|
||
assert result["value"] == ["a:second"]
|
||
|
||
final = _checkpoint_summary([s async for s in graph.aget_state_history(config)])
|
||
assert [(s["source"], s["next"], s["values"]) for s in final] == [
|
||
# New branch
|
||
("loop", (), {"value": ["a:second"]}),
|
||
("fork", ("executor",), {"value": []}),
|
||
# Original branch
|
||
("loop", (), {"value": ["a:first"]}),
|
||
("loop", ("executor",), {"value": []}),
|
||
("input", ("__start__",), {"value": []}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_3_levels_deep_time_travel_to_first_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to innermost subgraph checkpoint at FIRST interrupt (async, 3 levels)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result["__interrupt__"][0].value == "Question 1?"
|
||
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
mid_state = parent_state.tasks[0].state
|
||
inner_config = mid_state.tasks[0].state.config
|
||
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
# --- Scenario 1: Replay from innermost checkpoint at 1st interrupt ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(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 = await graph.aupdate_state(inner_config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_3_levels_deep_time_travel_to_second_interrupt_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to innermost subgraph checkpoint at SECOND interrupt (async, 3 levels)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||
assert result["__interrupt__"][0].value == "Question 2?"
|
||
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
mid_state = parent_state.tasks[0].state
|
||
inner_config = mid_state.tasks[0].state.config
|
||
|
||
await graph.ainvoke(Command(resume="answer_2"), config)
|
||
|
||
# --- Scenario 1: Replay ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(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 ---
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(inner_config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_3_levels_deep_time_travel_to_middle_subgraph_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel to the MIDDLE-level subgraph checkpoint (async, 3 levels)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
async def ask_2(state: State) -> State:
|
||
called.append("ask_2")
|
||
answer = interrupt("Question 2?")
|
||
return {"value": [f"ask_2:{answer}"]}
|
||
|
||
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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="answer_1"), config)
|
||
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
mid_config = parent_state.tasks[0].state.config
|
||
|
||
await graph.ainvoke(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 = await graph.ainvoke(None, mid_config)
|
||
assert "__interrupt__" in replay_result
|
||
|
||
# --- Scenario 2: Fork from middle-level subgraph checkpoint ---
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(mid_config, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
assert "__interrupt__" in fork_result
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_3_levels_deep_middle_has_interrupts_async(
|
||
async_checkpointer: BaseCheckpointSaver,
|
||
) -> None:
|
||
"""Time travel when the MIDDLE subgraph itself has interrupts (async)."""
|
||
|
||
called: list[str] = []
|
||
|
||
async def pre(state: State) -> State:
|
||
called.append("pre")
|
||
answer = interrupt("Pre-question?")
|
||
return {"value": [f"pre:{answer}"]}
|
||
|
||
async def step_a(state: State) -> State:
|
||
called.append("step_a")
|
||
return {"value": ["step_a_done"]}
|
||
|
||
async def ask_1(state: State) -> State:
|
||
called.append("ask_1")
|
||
answer = interrupt("Question 1?")
|
||
return {"value": [f"ask_1:{answer}"]}
|
||
|
||
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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run until first interrupt (pre in middle subgraph)
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result["__interrupt__"][0].value == "Pre-question?"
|
||
|
||
# Capture middle subgraph config at the pre interrupt
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
mid_config_at_pre = parent_state.tasks[0].state.config
|
||
|
||
# Resume pre, hits ask_1 in inner subgraph
|
||
result = await graph.ainvoke(Command(resume="pre_answer"), config)
|
||
assert result["__interrupt__"][0].value == "Question 1?"
|
||
|
||
# Capture middle subgraph config at the ask_1 interrupt
|
||
parent_state = await graph.aget_state(config, subgraphs=True)
|
||
mid_config_at_ask1 = parent_state.tasks[0].state.config
|
||
|
||
# Resume ask_1 to complete
|
||
result = await graph.ainvoke(Command(resume="answer_1"), config)
|
||
assert "__interrupt__" not in result
|
||
|
||
# --- Time travel to middle checkpoint at pre interrupt ---
|
||
called.clear()
|
||
replay_result = await graph.ainvoke(None, mid_config_at_pre)
|
||
assert "__interrupt__" in replay_result
|
||
assert replay_result["__interrupt__"][0].value == "Pre-question?"
|
||
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 = await graph.aupdate_state(mid_config_at_pre, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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 = await graph.ainvoke(None, mid_config_at_ask1)
|
||
assert "__interrupt__" in replay_result
|
||
assert replay_result["__interrupt__"][0].value == "Question 1?"
|
||
assert "pre" not in called
|
||
assert "ask_1" in called
|
||
|
||
# Fork from middle checkpoint at ask_1 interrupt
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(mid_config_at_ask1, {"value": ["forked"]})
|
||
fork_result = await graph.ainvoke(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)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_copy_fork_retriggers_interrupt(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="hello"), config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||
|
||
called.clear()
|
||
fork_config = await graph.aupdate_state(before_ask.config, None, as_node="__copy__")
|
||
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
assert "__interrupt__" in fork_result
|
||
assert fork_result["__interrupt__"][0].value == "What is your input?"
|
||
|
||
final = await graph.ainvoke(Command(resume="new_answer"), fork_config)
|
||
assert final == {"value": ["a", "human:new_answer", "b"]}
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_copy_fork_creates_sibling_checkpoint(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"value": []}, config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_b = next(s for s in history if s.next == ("node_b",))
|
||
|
||
# __copy__ fork → source="fork"
|
||
copy_config = await graph.aupdate_state(before_b.config, None, as_node="__copy__")
|
||
copy_state = await graph.aget_state(copy_config)
|
||
assert copy_state.metadata["source"] == "fork"
|
||
|
||
# Regular update → source="update"
|
||
regular_config = await graph.aupdate_state(before_b.config, {"value": ["x"]})
|
||
regular_state = await graph.aget_state(regular_config)
|
||
assert regular_state.metadata["source"] == "update"
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_update_state_with_none_values(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
await graph.ainvoke({"value": []}, config)
|
||
await graph.ainvoke(Command(resume="hello"), config)
|
||
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_ask = [s for s in history if s.next == ("ask_human",)][-1]
|
||
|
||
fork_config = await graph.aupdate_state(before_ask.config, None)
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
assert "__interrupt__" in fork_result
|
||
assert fork_result["__interrupt__"][0].value == "What is your input?"
|
||
|
||
fork_state = await graph.aget_state(fork_config)
|
||
assert fork_state.metadata["source"] == "update"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Section 7: Observability (get_state, config access)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_get_state_with_subgraphs_returns_subgraph_state(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"data": "input"}, config)
|
||
|
||
state = await graph.aget_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"]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_checkpoint_ns_accessible_in_subgraph(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
await graph.ainvoke({"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
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="Python 3.11+ required for async test",
|
||
)
|
||
async def test_stateful_subgraph_retains_state_on_parent_replay(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# === 1st invocation: answer "a1" and "b1" ===
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt
|
||
await graph.ainvoke(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()
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt
|
||
await graph.ainvoke(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 = [s async for s in graph.aget_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 = await graph.ainvoke(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"]})
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="Python 3.11+ required for async test",
|
||
)
|
||
async def test_stateful_subgraph_retains_state_on_parent_fork(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# === 1st invocation: answer "a1" and "b1" ===
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt
|
||
await graph.ainvoke(Command(resume="b1"), config) # completes
|
||
|
||
# === 2nd invocation: answer "a2" and "b2" ===
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt
|
||
await graph.ainvoke(Command(resume="b2"), config) # completes
|
||
|
||
# === Fork from checkpoint before sub_node in 2nd invocation ===
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0]
|
||
|
||
fork_config = await graph.aupdate_state(
|
||
before_sub_2nd.config, {"results": ["forked"]}
|
||
)
|
||
|
||
observed.clear()
|
||
fork_result = await graph.ainvoke(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"]})
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.version_info < (3, 11),
|
||
reason="Python 3.11+ required for async test",
|
||
)
|
||
# ---------------------------------------------------------------------------
|
||
# Section 8: Append-only checkpoint history (branching / forking)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_creates_branch_preserving_old_checkpoints(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result == {"value": ["a", "b1", "c"]}
|
||
|
||
# -- Original checkpoint history (newest first) --
|
||
original_history = [s async for s in graph.aget_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}
|
||
|
||
# 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 = await graph.ainvoke(None, before_b.config)
|
||
assert replay_result == {"value": ["a", "b2", "c"]}
|
||
|
||
# -- Post-replay checkpoint history (newest first) --
|
||
post_replay_history = [s async for s in graph.aget_state_history(config)]
|
||
post_summary = _checkpoint_summary(post_replay_history)
|
||
# 5 original + 1 fork + 2 new branch checkpoints = 8
|
||
assert len(post_summary) == 8
|
||
|
||
assert [s["next"] for s in post_summary] == [
|
||
(), # new branch tip
|
||
("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 = await graph.aget_state(config)
|
||
assert latest.values == {"value": ["a", "b2", "c"]}
|
||
assert latest.config["configurable"]["checkpoint_id"] not in original_ids
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_creates_branch_in_subgraph(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": [], "sub_value": []}, config)
|
||
assert result == {"value": ["p_start", "p_end"], "sub_value": ["sub1"]}
|
||
|
||
# Capture original checkpoint IDs
|
||
original_history = [s async for s in graph.aget_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 = await graph.ainvoke(None, before_sub.config)
|
||
assert replay_result == {"value": ["p_start", "p_end"], "sub_value": ["sub2"]}
|
||
|
||
# Get full history after replay
|
||
post_replay_history = [s async for s in graph.aget_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 = await graph.aget_state(config)
|
||
assert latest.config["configurable"]["checkpoint_id"] in new_ids
|
||
assert latest.values == {"value": ["p_start", "p_end"], "sub_value": ["sub2"]}
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_fork_creates_branch_preserving_old_checkpoints(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": []}, config)
|
||
assert result == {"value": ["a", "b1", "c"]}
|
||
|
||
# -- Original checkpoint history (newest first) --
|
||
original_history = [s async for s in graph.aget_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 = await graph.aupdate_state(before_b.config, {"value": ["x"]})
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
assert fork_result == {"value": ["a", "x", "b2", "c"]}
|
||
|
||
# -- Post-fork checkpoint history (newest first) --
|
||
post_fork_history = [s async for s in graph.aget_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 = await graph.aget_state(config)
|
||
assert latest.values == {"value": ["a", "x", "b2", "c"]}
|
||
assert latest.config["configurable"]["checkpoint_id"] not in original_ids
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_fork_creates_branch_in_subgraph(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
result = await graph.ainvoke({"value": [], "sub_value": []}, config)
|
||
assert result == {"value": ["p_start", "p_end"], "sub_value": ["sub1"]}
|
||
|
||
# Capture original checkpoint IDs
|
||
original_history = [s async for s in graph.aget_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 = await graph.aupdate_state(before_sub.config, {"value": ["extra"]})
|
||
fork_result = await graph.ainvoke(None, fork_config)
|
||
assert fork_result == {
|
||
"value": ["p_start", "extra", "p_end"],
|
||
"sub_value": ["sub2"],
|
||
}
|
||
|
||
# Get full history after fork
|
||
post_fork_history = [s async for s in graph.aget_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 = await graph.aget_state(config)
|
||
assert latest.config["configurable"]["checkpoint_id"] in new_ids
|
||
assert latest.values == {
|
||
"value": ["p_start", "extra", "p_end"],
|
||
"sub_value": ["sub2"],
|
||
}
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_stateless_subgraph_starts_fresh_on_parent_replay(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# === 1st invocation: answer "a1" and "b1" ===
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a1"), config) # hits step_b interrupt
|
||
await graph.ainvoke(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()
|
||
await graph.ainvoke({"results": []}, config) # hits step_a interrupt
|
||
await graph.ainvoke(Command(resume="a2"), config) # hits step_b interrupt
|
||
await graph.ainvoke(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 = [s async for s in graph.aget_state_history(config)]
|
||
before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0]
|
||
|
||
started.clear()
|
||
replay = await graph.ainvoke(None, before_sub_2nd.config)
|
||
|
||
assert "__interrupt__" in replay
|
||
# Stateless subgraph starts completely fresh on replay
|
||
assert started[0] == ("step_a", {"value": []})
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_stateful_subgraph_loads_latest_after_replay(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# 1st invocation — subgraph starts fresh
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed[-1] == ("sub_step", {"value": []})
|
||
|
||
# 2nd invocation — subgraph sees state from 1st
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed[-1] == ("sub_step", {"value": ["s"]})
|
||
|
||
# Replay from checkpoint before parent_node in 2nd invocation
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_parent_2nd = [s for s in history if s.next == ("parent_node",)][0]
|
||
|
||
observed.clear()
|
||
await graph.ainvoke(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()
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed[0] == ("sub_step", {"value": ["s", "s"]})
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_three_level_nested_subgraph_loads_state_on_replay(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# 1st invocation — everything starts fresh
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed == [
|
||
("mid_step", {"mid_trail": []}),
|
||
("inner_step", {"inner_trail": []}),
|
||
]
|
||
|
||
# 2nd invocation — both levels see accumulated state
|
||
observed.clear()
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed == [
|
||
("mid_step", {"mid_trail": ["mid"]}),
|
||
("inner_step", {"inner_trail": ["inner"]}),
|
||
]
|
||
|
||
# Replay from checkpoint before parent_step in 2nd invocation
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_parent_2nd = [s for s in history if s.next == ("parent_step",)][0]
|
||
|
||
observed.clear()
|
||
await graph.ainvoke(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()
|
||
await graph.ainvoke({"results": []}, config)
|
||
assert observed == [
|
||
("mid_step", {"mid_trail": ["mid", "mid"]}),
|
||
("inner_step", {"inner_trail": ["inner", "inner"]}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_three_level_nested_subgraph_loads_state_on_fork(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# 1st invocation
|
||
await graph.ainvoke({"results": []}, config)
|
||
|
||
# 2nd invocation
|
||
await graph.ainvoke({"results": []}, config)
|
||
|
||
# Fork from checkpoint before parent_step in 2nd invocation
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_parent_2nd = [s for s in history if s.next == ("parent_step",)][0]
|
||
fork_config = await graph.aupdate_state(
|
||
before_parent_2nd.config, {"results": ["forked"]}
|
||
)
|
||
|
||
observed.clear()
|
||
await graph.ainvoke(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"]}),
|
||
]
|
||
|
||
|
||
@NEEDS_CONTEXTVARS
|
||
async def test_replay_from_first_invocation_checkpoint(
|
||
async_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=async_checkpointer)
|
||
)
|
||
|
||
config = {"configurable": {"thread_id": "1"}}
|
||
|
||
# Run twice so subgraph accumulates state
|
||
await graph.ainvoke({"results": []}, config)
|
||
await graph.ainvoke({"results": []}, config)
|
||
|
||
# Replay from before sub_node in 1st invocation (furthest back)
|
||
history = [s async for s in graph.aget_state_history(config)]
|
||
before_sub_1st = [s for s in history if s.next == ("sub_node",)][-1]
|
||
|
||
observed.clear()
|
||
await graph.ainvoke(None, before_sub_1st.config)
|
||
# Should see empty state — no prior subgraph checkpoints exist
|
||
assert observed[0] == ("sub_step", {"value": []})
|