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

eh

Browse Source
This commit is contained in:
Sydney Runkle
2026-03-05 18:08:07 -08:00
parent 625c51e74c
commit 014f9d2a09
3 changed files with 118 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_loop.py
+54 -61
View File
@@ -1133,35 +1133,40 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
return parent_saved.parent_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
return parent_checkpoint_id
def _get_checkpoint_before_parent(self) -> CheckpointTuple | None:
"""Find the subgraph checkpoint that was current at the parent's
checkpoint time, using the parent checkpoint_id as an upper bound.
def _get_checkpoint_after_parent(self) -> CheckpointTuple | None:
"""Find the right subgraph checkpoint to restore when the parent replays.
We return a checkpoint with the historical channel_values but empty
channel_versions and versions_seen (via `empty_checkpoint()`). This is
intentional: the node scheduling logic in `_triggers()` decides whether
to run a node by comparing channel_versions against that node's
versions_seen. By clearing both, every channel looks "new" to every
node, which forces all nodes to re-trigger. Without this, the restored
checkpoint's versions would show all nodes as up-to-date and nothing
would be scheduled to run.
Each time the parent invokes a subgraph, the subgraph creates a series
of checkpoints. Every checkpoint records which parent checkpoint was
active when it was created (in `metadata["parents"]`). The first
checkpoint in each invocation has `source="input"` and contains the
accumulated channel_values from prior invocations but hasn't run any
nodes yet.
We query for `source="input"` + `parents={parent_ns: parent_id}` to
find the starting checkpoint from the invocation that ran under the
given parent checkpoint — one bounded query, one result.
We then clear `versions_seen` so all nodes re-trigger from that state.
The existing `is_replaying` logic in `_first()` handles dropping any
cached RESUME writes so that interrupts re-fire.
Returns None to start fresh if no such checkpoint exists."""
parent_checkpoint_id = self._get_parent_checkpoint_id()
if parent_checkpoint_id and self.checkpointer:
before_config: RunnableConfig = {
CONF: {"checkpoint_id": parent_checkpoint_id}
}
parent_ns = (
NS_SEP.join(self.checkpoint_ns[:-1]) if self.checkpoint_ns else ""
)
for saved in self.checkpointer.list(
self.checkpoint_config, before=before_config, limit=1
self.checkpoint_config,
filter={
"source": "input",
"parents": {parent_ns: parent_checkpoint_id},
},
limit=1,
):
checkpoint = empty_checkpoint()
checkpoint["channel_values"] = saved.checkpoint.get(
"channel_values", {}
)
return CheckpointTuple(
self.checkpoint_config, checkpoint, {"step": -2}, None, []
)
saved.checkpoint["versions_seen"] = {}
return saved
return None
# context manager
@@ -1171,16 +1176,13 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
saved = self.checkpointer.get_tuple(self.checkpoint_config)
else:
saved = None
# When replaying a subgraph that wasn't in the checkpoint map
# (parent checkpoint predates this subgraph), start fresh.
# For stateful subgraphs (checkpointer=True), find the checkpoint
# that was current at the parent's checkpoint time.
if (
saved is not None
and self.config[CONF].get(CONFIG_KEY_REPLAYING)
and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID)
):
saved = self._get_checkpoint_before_parent()
# When replaying a subgraph, find the checkpoint that was current
# at the parent's checkpoint time. For stateless subgraphs (no
# checkpointer), this returns None and we start fresh as usual.
if self.config[CONF].get(
CONFIG_KEY_REPLAYING
) and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID):
saved = self._get_checkpoint_after_parent()
if saved is None:
saved = CheckpointTuple(
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
@@ -1372,29 +1374,23 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
return parent_saved.parent_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
return parent_checkpoint_id
async def _aget_checkpoint_before_parent(self) -> CheckpointTuple | None:
"""Async version of `_get_checkpoint_before_parent`. See that method's
docstring for details on why we use empty_checkpoint() with only
channel_values restored."""
async def _aget_checkpoint_after_parent(self) -> CheckpointTuple | None:
"""Async version of `_get_checkpoint_after_parent`."""
parent_checkpoint_id = await self._aget_parent_checkpoint_id()
if parent_checkpoint_id and self.checkpointer:
before_config: RunnableConfig = {
CONF: {"checkpoint_id": parent_checkpoint_id}
}
parent_ns = (
NS_SEP.join(self.checkpoint_ns[:-1]) if self.checkpoint_ns else ""
)
async for saved in self.checkpointer.alist(
self.checkpoint_config, before=before_config, limit=1
self.checkpoint_config,
filter={
"source": "input",
"parents": {parent_ns: parent_checkpoint_id},
},
limit=1,
):
checkpoint = empty_checkpoint()
checkpoint["channel_values"] = saved.checkpoint.get(
"channel_values", {}
)
return CheckpointTuple(
self.checkpoint_config,
checkpoint,
{"step": -2},
None,
[],
)
saved.checkpoint["versions_seen"] = {}
return saved
return None
# context manager
@@ -1404,16 +1400,13 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
else:
saved = None
# When replaying a subgraph that wasn't in the checkpoint map
# (parent checkpoint predates this subgraph), start fresh.
# For stateful subgraphs (checkpointer=True), find the checkpoint
# that was current at the parent's checkpoint time.
if (
saved is not None
and self.config[CONF].get(CONFIG_KEY_REPLAYING)
and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID)
):
saved = await self._aget_checkpoint_before_parent()
# When replaying a subgraph, find the checkpoint that was current
# at the parent's checkpoint time. For stateless subgraphs (no
# checkpointer), this returns None and we start fresh as usual.
if self.config[CONF].get(
CONFIG_KEY_REPLAYING
) and not self.checkpoint_config.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_ID):
saved = await self._aget_checkpoint_after_parent()
if saved is None:
saved = CheckpointTuple(
self.checkpoint_config, empty_checkpoint(), {"step": -2}, None, []
libs/langgraph/tests/test_time_travel.py
+32 -26
View File
@@ -1570,6 +1570,7 @@ def test_stateful_subgraph_retains_state_on_parent_replay(
) -> 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):
@@ -1582,13 +1583,15 @@ def test_stateful_subgraph_retains_state_on_parent_replay(
return {"results": ["p"]}
def step_a(state: SubState) -> SubState:
observed.append(("step_a", dict(state)))
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:
observed.append(("step_b", dict(state)))
started.append(("step_b", dict(state)))
answer = interrupt("question_b")
observed.append(("step_b", dict(state)))
return {"value": [f"b:{answer}"]}
sub = (
@@ -1612,9 +1615,9 @@ def test_stateful_subgraph_retains_state_on_parent_replay(
config = {"configurable": {"thread_id": "1"}}
# === 1st invocation: answer "a1" and "b1" ===
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
# step_a saw empty state (fresh subgraph)
assert observed[0] == ("step_a", {"value": []})
@@ -1623,9 +1626,9 @@ def test_stateful_subgraph_retains_state_on_parent_replay(
# === 2nd invocation: answer "a2" and "b2" ===
observed.clear()
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
# Stateful subgraph retained state from 1st invocation
assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]})
@@ -1636,12 +1639,12 @@ def test_stateful_subgraph_retains_state_on_parent_replay(
# History is newest-first, so first match = 2nd invocation
before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0]
observed.clear()
started.clear()
replay = graph.invoke(None, before_sub_2nd.config)
assert "__interrupt__" in replay
# Replay sees 1st invocation's final state, NOT 2nd invocation's
assert observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]})
assert started[0] == ("step_a", {"value": ["a:a1", "b:b1"]})
def test_stateful_subgraph_retains_state_on_parent_fork(
@@ -1691,14 +1694,14 @@ def test_stateful_subgraph_retains_state_on_parent_fork(
config = {"configurable": {"thread_id": "1"}}
# === 1st invocation: answer "a1" and "b1" ===
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
# === 2nd invocation: answer "a2" and "b2" ===
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
# === Fork from checkpoint before sub_node in 2nd invocation ===
history = list(graph.get_state_history(config))
@@ -1718,6 +1721,7 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay(
) -> 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):
@@ -1730,13 +1734,15 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay(
return {"results": ["p"]}
def step_a(state: SubState) -> SubState:
observed.append(("step_a", dict(state)))
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:
observed.append(("step_b", dict(state)))
started.append(("step_b", dict(state)))
answer = interrupt("question_b")
observed.append(("step_b", dict(state)))
return {"value": [f"b:{answer}"]}
sub = (
@@ -1760,9 +1766,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay(
config = {"configurable": {"thread_id": "1"}}
# === 1st invocation: answer "a1" and "b1" ===
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a1"), config) # hits step_b interrupt
graph.invoke(Command(resume="b1"), config) # completes
# step_a saw empty state, step_b saw only step_a's answer
assert observed[0] == ("step_a", {"value": []})
@@ -1770,9 +1776,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay(
# === 2nd invocation: answer "a2" and "b2" ===
observed.clear()
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
graph.invoke({"results": []}, config) # hits step_a interrupt
graph.invoke(Command(resume="a2"), config) # hits step_b interrupt
graph.invoke(Command(resume="b2"), config) # completes
# Stateless subgraph starts fresh — no memory of 1st invocation
assert observed[0] == ("step_a", {"value": []})
@@ -1782,9 +1788,9 @@ def test_stateless_subgraph_starts_fresh_on_parent_replay(
history = list(graph.get_state_history(config))
before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0]
observed.clear()
started.clear()
replay = graph.invoke(None, before_sub_2nd.config)
assert "__interrupt__" in replay
# Stateless subgraph starts completely fresh on replay
assert observed[0] == ("step_a", {"value": []})
assert started[0] == ("step_a", {"value": []})
libs/langgraph/tests/test_time_travel_async.py
+32 -26
View File
@@ -1603,6 +1603,7 @@ async def test_stateful_subgraph_retains_state_on_parent_replay(
) -> 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):
@@ -1615,13 +1616,15 @@ async def test_stateful_subgraph_retains_state_on_parent_replay(
return {"results": ["p"]}
def step_a(state: SubState) -> SubState:
observed.append(("step_a", dict(state)))
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:
observed.append(("step_b", dict(state)))
started.append(("step_b", dict(state)))
answer = interrupt("question_b")
observed.append(("step_b", dict(state)))
return {"value": [f"b:{answer}"]}
sub = (
@@ -1645,9 +1648,9 @@ async def test_stateful_subgraph_retains_state_on_parent_replay(
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
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": []})
@@ -1656,9 +1659,9 @@ async def test_stateful_subgraph_retains_state_on_parent_replay(
# === 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
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"]})
@@ -1669,12 +1672,12 @@ async def test_stateful_subgraph_retains_state_on_parent_replay(
# History is newest-first, so first match = 2nd invocation
before_sub_2nd = [s for s in history if s.next == ("sub_node",)][0]
observed.clear()
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 observed[0] == ("step_a", {"value": ["a:a1", "b:b1"]})
assert started[0] == ("step_a", {"value": ["a:a1", "b:b1"]})
@pytest.mark.skipif(
@@ -1728,14 +1731,14 @@ async def test_stateful_subgraph_retains_state_on_parent_fork(
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
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
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)]
@@ -1762,6 +1765,7 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay(
) -> 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):
@@ -1774,13 +1778,15 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay(
return {"results": ["p"]}
def step_a(state: SubState) -> SubState:
observed.append(("step_a", dict(state)))
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:
observed.append(("step_b", dict(state)))
started.append(("step_b", dict(state)))
answer = interrupt("question_b")
observed.append(("step_b", dict(state)))
return {"value": [f"b:{answer}"]}
sub = (
@@ -1804,9 +1810,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay(
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
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": []})
@@ -1814,9 +1820,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay(
# === 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
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": []})
@@ -1826,9 +1832,9 @@ async def test_stateless_subgraph_starts_fresh_on_parent_replay(
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]
observed.clear()
started.clear()
replay = await graph.ainvoke(None, before_sub_2nd.config)
assert "__interrupt__" in replay
# Stateless subgraph starts completely fresh on replay
assert observed[0] == ("step_a", {"value": []})
assert started[0] == ("step_a", {"value": []})