WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 09:47:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: time travel when going back to interrupt node (#7498)

Browse Source 
# Fix: Create fork checkpoint on subgraph time travel

## Problem

When time-traveling to a subgraph checkpoint that has an interrupt, and
then resuming, the resume would load the **wrong state** — it would pick
up the original execution's latest checkpoint instead of the
time-traveled one.

This happened because replaying from a subgraph checkpoint never created
a new parent checkpoint. If the replay hit an interrupt before
`after_tick()` ran, no checkpoint was written at all, so the parent's
"latest" checkpoint was still the old one from the original execution.

## Fix

When the loop detects a time-travel replay (not an `update_state` fork),
it now **eagerly writes a fork checkpoint** at the start of the tick.
This ensures:

1. The parent thread's latest checkpoint points to the replayed state
2. Subsequent `Command(resume=...)` calls find the correct checkpoint
3. Stale `INTERRUPT` pending writes from the old checkpoint are cleared
(they reference old task IDs)

Additionally, the subgraph replay logic now uses the **parent checkpoint
ID** (from `prev_checkpoint_config`) when resolving subgraph checkpoints
during time-travel, matching the existing behavior for `update_state`
forks.

## Checkpoint flow diagrams

### Before fix: time travel leaves no fork

```
Original execution:

  C0 (start) --> C1 (step_a) --> C2 (ask_1 interrupt) --> C3 (resume) --> C4 (ask_2 interrupt) --> C5 (done)

Time travel to C2 (subgraph config):

  Replay runs... hits interrupt... no new checkpoint written.
  Parent "latest" is still C5.

  Command(resume="new_answer"):
    Loads C5 (wrong!) instead of the replayed C2 state.
```

### After fix: time travel creates a fork

```
Original execution:

  C0 --> C1 --> C2 --> C3 --> C4 --> C5 (done)

Time travel to C2 (subgraph config):

  C0 --> C1 --> C2 --> C3 --> C4 --> C5
                  \
                   F1 (fork, source="fork")  <-- new latest

  Command(resume="new_answer"):
    Loads F1 (correct!) --> resumes from the right state.

  After full resume:

  C0 --> C1 --> C2 --> C3 --> C4 --> C5
                  \
                   F1 --> F2 (ask_1 result) --> F3 (ask_2 interrupt) --> F4 (done)
```

### Manual fork via `update_state` (unchanged)

```
  C0 --> C1 --> C2 --> C3
                  \
                   U1 (source="update")  <-- created by update_state()

  This path already worked. The fix skips update/fork sources
  so existing behavior is preserved.
```

## Changes

- **`libs/langgraph/langgraph/pregel/_loop.py`**:
- Extract `is_time_traveling` flag from the existing replay detection
logic for reuse
- Write a fork checkpoint (`source="fork"`) eagerly at the start of a
time-travel tick, before execution begins
- Clear stale `INTERRUPT` pending writes when creating the fork (they
reference old task IDs that won't match the new checkpoint)
- Unify subgraph replay ID resolution: check `source in ("update",
"fork")` instead of a separate `is_time_traveling` condition, since the
new fork checkpoint now has `source="fork"`
- **`libs/langgraph/tests/test_time_travel.py`** and
**`test_time_travel_async.py`**: Added 4 new test cases (sync + async):
- `test_replay_from_before_interrupt_then_resume` — replays from a
checkpoint before an interrupt, resumes with a new answer, and verifies
the full checkpoint history (source, next, values) at each stage
- `test_subgraph_time_travel_resume_from_first_interrupt` — time-travels
to a subgraph's first interrupt, resumes both interrupts with new
answers, and verifies the fork creates a new branch while preserving the
original
- `test_subgraph_time_travel_resume_from_second_interrupt` —
time-travels to a subgraph's second interrupt, resumes with a new
answer, and verifies the first interrupt's original answer is preserved
- `test_subgraph_time_travel_checkpoint_pattern` — verifies the fork
checkpoint branches from the correct replay point and that the full
checkpoint tree is correct after resume
- **`libs/langgraph/tests/test_pregel.py`** /
**`test_pregel_async.py`**: Updated existing
`test_weather_subgraph_state` to account for the new fork checkpoint
appearing in history (history length increases by 1)
This commit is contained in:
Sydney Runkle
2026-04-16 08:29:48 -04:00
committed by GitHub
parent 4d64227c13
commit 51cbdbd5cd
5 changed files with 1171 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_loop.py
+29 -4
View File
@@ -692,7 +692,7 @@ class PregelLoop:
# writes so that interrupt() calls re-fire instead of returning
# stale values. But if we're actively resuming, keep them —
# multi-interrupt scenarios need previously resolved values preserved.
if self.is_replaying and (
is_time_traveling = self.is_replaying and (
# Time-travel to a subgraph checkpoint: the parent sets
# RESUMING=True (it can't distinguish time-travel from resume),
# so we check if this subgraph's own ns is in checkpoint_map.
@@ -710,7 +710,8 @@ class PregelLoop:
# (subgraph input is a Send arg, not a Command)
or configurable.get(CONFIG_KEY_RESUMING, False)
)
):
)
if is_time_traveling:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
]
@@ -765,6 +766,26 @@ class PregelLoop:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
# When time-traveling (replaying from a specific checkpoint),
# save a fork checkpoint so the replayed execution creates a
# new branch. Without this, if the execution hits an interrupt
# before after_tick() runs, no new checkpoint is created —
# the parent's latest checkpoint remains the old one and
# subsequent resumes load the wrong state.
# Skip for update_state forks (source=update/fork) since they
# already have their own fork checkpoint.
if is_time_traveling and self.checkpoint_metadata.get("source") not in (
"update",
"fork",
):
# Clear old INTERRUPT writes from the loaded checkpoint.
# The fork will have a new checkpoint_id which changes
# task IDs — stale interrupt writes would accumulate and
# confuse the multiple-interrupt check in future resumes.
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != INTERRUPT
]
self._put_checkpoint({"source": "fork"})
# produce values output
self._emit(
"values", map_output_values, self.output_keys, True, self.channels
@@ -807,14 +828,18 @@ class PregelLoop:
if not self.is_nested:
# Pass the resolved before-bound checkpoint ID so subgraphs can
# find their corresponding checkpoint without re-fetching the
# parent. For forks (source=update), use the fork's parent
# parent. For forks (source=update/fork), use the fork's parent
# checkpoint ID since the fork was created after the subgraph's
# checkpoints from the original execution.
replay_state: ReplayState | None = None
if self.is_replaying:
replay_checkpoint_id = self.checkpoint["id"]
if (
self.checkpoint_metadata.get("source") == "update"
self.checkpoint_metadata.get("source")
in (
"update",
"fork",
)
and self.prev_checkpoint_config
):
replay_checkpoint_id = self.prev_checkpoint_config[CONF].get(
libs/langgraph/tests/test_pregel.py
+6 -3
View File
@@ -615,8 +615,11 @@ def test_run_from_checkpoint_id_retains_previous_writes(
)
]
assert len(new_history) == len(history) + 1
for original, new in zip(history, new_history[1:]):
# +2: one fork checkpoint from time travel, one from the new execution
assert len(new_history) == len(history) + 2
# new_history[0] is the new execution result, new_history[1] is the fork
assert new_history[1].metadata["source"] == "fork"
for original, new in zip(history, new_history[2:]):
assert original.values == new.values
assert original.next == new.next
assert original.metadata["step"] == new.metadata["step"]
@@ -624,7 +627,7 @@ def test_run_from_checkpoint_id_retains_previous_writes(
def _get_tasks(hist: list, start: int):
return [h.tasks for h in hist[start:]]
assert _get_tasks(new_history, 1) == _get_tasks(history, 0)
assert _get_tasks(new_history, 2) == _get_tasks(history, 0)
def test_batch_two_processes_in_out() -> None:
libs/langgraph/tests/test_pregel_async.py
+6 -3
View File
@@ -2086,8 +2086,11 @@ async def test_run_from_checkpoint_id_retains_previous_writes(
)
]
assert len(new_history) == len(history) + 1
for original, new in zip(history, new_history[1:]):
# +2: one fork checkpoint from time travel, one from the new execution
assert len(new_history) == len(history) + 2
# new_history[0] is the new execution result, new_history[1] is the fork
assert new_history[1].metadata["source"] == "fork"
for original, new in zip(history, new_history[2:]):
assert original.values == new.values
assert original.next == new.next
assert original.metadata["step"] == new.metadata["step"]
@@ -2095,7 +2098,7 @@ async def test_run_from_checkpoint_id_retains_previous_writes(
def _get_tasks(hist: list, start: int):
return [h.tasks for h in hist[start:]]
assert _get_tasks(new_history, 1) == _get_tasks(history, 0)
assert _get_tasks(new_history, 2) == _get_tasks(history, 0)
async def test_cond_edge_after_send() -> None:
libs/langgraph/tests/test_time_travel.py
+727 -7
View File
@@ -37,6 +37,7 @@ def _checkpoint_summary(history: list) -> list[dict]:
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
"""
@@ -52,6 +53,7 @@ def _checkpoint_summary(history: list) -> list[dict]:
{
"id": cid[-6:],
"parent_id": pid[-6:] if pid else None,
"source": s.metadata.get("source"),
"next": s.next,
"values": s.values,
}
@@ -280,6 +282,116 @@ def test_replay_from_before_interrupt_refires(
assert call_count["node_b"] == 1 # NOT re-executed (after interrupt)
def test_replay_from_before_interrupt_then_resume(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Replay from checkpoint before interrupt node, then resume with a new
answer and verify the graph completes with the new value.
Graph: START --> node_a --> ask_human (interrupt) --> node_b --> END
Original run:
source=input next=(__start__,) values=[]
source=loop next=(node_a,) values=[]
source=loop next=(ask_human,) values=[a] <-- replay from here
source=loop next=(node_b,) values=[a, human:old_answer]
source=loop next=() values=[a, human:old_answer, b]
After replay (fork created) + resume with "new_answer":
source=input next=(__start__,) values=[]
source=loop next=(node_a,) values=[]
source=loop next=(ask_human,) values=[a] <-- branch point
source=loop next=(node_b,) values=[a, human:old_answer]
source=loop next=() values=[a, human:old_answer, b] (old branch)
source=fork next=(ask_human,) values=[a] <-- fork from branch point
source=loop next=(node_b,) values=[a, human:new_answer]
source=loop next=() values=[a, human:new_answer, b] (new branch)
"""
called: list[str] = []
def node_a(state: State) -> State:
called.append("node_a")
return {"value": ["a"]}
def ask_human(state: State) -> State:
called.append("ask_human")
answer = interrupt("What is your input?")
return {"value": [f"human:{answer}"]}
def node_b(state: State) -> State:
called.append("node_b")
return {"value": ["b"]}
graph = (
StateGraph(State)
.add_node("node_a", node_a)
.add_node("ask_human", ask_human)
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge("node_a", "ask_human")
.add_edge("ask_human", "node_b")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# --- Original run: invoke until interrupt, then resume to complete ---
graph.invoke({"value": []}, config)
graph.invoke(Command(resume="old_answer"), config)
original_history = list(graph.get_state_history(config))
original = _checkpoint_summary(original_history)
assert [(s["source"], s["next"], s["values"]) for s in original] == [
("loop", (), {"value": ["a", "human:old_answer", "b"]}),
("loop", ("node_b",), {"value": ["a", "human:old_answer"]}),
("loop", ("ask_human",), {"value": ["a"]}),
("loop", ("node_a",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
# --- Replay from checkpoint before ask_human ---
before_ask = next(s for s in original_history if s.next == ("ask_human",))
called.clear()
replay_result = graph.invoke(None, before_ask.config)
assert replay_result["__interrupt__"][0].value == "What is your input?"
assert "ask_human" in called
assert "node_a" not in called # before the replay point, not re-executed
# A fork checkpoint is now the latest — it branches from the replay point
post_replay = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"]) for s in post_replay] == [
("fork", ("ask_human",)), # <-- new fork (latest)
("loop", ()), # original done
("loop", ("node_b",)),
("loop", ("ask_human",)), # branch point
("loop", ("node_a",)),
("input", ("__start__",)),
]
# --- Resume with a new answer ---
called.clear()
final_result = graph.invoke(Command(resume="new_answer"), config)
assert final_result["value"] == ["a", "human:new_answer", "b"]
assert "ask_human" in called
assert "node_b" in called
final = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in final] == [
# New branch (from fork)
("loop", (), {"value": ["a", "human:new_answer", "b"]}),
("loop", ("node_b",), {"value": ["a", "human:new_answer"]}),
("fork", ("ask_human",), {"value": ["a"]}),
# Original branch (preserved)
("loop", (), {"value": ["a", "human:old_answer", "b"]}),
("loop", ("node_b",), {"value": ["a", "human:old_answer"]}),
("loop", ("ask_human",), {"value": ["a"]}),
("loop", ("node_a",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
def test_replay_interrupt_stable_across_replays(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
@@ -320,8 +432,14 @@ def test_replay_interrupt_stable_across_replays(
r = graph.invoke(None, before_ask.config)
results.append(r)
assert all(r == results[0] for r in results)
assert "__interrupt__" in results[0]
# Each replay creates a fork with a unique interrupt ID, so we compare
# interrupt values and state values rather than full equality.
assert all("__interrupt__" in r for r in results)
assert all(
r["__interrupt__"][0].value == results[0]["__interrupt__"][0].value
for r in results
)
assert all(r["value"] == results[0]["value"] for r in results)
def test_fork_from_before_interrupt_refires(
@@ -854,6 +972,290 @@ def test_subgraph_interrupt_replay_from_interrupt_checkpoint(
assert "step_b" not in called
def test_subgraph_interrupt_replay_from_parent_then_resume(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Replay from the parent checkpoint where a subgraph interrupt fired,
then resume with a new answer. Verifies that a fork is created and the
full graph completes. Checks full checkpoint history at each stage."""
called: list[str] = []
def router(state: State) -> State:
called.append("router")
return {"value": ["routed"]}
def step_a(state: State) -> State:
called.append("step_a")
return {"value": ["sub_a"]}
def ask_human(state: State) -> State:
called.append("ask_human")
answer = interrupt("Provide input:")
return {"value": [f"human:{answer}"]}
def step_b(state: State) -> State:
called.append("step_b")
return {"value": ["sub_b"]}
subgraph = (
StateGraph(State)
.add_node("step_a", step_a)
.add_node("ask_human", ask_human)
.add_node("step_b", step_b)
.add_edge(START, "step_a")
.add_edge("step_a", "ask_human")
.add_edge("ask_human", "step_b")
.compile(checkpointer=True)
)
def post_process(state: State) -> State:
called.append("post_process")
return {"value": ["post"]}
graph = (
StateGraph(State)
.add_node("router", router)
.add_node("subgraph_node", subgraph)
.add_node("post_process", post_process)
.add_edge(START, "router")
.add_edge("router", "subgraph_node")
.add_edge("subgraph_node", "post_process")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# Run until interrupt, then resume to complete
graph.invoke({"value": []}, config)
graph.invoke(Command(resume="old_answer"), config)
# Original parent history (newest first)
original_history = list(graph.get_state_history(config))
assert [s.next for s in original_history] == [
(), # done
("post_process",),
("subgraph_node",), # subgraph ran, interrupt fired here
("router",),
("__start__",),
]
# Find the parent checkpoint where the interrupt fired
interrupt_checkpoint = next(
s for s in original_history if s.next == ("subgraph_node",)
)
# Replay from parent checkpoint — subgraph re-executes, interrupt re-fires
called.clear()
replay_result = graph.invoke(None, interrupt_checkpoint.config)
assert "__interrupt__" in replay_result
assert replay_result["__interrupt__"][0].value == "Provide input:"
assert "step_a" in called
assert "ask_human" in called
assert "step_b" not in called
# Verify fork checkpoint was created
post_replay_history = list(graph.get_state_history(config))
assert [s.next for s in post_replay_history] == [
("subgraph_node",), # fork (interrupt pending)
(), # original done
("post_process",),
("subgraph_node",),
("router",),
("__start__",),
]
assert [s.metadata["source"] for s in post_replay_history] == [
"fork",
"loop",
"loop",
"loop",
"loop",
"input",
]
fork = post_replay_history[0]
assert (
fork.parent_config["configurable"]["checkpoint_id"]
== interrupt_checkpoint.config["configurable"]["checkpoint_id"]
)
# Resume with a new answer — full graph should complete
called.clear()
final_result = graph.invoke(Command(resume="new_answer"), config)
assert "__interrupt__" not in final_result
assert "human:new_answer" in final_result["value"]
assert "sub_b" in final_result["value"]
assert "post" in final_result["value"]
assert "ask_human" in called
assert "step_b" in called
assert "post_process" in called
# Final checkpoint history
final_history = list(graph.get_state_history(config))
assert [s.next for s in final_history] == [
(), # new branch done
("post_process",), # new branch post_process
("subgraph_node",), # fork
(), # original done
("post_process",),
("subgraph_node",),
("router",),
("__start__",),
]
assert [s.metadata["source"] for s in final_history] == [
"loop",
"loop",
"fork",
"loop",
"loop",
"loop",
"loop",
"input",
]
def test_subgraph_replay_loads_accumulated_state_then_resume(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Two parent invocations, then replay from before the subgraph in the
2nd invocation. The subgraph (checkpointer=True) should load its
accumulated state from the 1st invocation via ReplayState, re-fire
the interrupt, and then resume + complete.
This tests the ReplayState path: the parent is replaying and the
subgraph uses list(before=parent_checkpoint_id) to find its
corresponding checkpoint from the original execution.
"""
class SubState(TypedDict):
value: Annotated[list[str], operator.add]
class ParentState(TypedDict):
results: Annotated[list[str], operator.add]
started_state: list[dict] = []
def step_a(state: SubState) -> SubState:
started_state.append(dict(state))
answer = interrupt("question_a")
return {"value": [f"a:{answer}"]}
subgraph = (
StateGraph(SubState)
.add_node("step_a", step_a)
.add_edge(START, "step_a")
.compile(checkpointer=True)
)
def parent_node(state: ParentState) -> ParentState:
return {"results": ["p"]}
graph = (
StateGraph(ParentState)
.add_node("parent_node", parent_node)
.add_node("sub_node", subgraph)
.add_edge(START, "parent_node")
.add_edge("parent_node", "sub_node")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# === 1st invocation: complete with answer "a1" ===
graph.invoke({"results": []}, config)
graph.invoke(Command(resume="a1"), config)
# step_a saw empty state (fresh subgraph)
assert started_state[0] == {"value": []}
# === 2nd invocation: complete with answer "a2" ===
started_state.clear()
graph.invoke({"results": []}, config)
graph.invoke(Command(resume="a2"), config)
# Stateful subgraph retained state from 1st invocation
assert started_state[0] == {"value": ["a:a1"]}
# Original history (newest first)
original_history = list(graph.get_state_history(config))
assert [s.next for s in original_history] == [
(), # 2nd done
("sub_node",), # 2nd sub_node
("parent_node",), # 2nd parent_node
("__start__",), # 2nd input
(), # 1st done
("sub_node",), # 1st sub_node
("parent_node",), # 1st parent_node
("__start__",), # 1st input
]
# Replay from before sub_node in 2nd invocation (newest match)
before_sub_2nd = [s for s in original_history if s.next == ("sub_node",)][0]
started_state.clear()
replay = graph.invoke(None, before_sub_2nd.config)
assert "__interrupt__" in replay
# Subgraph should see accumulated state from END of 1st invocation
assert started_state[0] == {"value": ["a:a1"]}
# Verify fork was created
post_replay_history = list(graph.get_state_history(config))
assert [s.next for s in post_replay_history] == [
("sub_node",), # fork (interrupt pending)
(), # 2nd done
("sub_node",), # 2nd sub_node
("parent_node",), # 2nd parent_node
("__start__",), # 2nd input
(), # 1st done
("sub_node",), # 1st sub_node
("parent_node",), # 1st parent_node
("__start__",), # 1st input
]
assert [s.metadata["source"] for s in post_replay_history] == [
"fork",
"loop",
"loop",
"loop",
"input",
"loop",
"loop",
"loop",
"input",
]
# Resume with a new answer
started_state.clear()
final = graph.invoke(Command(resume="a3"), config)
assert "__interrupt__" not in final
assert final["results"] == ["p", "p"]
# Final history
final_history = list(graph.get_state_history(config))
assert [s.next for s in final_history] == [
(), # new branch done
("sub_node",), # fork
(), # 2nd done
("sub_node",), # 2nd sub_node
("parent_node",), # 2nd parent_node
("__start__",), # 2nd input
(), # 1st done
("sub_node",), # 1st sub_node
("parent_node",), # 1st parent_node
("__start__",), # 1st input
]
assert [s.metadata["source"] for s in final_history] == [
"loop",
"fork",
"loop",
"loop",
"loop",
"input",
"loop",
"loop",
"loop",
"input",
]
def test_subgraph_interrupt_full_flow(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
@@ -1290,6 +1692,321 @@ def test_subgraph_time_travel_to_second_interrupt(
assert "ask_1" not in called
def test_subgraph_time_travel_resume_from_first_interrupt(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Time travel to a subgraph checkpoint at the first interrupt, then
resume through both interrupts with new answers.
This verifies the key bug fix: after time-traveling to a subgraph
checkpoint with an interrupt, a fork checkpoint is created so that
subsequent resumes find the correct state (not the old branch tip).
Parent: START --> executor (subgraph, checkpointer=True) --> END
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
Parent history after original run completes:
source=input next=(__start__,) values=[]
source=loop next=(executor,) values=[]
source=loop next=() values=[step_a_done, ask_1:answer_1, ask_2:answer_2]
After time-traveling to 1st interrupt + resuming with new answers:
source=input next=(__start__,) values=[]
source=loop next=(executor,) values=[] <-- branch point
source=loop next=() values=[..., ask_2:answer_2] (old branch)
source=fork next=(executor,) values=[] <-- fork from time travel
source=loop next=() values=[..., ask_2:new_answer_2] (new branch)
"""
called: list[str] = []
def step_a(state: State) -> State:
called.append("step_a")
return {"value": ["step_a_done"]}
def ask_1(state: State) -> State:
called.append("ask_1")
answer = interrupt("Question 1?")
return {"value": [f"ask_1:{answer}"]}
def ask_2(state: State) -> State:
called.append("ask_2")
answer = interrupt("Question 2?")
return {"value": [f"ask_2:{answer}"]}
executor = (
StateGraph(State)
.add_node("step_a", step_a)
.add_node("ask_1", ask_1)
.add_node("ask_2", ask_2)
.add_edge(START, "step_a")
.add_edge("step_a", "ask_1")
.add_edge("ask_1", "ask_2")
.add_edge("ask_2", "__end__")
.compile(checkpointer=True)
)
graph = (
StateGraph(State)
.add_node("executor", executor)
.add_edge(START, "executor")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# --- Original run: hit both interrupts and resume ---
graph.invoke({"value": []}, config)
sub_config_at_first = graph.get_state(config, subgraphs=True).tasks[0].state.config
graph.invoke(Command(resume="answer_1"), config)
graph.invoke(Command(resume="answer_2"), config)
original = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in original] == [
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
# --- Time travel to first interrupt's subgraph checkpoint ---
called.clear()
replay_result = graph.invoke(None, sub_config_at_first)
assert replay_result["__interrupt__"][0].value == "Question 1?"
assert "step_a" not in called # before interrupt, not re-executed
# Fork is now the latest parent checkpoint
post_tt = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"]) for s in post_tt] == [
("fork", ("executor",)), # <-- new fork (latest)
("loop", ()), # original done
("loop", ("executor",)),
("input", ("__start__",)),
]
# --- Resume both interrupts with new answers ---
called.clear()
resume_1 = graph.invoke(Command(resume="new_answer_1"), config)
assert resume_1["__interrupt__"][0].value == "Question 2?"
assert "ask_1" in called
called.clear()
resume_2 = graph.invoke(Command(resume="new_answer_2"), config)
assert resume_2["value"] == [
"step_a_done",
"ask_1:new_answer_1",
"ask_2:new_answer_2",
]
# Verify final history: original branch preserved, new branch appended
final = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in final] == [
# New branch (from time travel fork)
(
"loop",
(),
{"value": ["step_a_done", "ask_1:new_answer_1", "ask_2:new_answer_2"]},
),
("fork", ("executor",), {"value": []}),
# Original branch (preserved)
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
def test_subgraph_time_travel_resume_from_second_interrupt(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Time travel to a subgraph checkpoint at the second interrupt, then
resume with a new answer. The first interrupt's answer should be preserved.
Parent: START --> executor (subgraph, checkpointer=True) --> END
Executor: START --> step_a --> ask_1 (interrupt) --> ask_2 (interrupt) --> END
Key assertion: after resuming from a time-travel to the 2nd interrupt,
the final state keeps ask_1's original answer but uses the new ask_2 answer.
"""
called: list[str] = []
def step_a(state: State) -> State:
called.append("step_a")
return {"value": ["step_a_done"]}
def ask_1(state: State) -> State:
called.append("ask_1")
answer = interrupt("Question 1?")
return {"value": [f"ask_1:{answer}"]}
def ask_2(state: State) -> State:
called.append("ask_2")
answer = interrupt("Question 2?")
return {"value": [f"ask_2:{answer}"]}
executor = (
StateGraph(State)
.add_node("step_a", step_a)
.add_node("ask_1", ask_1)
.add_node("ask_2", ask_2)
.add_edge(START, "step_a")
.add_edge("step_a", "ask_1")
.add_edge("ask_1", "ask_2")
.add_edge("ask_2", "__end__")
.compile(checkpointer=True)
)
graph = (
StateGraph(State)
.add_node("executor", executor)
.add_edge(START, "executor")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# --- Original run: hit both interrupts and resume ---
graph.invoke({"value": []}, config)
graph.invoke(Command(resume="answer_1"), config)
sub_config_at_second = graph.get_state(config, subgraphs=True).tasks[0].state.config
graph.invoke(Command(resume="answer_2"), config)
original = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in original] == [
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
# --- Time travel to second interrupt ---
called.clear()
replay_result = graph.invoke(None, sub_config_at_second)
assert replay_result["__interrupt__"][0].value == "Question 2?"
assert "step_a" not in called
assert "ask_1" not in called # already resolved, not re-executed
# Fork is now the latest parent checkpoint
post_tt = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"]) for s in post_tt] == [
("fork", ("executor",)), # <-- new fork (latest)
("loop", ()), # original done
("loop", ("executor",)),
("input", ("__start__",)),
]
# --- Resume with a new answer for ask_2 only ---
called.clear()
resume_result = graph.invoke(Command(resume="new_answer_2"), config)
# ask_1's original answer preserved, ask_2 uses the new answer
assert resume_result["value"] == [
"step_a_done",
"ask_1:answer_1",
"ask_2:new_answer_2",
]
# Verify final history: original branch preserved, new branch appended
final = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in final] == [
# New branch (from time travel fork)
(
"loop",
(),
{"value": ["step_a_done", "ask_1:answer_1", "ask_2:new_answer_2"]},
),
("fork", ("executor",), {"value": []}),
# Original branch (preserved)
("loop", (), {"value": ["step_a_done", "ask_1:answer_1", "ask_2:answer_2"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
def test_subgraph_time_travel_checkpoint_pattern(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Verify the checkpoint pattern created by time travel to a subgraph
interrupt. A fork checkpoint should branch from the replay point and
become the latest parent checkpoint.
Parent: START --> executor (subgraph, checkpointer=True) --> END
Executor: START --> ask (interrupt) --> END
Original run (after completing):
source=input next=(__start__,) values=[]
source=loop next=(executor,) values=[] <-- replay point
source=loop next=() values=[a:first]
After time travel to interrupt + resume with "second":
source=input next=(__start__,) values=[]
source=loop next=(executor,) values=[] <-- branch point
source=loop next=() values=[a:first] (old branch)
source=fork next=(executor,) values=[] <-- fork
source=loop next=() values=[a:second] (new branch)
"""
def ask(state: State) -> State:
answer = interrupt("Q?")
return {"value": [f"a:{answer}"]}
executor = (
StateGraph(State)
.add_node("ask", ask)
.add_edge(START, "ask")
.compile(checkpointer=True)
)
graph = (
StateGraph(State)
.add_node("executor", executor)
.add_edge(START, "executor")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# Run until interrupt, then complete
graph.invoke({"value": []}, config)
sub_config = graph.get_state(config, subgraphs=True).tasks[0].state.config
graph.invoke(Command(resume="first"), config)
original = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in original] == [
("loop", (), {"value": ["a:first"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
# Time travel to the interrupt
graph.invoke(None, sub_config)
# Fork is now the latest, branching from the original replay point
post_tt = list(graph.get_state_history(config))
post_tt_summary = _checkpoint_summary(post_tt)
assert [(s["source"], s["next"]) for s in post_tt_summary] == [
("fork", ("executor",)), # <-- new fork (latest)
("loop", ()),
("loop", ("executor",)), # <-- replay point / fork parent
("input", ("__start__",)),
]
# Verify the fork's parent is the original replay point
replay_point_id = sub_config["configurable"]["checkpoint_map"][""]
assert post_tt[0].parent_config["configurable"]["checkpoint_id"] == replay_point_id
# Resume from the fork — graph completes with new answer
result = graph.invoke(Command(resume="second"), config)
assert result["value"] == ["a:second"]
final = _checkpoint_summary(list(graph.get_state_history(config)))
assert [(s["source"], s["next"], s["values"]) for s in final] == [
# New branch
("loop", (), {"value": ["a:second"]}),
("fork", ("executor",), {"value": []}),
# Original branch
("loop", (), {"value": ["a:first"]}),
("loop", ("executor",), {"value": []}),
("input", ("__start__",), {"value": []}),
]
def test_subgraph_time_travel_after_completion(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
@@ -2283,14 +3000,16 @@ def test_replay_creates_branch_preserving_old_checkpoints(
# -- Post-replay checkpoint history (newest first) --
post_replay_history = list(graph.get_state_history(config))
post_summary = _checkpoint_summary(post_replay_history)
assert len(post_summary) == 7 # 5 original + 2 new branch checkpoints
# 5 original + 1 fork + 2 new branch checkpoints = 8
assert len(post_summary) == 8
# Verify the full shape after replay
assert [s["next"] for s in post_summary] == [
(), # new branch tip (C6)
("node_c",), # new branch (C5)
(), # old branch tip (C4)
("node_c",), # old (C3)
(), # 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)
@@ -2298,6 +3017,7 @@ def test_replay_creates_branch_preserving_old_checkpoints(
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
libs/langgraph/tests/test_time_travel_async.py
+403 -7
View File
@@ -46,6 +46,7 @@ def _checkpoint_summary(history: list) -> list[dict]:
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
"""
@@ -61,6 +62,7 @@ def _checkpoint_summary(history: list) -> list[dict]:
{
"id": cid[-6:],
"parent_id": pid[-6:] if pid else None,
"source": s.metadata.get("source"),
"next": s.next,
"values": s.values,
}
@@ -335,8 +337,14 @@ async def test_replay_interrupt_stable_across_replays(
r = await graph.ainvoke(None, before_ask.config)
results.append(r)
assert all(r == results[0] for r in results)
assert "__interrupt__" in results[0]
# 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
@@ -1261,6 +1269,391 @@ async def test_subgraph_time_travel_after_completion_async(
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,
@@ -2088,13 +2481,15 @@ async def test_replay_creates_branch_preserving_old_checkpoints(
# -- 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)
assert len(post_summary) == 7 # 5 original + 2 new branch checkpoints
# 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 (C6)
("node_c",), # new branch (C5)
(), # old branch tip (C4)
("node_c",), # old (C3)
(), # 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)
@@ -2102,6 +2497,7 @@ async def test_replay_creates_branch_preserving_old_checkpoints(
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