Compare commits

..
Author SHA1 Message Date
Elior Nataf LackritzandGitHub e857fede04 Merge branch 'main' into fix/sqlite-delta-parent-walk 2026-08-07 09:51:40 -04:00
Elior Nataf Lackritzandlylelllll ecc420e7fe fix(checkpoint-sqlite): walk delta ancestors by parent pointer
Stage 1 of the sqlite delta history filtered `checkpoint_id <= target` and
streamed `ORDER BY checkpoint_id DESC`. Both predicates encode the same extra
assumption: that every child's checkpoint id sorts above its parent's.

Ancestry is defined by `parent_checkpoint_id`, and nothing in the contract
requires ids to be monotonic. A parent whose id sorted above its child's was
dropped from the stream, so its stored value and its writes were lost with no
error raised. Removing the range filter alone would not help: in DESC order
that parent arrives before the target, so the walk passes it before it has
started. A single-pass ordered stream cannot express this walk.

Replace it with a recursive CTE anchored at the target that follows
`parent_checkpoint_id`. Rows arrive in walk order, so `step_walk_with_row`
keeps its existing shape, and the query now reads only true ancestors instead
of every row at or below the target, which is strictly less IO than before.

Following pointers can loop where a bounded id scan could not, and a loop is
reachable through `put` alone rather than only by corruption: it writes with
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
descendant's config repoints that checkpoint at its own descendant. The walk
therefore stops on a repeated checkpoint id. sqlite yields recursive rows
lazily, so abandoning the cursor ends the recursion instead of waiting on it.

Postgres needs no equivalent change. It pages the whole thread and follows
parent pointers already, so it returns the correct history for this scenario.

Fixes #8550

Co-authored-by: lylelllll <59271327+lylelllll@users.noreply.github.com>
2026-08-06 10:56:58 -04:00
9 changed files with 211 additions and 678 deletions
@@ -538,7 +538,10 @@ class SqliteSaver(BaseCheckpointSaver[str]):
seeded: set[str] = set()
with self.cursor(transaction=False) as cur:
cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id))
cur.execute(
DELTA_STAGE1_SQL,
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
)
for row in cur:
cid, parent_cid, type_tag, blob = row
if step_walk_with_row(
@@ -26,16 +26,38 @@ from typing import Any
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite
# Stage 1 streams ancestors of `target_cid` newest-first. The `<=`
# predicate keeps target itself in the stream so we can read its
# `parent_checkpoint_id` from the first row without a separate lookup;
# the caller skips target's own writes/seed (matches the
# `BaseCheckpointSaver` contract).
# Stage 1 streams target followed by its ancestors, oldest step last, by
# following `parent_checkpoint_id` in a recursive CTE. Target is the anchor
# row, so the caller reads its `parent_checkpoint_id` without a separate
# lookup and skips its own writes/seed (matches the `BaseCheckpointSaver`
# contract).
#
# Ancestry is defined by `parent_checkpoint_id` alone. An earlier form
# filtered `checkpoint_id <= target` and ordered by `checkpoint_id DESC`,
# which additionally required every child's id to sort above its parent's.
# Nothing in the contract promises that, and a parent sorting above its
# child was dropped from the stream entirely, silently costing that
# parent's seed and writes. See #8550.
#
# A parent chain can cycle: `put` writes with `INSERT OR REPLACE`, so
# re-putting an existing checkpoint id under a descendant's config rewrites
# its parent. The old id-range scan read a finite row set and could not
# loop; this one can, so `step_walk_with_row` stops on a repeated
# checkpoint_id. sqlite yields recursive rows lazily, so abandoning the
# cursor ends the recursion rather than waiting on it.
DELTA_STAGE1_SQL = (
"WITH RECURSIVE ancestors(checkpoint_id, parent_checkpoint_id, type, "
"checkpoint) AS ("
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
"FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
"ORDER BY checkpoint_id DESC"
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? "
"UNION ALL "
"SELECT c.checkpoint_id, c.parent_checkpoint_id, c.type, c.checkpoint "
"FROM checkpoints c JOIN ancestors a "
"ON c.checkpoint_id = a.parent_checkpoint_id "
"WHERE c.thread_id = ? AND c.checkpoint_ns = ?"
") "
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint FROM ancestors"
)
@@ -95,14 +117,20 @@ def step_walk_with_row(
Off-path rows (different branch on the same thread) advance the
cursor without doing any work.
Returns True when every requested channel is seeded — the caller
can stop iterating and close the cursor.
Returns True when the caller can stop iterating and close the cursor,
either because every requested channel is seeded or because the parent
chain revisited a checkpoint it had already walked. `put` writes with
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
descendant's config points it at its own descendant and makes the chain
a loop; without this check the recursive stage-1 query would feed rows
forever.
"""
if "started" not in walk_state:
if cid == target_id:
walk_state["started"] = True
walk_state["cur_cid"] = parent_cid
walk_state["active"] = {ch for ch in channels if ch not in seeded}
walk_state["walked"] = {cid}
# Not target yet (or target not present): keep streaming.
return False
active: set[str] = walk_state["active"]
@@ -111,6 +139,10 @@ def step_walk_with_row(
if cid != walk_state["cur_cid"]:
# Off-path row from a sibling branch — skip without deserializing.
return False
walked: set[str] = walk_state["walked"]
if cid in walked:
return True
walked.add(cid)
for ch in active:
chain_by_ch[ch].append(cid)
ckpt = serde.loads_typed((type_tag, blob))
@@ -650,7 +650,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
async with self.lock, self.conn.cursor() as cur:
await cur.execute(
DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)
DELTA_STAGE1_SQL,
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
)
async for row in cur:
cid, parent_cid, type_tag, blob = row
@@ -0,0 +1,129 @@
"""The stage-1 ancestor walk follows `parent_checkpoint_id`, not id order.
Ancestry is defined by the `parent_checkpoint_id` column. An earlier stage-1
query filtered `checkpoint_id <= target` and streamed `ORDER BY checkpoint_id
DESC`, so it also required every child's id to sort above its parent's. Real
ids are `uuid6` and happen to satisfy that, but the contract never promised it,
and a parent sorting above its child was dropped from the stream: its seed and
its writes vanished with no error. See #8550.
Following parent pointers in a recursive CTE removes both assumptions, at the
cost of needing a cycle guard, which the bounded id scan got for free.
"""
from __future__ import annotations
from typing import Any
import pytest
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
DeltaChannelHistory,
empty_checkpoint,
)
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
CHANNEL = "ch"
CONFIG: dict[str, Any] = {"configurable": {"thread_id": "t", "checkpoint_ns": ""}}
EXPECTED: DeltaChannelHistory = {
"writes": [("task", CHANNEL, "write-root")],
"seed": "seed",
}
def _checkpoint(checkpoint_id: str, values: dict[str, Any]) -> Checkpoint:
value = empty_checkpoint()
value["id"] = checkpoint_id
value["channel_values"] = values
return value
# `z` sorts above `a`, so the parent's id sorts above its child's. Real uuid6
# ids never do this; nothing in the contract stops a caller supplying ids that
# do, and clock skew between two processes writing one thread produces it.
PARENT_ID_ORDERS = [
pytest.param("z-older", "a-newer", id="parent_id_sorts_above_child"),
pytest.param("a-older", "z-newer", id="parent_id_sorts_below_child"),
]
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
def test_sync_walk_reaches_parent_whatever_the_id_order(
root_id: str, child_id: str
) -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
root = saver.put(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
child = saver.put(root, _checkpoint(child_id, {}), {}, {})
got = saver.get_delta_channel_history(config=child, channels=[CHANNEL])
assert got[CHANNEL] == EXPECTED
# The unoptimised implementation is the contract; the fast path must
# agree with it on the same rows.
assert (
got[CHANNEL]
== BaseCheckpointSaver.get_delta_channel_history(
saver, config=child, channels=[CHANNEL]
)[CHANNEL]
)
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
async def test_async_walk_reaches_parent_whatever_the_id_order(
root_id: str, child_id: str
) -> None:
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
root = await saver.aput(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
child = await saver.aput(root, _checkpoint(child_id, {}), {}, {})
got = await saver.aget_delta_channel_history(config=child, channels=[CHANNEL])
assert got[CHANNEL] == EXPECTED
def test_long_chain_walks_the_whole_thread() -> None:
"""A migrated thread can hold its only stored value at the root.
The walk then legitimately runs the length of the thread, so nothing in
the cycle check may cut it short. Ids descend as the chain grows here, so
id order fights the walk at every step.
"""
steps = 40
with SqliteSaver.from_conn_string(":memory:") as saver:
parent = saver.put(
CONFIG, _checkpoint(f"id-{steps:03d}", {CHANNEL: "seed"}), {}, {}
)
saver.put_writes(parent, [(CHANNEL, "write-root")], "task")
for step in range(steps - 1, 0, -1):
parent = saver.put(parent, _checkpoint(f"id-{step:03d}", {}), {}, {})
got = saver.get_delta_channel_history(config=parent, channels=[CHANNEL])
assert got[CHANNEL] == EXPECTED
def test_cyclic_parent_chain_terminates() -> None:
"""A parent chain that loops must not feed the walk forever.
Reachable through `put` alone, no corruption needed: it writes with
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
descendant's config repoints that checkpoint at its own descendant. The
old id-range scan read a finite row set and could not loop; a recursive
parent-pointer query can, so the walk stops on a repeated id.
Note this one fails by hanging, not by asserting, since a regression
means the row stream never ends. The package has no timeout plugin, so
the CI job timeout is the backstop.
"""
with SqliteSaver.from_conn_string(":memory:") as saver:
a = saver.put(CONFIG, _checkpoint("cid-a", {}), {}, {})
b = saver.put(a, _checkpoint("cid-b", {}), {}, {})
# Re-put "cid-a" with "cid-b" as its parent: a -> b -> a.
saver.put(b, _checkpoint("cid-a", {}), {}, {})
got = saver.get_delta_channel_history(config=b, channels=[CHANNEL])
# Nothing on the cycle stores a value, so no seed. The point of the
# test is that the call returns at all.
assert "seed" not in got[CHANNEL]
+9 -79
View File
@@ -80,20 +80,12 @@ def get_updated_channels_from_tasks(
def get_delta_channels_from_all_channels(
channels: Mapping[str, BaseChannel],
*,
include_unavailable: bool = False,
) -> set[str]:
"""Every available DeltaChannel.
The set to snapshot whenever no ancestor walk can reconstruct these
channels: the first update_state of a fresh thread (no ancestors at all),
and the first checkpoint of a fork (whose base also holds the writes of the
branch the fork abandons).
"""
"""DeltaChannels to snapshot on the first update_state of a fresh thread."""
return {
k
for k, ch in channels.items()
if isinstance(ch, DeltaChannel) and (include_unavailable or ch.is_available())
if isinstance(ch, DeltaChannel) and ch.is_available()
}
@@ -130,27 +122,15 @@ def create_checkpoint_plan_for_update_state_api(
parents: dict[str, Any],
saved_metadata: Mapping[str, Any] | None,
is_fresh_thread: bool,
is_fork: bool,
) -> tuple[set[str], dict[str, Any]]:
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head.
``is_fork`` (the update was addressed at an explicit checkpoint) forces a
full snapshot for the same reason ``is_fresh_thread`` does: the ancestor
walk cannot reconstruct this head. The base a fork branches off keeps the
pending writes of the branch being abandoned, and nothing records which
child consumed which write, so the walk would replay them here too.
Snapshotting terminates the walk at this checkpoint. Every delta channel
snapshots, so no counters carry over.
"""
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head."""
metadata: dict[str, Any] = {
"source": "update",
"step": step,
"parents": parents,
}
if is_fresh_thread or is_fork:
return get_delta_channels_from_all_channels(
channels, include_unavailable=is_fork
), metadata
if is_fresh_thread:
return get_delta_channels_from_all_channels(channels), metadata
new_counters = create_metadata_for_update_state_api(
channels,
@@ -166,43 +146,6 @@ def create_checkpoint_plan_for_update_state_api(
return channels_to_snapshot, metadata
def create_fork_checkpoint(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
step: int,
*,
is_fork: bool,
get_next_version: GetNextVersion,
) -> Checkpoint:
"""``create_checkpoint`` for an update_state path that bypasses the plan.
The ``as_node`` INPUT and END paths write the fork's first checkpoint and
return before ``create_checkpoint_plan_for_update_state_api`` runs. Left
without a snapshot that checkpoint does not seal the fork, and the next
superstep reconstructs its delta channels by walking through the shared
base, picking up the abandoned branch's writes and then baking them into
whatever it snapshots. Sealing has to happen on the fork's *first*
checkpoint, which is this one.
``get_next_version`` is required for the same reason exit mode needs it:
these paths apply writes to the input channel, not to the delta channel,
so nothing bumps the delta channel's version and ``put`` would drop the
blob as not-a-new-version. Callers must derive ``new_versions`` from the
returned checkpoint rather than the one they passed in.
"""
if not is_fork:
return create_checkpoint(checkpoint, channels, step)
return create_checkpoint(
checkpoint,
channels,
step,
get_next_version=get_next_version,
channels_to_snapshot=get_delta_channels_from_all_channels(
channels, include_unavailable=True
),
)
def create_checkpoint(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel] | None,
@@ -231,27 +174,14 @@ def create_checkpoint(
values = {}
channel_versions = dict(checkpoint["channel_versions"])
for k in channels:
ch = channels[k]
if k not in channel_versions:
# Nothing was ever written to this channel on this branch, so
# it has no version and `put` would drop any blob stored for
# it. A *forced* snapshot still has to land: it is the only
# thing that stops the ancestor walk running past this
# checkpoint into a fork base that holds another branch's
# writes. Mint a first version so the blob survives.
if k in channels_to_snapshot and get_next_version is not None:
channel_versions[k] = get_next_version(None, None)
values[k] = _DeltaSnapshot(
ch.get() if ch.is_available() else ch.typ()
)
continue
ch = channels[k]
if k in channels_to_snapshot:
# Callers force a full snapshot blob here: exit mode when a
# delta channel reaches its snapshot cadence, update_state on
# a fresh thread (no ancestor to replay writes from), and a
# fork (whose base also holds the abandoned branch's writes).
# The manual version-bump below only applies to the exit-mode
# case.
# delta channel reaches its snapshot cadence, and update_state
# on a fresh thread (no ancestor to replay writes from). The
# manual version-bump below only applies to the exit-mode case.
#
# In exit mode, the snapshot decision is deferred to exit
# time (intermediate steps have do_checkpoint=False). The
+11 -60
View File
@@ -222,32 +222,10 @@ class PregelLoop:
# under the saver's `ORDER BY task_id, idx` sorting.
_exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
# Delta channels that must snapshot at the next checkpoint, whatever their
# cadence counters say. Two sources:
# * an Overwrite arrived since the last checkpoint, so the snapshot has to
# happen after live update applied overwrite semantics and sparse replay
# starts from the same post-overwrite value;
# * this run forked off an explicitly addressed checkpoint, see
# `_delta_channels_awaiting_fork_snapshot`.
_delta_channels_forced_snapshot: set[str]
# Delta channels still owed a fork snapshot, when this run was launched
# against an explicitly addressed checkpoint (time travel / fork). That
# base keeps the pending writes of the branch the fork abandons, and
# nothing records which child consumed which write, so the ancestor walk
# would replay them into this branch too. Snapshotting terminates the walk
# inside the fork instead of at the shared base. Names drop out once the
# blob has landed; a channel with no value yet has nothing to snapshot, so
# it waits for the superstep that gives it one.
#
# The trigger is deliberately coarse: any addressed checkpoint, not only
# one that turns out to have abandoned writes on it. Which writes belong to
# which child is exactly what is not recorded, so a narrower test would
# have to trust the base's `pending_writes` to be complete, and a saver
# that leaves them out would silently go back to leaking. Snapshotting when
# it was not needed costs one blob per addressed run; not snapshotting when
# it was needed is silent corruption.
_delta_channels_awaiting_fork_snapshot: set[str]
# Delta channels that saw an Overwrite since the last checkpoint. These
# channels must snapshot after live update applies overwrite semantics so
# sparse replay starts from the same post-overwrite value.
_delta_channels_with_overwrite: set[str]
# The checkpoint_config that points at the parent loaded at `__enter__`
# (or the synthetic-empty checkpoint, on first run). We capture it
@@ -391,16 +369,6 @@ class PregelLoop:
if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
else ()
)
# Checks the value, not just key presence like `is_replaying` above:
# subgraph task configs always carry an explicit `None` here, and only
# a real id means the caller addressed one specific checkpoint. Read
# off `checkpoint_config` so subgraphs resolved through a checkpoint
# map during time travel are covered too, matching `__enter__`.
self._delta_channels_awaiting_fork_snapshot = (
{k for k, spec in specs.items() if isinstance(spec, DeltaChannel)}
if self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
else set()
)
self.prev_checkpoint_config = None
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
self.control = runtime.control if isinstance(runtime, Runtime) else None
@@ -715,7 +683,7 @@ class PregelLoop:
def after_tick(self) -> None:
# finish superstep
writes = [w for t in self.tasks.values() for w in t.writes]
self._delta_channels_forced_snapshot.update(
self._delta_channels_with_overwrite.update(
ch
for ch, v in writes
if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
@@ -1023,7 +991,7 @@ class PregelLoop:
manager=None,
updated_channels=updated_channels,
)
self._delta_channels_forced_snapshot.update(
self._delta_channels_with_overwrite.update(
c
for c, v in input_writes
if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0]
@@ -1165,21 +1133,10 @@ class PregelLoop:
do_checkpoint = self._checkpointer_put_after_previous is not None and (
exiting or self.durability != "exit"
)
# Fork: make this checkpoint self-contained, so the ancestor walk stops
# inside the fork instead of reaching the base this run forked off and
# collecting the abandoned branch's writes from it. Resolved here
# rather than in `_first` so channels that only got a value this
# superstep are covered too.
if self._delta_channels_awaiting_fork_snapshot:
self._delta_channels_forced_snapshot.update(
k
for k in self._delta_channels_awaiting_fork_snapshot
if k in self.channels
)
# create new checkpoint
channels_to_snapshot = (
delta_channels_to_snapshot(self.channels, new_counters)
| self._delta_channels_forced_snapshot
| self._delta_channels_with_overwrite
if do_checkpoint
else set()
)
@@ -1197,13 +1154,7 @@ class PregelLoop:
for k in channels_to_snapshot:
new_counters[k] = (0, 0)
if do_checkpoint:
self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot)
# `create_checkpoint` drops a requested snapshot for a channel with
# no version in this checkpoint yet (nothing was ever written to it
# on this branch), so keep asking until the blob really landed.
self._delta_channels_awaiting_fork_snapshot.difference_update(
self.checkpoint["channel_values"]
)
self._delta_channels_with_overwrite.difference_update(channels_to_snapshot)
non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)}
if non_zero:
self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero
@@ -1288,7 +1239,7 @@ class PregelLoop:
)
channels_to_snapshot = (
delta_channels_to_snapshot(self.channels, counters)
| self._delta_channels_forced_snapshot
| self._delta_channels_with_overwrite
)
pending = [
@@ -1733,7 +1684,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
)
self._delta_write_futs = []
self._error_handler_write_futs = []
self._delta_channels_forced_snapshot = set()
self._delta_channels_with_overwrite = set()
self._exit_delta_writes = (
[] if self.durability == "exit" and self.checkpointer is not None else None
)
@@ -1991,7 +1942,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
)
self._delta_write_futs = []
self._error_handler_write_futs = []
self._delta_channels_forced_snapshot = set()
self._delta_channels_with_overwrite = set()
self._exit_delta_writes = (
[] if self.durability == "exit" and self.checkpointer is not None else None
)
+12 -103
View File
@@ -108,7 +108,6 @@ from langgraph.callbacks import (
get_sync_graph_callback_manager_for_config,
)
from langgraph.channels.base import BaseChannel
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.topic import Topic
from langgraph.config import get_config
from langgraph.constants import END
@@ -134,7 +133,6 @@ from langgraph.pregel._checkpoint import (
copy_checkpoint,
create_checkpoint,
create_checkpoint_plan_for_update_state_api,
create_fork_checkpoint,
empty_checkpoint,
get_updated_channels_from_tasks,
)
@@ -1639,24 +1637,8 @@ class Pregel(
else:
raise ValueError(f"Subgraph {recast} not found")
# Delta channels still owed a fork snapshot. Mirrors
# `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only
# sealed once a checkpoint actually carries the blob. Several superstep
# paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and
# return before reaching the plan, so a flag cleared after the first
# superstep would leave the branch unsealed and the next superstep would
# reconstruct through the shared base.
fork_pending: set[str] = (
{k for k, v in self.channels.items() if isinstance(v, DeltaChannel)}
if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
else set()
)
def perform_superstep(
input_config: RunnableConfig,
updates: Sequence[StateUpdate],
*,
is_fork: bool,
input_config: RunnableConfig, updates: Sequence[StateUpdate]
) -> RunnableConfig:
# get last checkpoint
config = ensure_config(self.config, input_config)
@@ -1744,17 +1726,9 @@ class Pregel(
self.trigger_to_nodes,
)
# save checkpoint
next_checkpoint = create_fork_checkpoint(
checkpoint,
channels,
step,
is_fork=is_fork,
get_next_version=checkpointer.get_next_version,
)
fork_pending.difference_update(next_checkpoint["channel_values"])
next_config = checkpointer.put(
checkpoint_config,
next_checkpoint,
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
@@ -1762,7 +1736,7 @@ class Pregel(
},
get_new_channel_versions(
checkpoint_previous_versions,
next_checkpoint["channel_versions"],
checkpoint["channel_versions"],
),
)
return patch_checkpoint_map(
@@ -1791,17 +1765,9 @@ class Pregel(
if saved and saved.metadata.get("step") is not None
else -1
)
next_checkpoint = create_fork_checkpoint(
checkpoint,
channels,
next_step,
is_fork=is_fork,
get_next_version=checkpointer.get_next_version,
)
fork_pending.difference_update(next_checkpoint["channel_values"])
next_config = checkpointer.put(
checkpoint_config,
next_checkpoint,
create_checkpoint(checkpoint, channels, next_step),
{
"source": "input",
"step": next_step,
@@ -1811,7 +1777,7 @@ class Pregel(
},
get_new_channel_versions(
checkpoint_previous_versions,
next_checkpoint["channel_versions"],
checkpoint["channel_versions"],
),
)
@@ -1907,9 +1873,6 @@ class Pregel(
return perform_superstep(
patch_checkpoint_map(next_config, saved.metadata),
[item for lst in user_group_by.values() for item in lst],
# The checkpoint just written clears tasks and carries
# no delta snapshot, so a fork is still unsealed here.
is_fork=is_fork,
)
return patch_checkpoint_map(next_config, saved.metadata)
@@ -2057,7 +2020,6 @@ class Pregel(
parents=saved.metadata.get("parents", {}) if saved else {},
saved_metadata=saved.metadata if saved else None,
is_fresh_thread=saved is None,
is_fork=is_fork,
)
)
checkpoint = create_checkpoint(
@@ -2070,8 +2032,6 @@ class Pregel(
else None,
channels_to_snapshot=channels_to_snapshot,
)
if is_fork:
fork_pending.difference_update(checkpoint["channel_values"])
next_config = checkpointer.put(
checkpoint_config,
checkpoint,
@@ -2089,14 +2049,8 @@ class Pregel(
current_config = patch_configurable(
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
)
# The flag cannot be derived from `current_config`: `perform_superstep`
# returns the config of the checkpoint it just wrote and the loop feeds
# that back in, so from the second superstep on it always names a
# checkpoint whether or not the caller addressed one.
for superstep in supersteps:
current_config = perform_superstep(
current_config, superstep, is_fork=bool(fork_pending)
)
current_config = perform_superstep(current_config, superstep)
return current_config
async def abulk_update_state(
@@ -2149,24 +2103,8 @@ class Pregel(
else:
raise ValueError(f"Subgraph {recast} not found")
# Delta channels still owed a fork snapshot. Mirrors
# `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only
# sealed once a checkpoint actually carries the blob. Several superstep
# paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and
# return before reaching the plan, so a flag cleared after the first
# superstep would leave the branch unsealed and the next superstep would
# reconstruct through the shared base.
fork_pending: set[str] = (
{k for k, v in self.channels.items() if isinstance(v, DeltaChannel)}
if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
else set()
)
async def aperform_superstep(
input_config: RunnableConfig,
updates: Sequence[StateUpdate],
*,
is_fork: bool,
input_config: RunnableConfig, updates: Sequence[StateUpdate]
) -> RunnableConfig:
# get last checkpoint
config = ensure_config(self.config, input_config)
@@ -2252,25 +2190,16 @@ class Pregel(
self.trigger_to_nodes,
)
# save checkpoint
next_checkpoint = create_fork_checkpoint(
checkpoint,
channels,
step,
is_fork=is_fork,
get_next_version=checkpointer.get_next_version,
)
fork_pending.difference_update(next_checkpoint["channel_values"])
next_config = await checkpointer.aput(
checkpoint_config,
next_checkpoint,
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
},
get_new_channel_versions(
checkpoint_previous_versions,
next_checkpoint["channel_versions"],
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
return patch_checkpoint_map(
@@ -2299,17 +2228,9 @@ class Pregel(
if saved and saved.metadata.get("step") is not None
else -1
)
next_checkpoint = create_fork_checkpoint(
checkpoint,
channels,
next_step,
is_fork=is_fork,
get_next_version=checkpointer.get_next_version,
)
fork_pending.difference_update(next_checkpoint["channel_values"])
next_config = await checkpointer.aput(
checkpoint_config,
next_checkpoint,
create_checkpoint(checkpoint, channels, next_step),
{
"source": "input",
"step": next_step,
@@ -2319,7 +2240,7 @@ class Pregel(
},
get_new_channel_versions(
checkpoint_previous_versions,
next_checkpoint["channel_versions"],
checkpoint["channel_versions"],
),
)
@@ -2414,9 +2335,6 @@ class Pregel(
return await aperform_superstep(
patch_checkpoint_map(next_config, saved.metadata),
[item for lst in user_group_by.values() for item in lst],
# The checkpoint just written clears tasks and carries
# no delta snapshot, so a fork is still unsealed here.
is_fork=is_fork,
)
return patch_checkpoint_map(
@@ -2562,7 +2480,6 @@ class Pregel(
parents=saved.metadata.get("parents", {}) if saved else {},
saved_metadata=saved.metadata if saved else None,
is_fresh_thread=saved is None,
is_fork=is_fork,
)
)
checkpoint = create_checkpoint(
@@ -2575,8 +2492,6 @@ class Pregel(
else None,
channels_to_snapshot=channels_to_snapshot,
)
if is_fork:
fork_pending.difference_update(checkpoint["channel_values"])
next_config = await checkpointer.aput(
checkpoint_config,
checkpoint,
@@ -2593,14 +2508,8 @@ class Pregel(
current_config = patch_configurable(
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
)
# The flag cannot be derived from `current_config`: `aperform_superstep`
# returns the config of the checkpoint it just wrote and the loop feeds
# that back in, so from the second superstep on it always names a
# checkpoint whether or not the caller addressed one.
for superstep in supersteps:
current_config = await aperform_superstep(
current_config, superstep, is_fork=bool(fork_pending)
)
current_config = await aperform_superstep(current_config, superstep)
return current_config
def update_state(
+3 -11
View File
@@ -85,19 +85,11 @@ class MemorySaverAssertImmutable(InMemorySaver):
)
== saved
), config["configurable"]["checkpoint_ns"]
# call super to write checkpoint
next_config = super().put(config, checkpoint, metadata, new_versions)
# Record the checkpoint as the saver reads it back, not the object it
# was handed. `channel_values` are stored per (channel, version), so a
# channel a step did not write is refilled from the blob its inherited
# version still points at. A `DeltaChannel` omits its value except at a
# snapshot, which makes the two representations differ for reasons that
# are not mutation. Comparing read-back against read-back still catches
# a checkpoint whose stored data actually changed.
self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
self.serde.dumps_typed(super().get(next_config))
self.serde.dumps_typed(checkpoint)
)
return next_config
# call super to write checkpoint
return super().put(config, checkpoint, metadata, new_versions)
class MemorySaverNoPending(InMemorySaver):
@@ -1,414 +0,0 @@
"""Forking a thread must not replay the abandoned branch into the fork.
Regression suite for #8443. Addressing an older checkpoint creates a fork: the
shared base ends up with two children, and it keeps the ``checkpoint_writes``
of the branch the fork abandons. Nothing in the stored data records which child
consumed which write, so the ``DeltaChannel`` ancestor walk used to collect the
abandoned branch's writes as well.
Every graph here carries a ``DeltaChannel`` and a plain reducer channel fed the
same values. ``full`` channels store complete ``channel_values`` and need no
replay, so the plain channel is the oracle: after a fork the two must agree.
Coverage: fork by ``invoke`` with new input (sync/async, all durabilities),
fork off the checkpoint that predates the thread's first input (sync/async),
fork by ``update_state`` / ``aupdate_state``, fork before the delta channel
ever had a value, and guards that neither an
unaddressed run nor an unaddressed multi-superstep ``bulk_update_state``
departs from the normal ``snapshot_frequency`` cadence.
"""
from collections.abc import Sequence
from operator import add
from typing import Annotated, Any
import pytest
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph._internal._constants import INPUT
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import END, StateGraph
from langgraph.types import Durability, StateSnapshot, StateUpdate
pytestmark = pytest.mark.anyio
def _append(current: list | None, writes: Sequence[Any]) -> list:
"""DeltaChannel reducer: extend the list with every batched write."""
out = list(current or [])
for write in writes:
out.extend(write if isinstance(write, list) else [write])
return out
class _State(TypedDict):
log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)]
plain: Annotated[list, add]
# Written only by the tests that fork before `log` ever has a value, so the
# fork advances without touching the delta channel.
other: Annotated[list, add]
def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any:
"""Compile a one-node graph whose node appends ``{tag}-out`` to both channels.
``snapshot_frequency=1000`` keeps the cadence from masking the bug: without
a forced snapshot the fork's ancestor walk always runs past the fork base.
"""
def node(state: _State) -> dict:
return {"log": [f"{tag}-out"], "plain": [f"{tag}-out"]}
builder = StateGraph(_State)
builder.add_node("n", node)
builder.set_entry_point("n")
builder.set_finish_point("n")
return builder.compile(checkpointer=checkpointer)
def _build_without_delta_writes(checkpointer: BaseCheckpointSaver, tag: str) -> Any:
"""Compile a graph whose node writes only ``other``, never the delta channel."""
def node(state: _State) -> dict:
return {"other": [f"{tag}-other"]}
builder = StateGraph(_State)
builder.add_node("n", node)
builder.set_entry_point("n")
builder.set_finish_point("n")
return builder.compile(checkpointer=checkpointer)
def _thread(thread_id: str) -> RunnableConfig:
return {"configurable": {"thread_id": thread_id}}
def _at(config: RunnableConfig, snapshot: StateSnapshot) -> RunnableConfig:
"""Config addressing one specific checkpoint of ``config``'s thread."""
return {
"configurable": {
**config["configurable"],
"checkpoint_ns": "",
"checkpoint_id": snapshot.config["configurable"]["checkpoint_id"],
}
}
def _input(marker: str) -> dict:
return {"log": [marker], "plain": [marker]}
def _snapshotted_checkpoints(
checkpointer: BaseCheckpointSaver, config: RunnableConfig
) -> list[str]:
"""Ids of this thread's checkpoints carrying a ``log`` snapshot blob."""
return [
tuple_.config["configurable"]["checkpoint_id"]
for tuple_ in checkpointer.list(config)
if isinstance(tuple_.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
]
def _assert_fork_is_clean(state: StateSnapshot, abandoned: str) -> None:
"""The delta channel must match the plain channel and drop ``abandoned``."""
assert state.values["log"] == state.values["plain"], (
f"delta channel diverged from the plain channel: "
f"{state.values['log']} != {state.values['plain']}"
)
assert abandoned not in state.values["log"], (
f"{abandoned!r} belongs to the branch the fork replaced, "
f"but was replayed into {state.values['log']}"
)
def test_fork_by_invoke(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
config = _thread("t")
_build(sync_checkpointer, "first").invoke(
_input("in-1"), config, durability=durability
)
graph = _build(sync_checkpointer, "second")
graph.invoke(_input("in-2"), config, durability=durability)
# The last checkpoint that predates "in-2" entering state: forking here
# abandons the "in-2" branch, whose writes still hang off this checkpoint.
base = next(
snapshot
for snapshot in graph.get_state_history(config)
if "in-2" not in snapshot.values["log"]
)
_build(sync_checkpointer, "third").invoke(
_input("in-3"), _at(config, base), durability=durability
)
state = graph.get_state(config)
_assert_fork_is_clean(state, "in-2")
assert state.values["log"] == [*base.values["log"], "in-3", "third-out"]
async def test_afork_by_invoke(
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
config = _thread("t")
await _build(async_checkpointer, "first").ainvoke(
_input("in-1"), config, durability=durability
)
graph = _build(async_checkpointer, "second")
await graph.ainvoke(_input("in-2"), config, durability=durability)
base = await anext(
snapshot
async for snapshot in graph.aget_state_history(config)
if "in-2" not in snapshot.values["log"]
)
await _build(async_checkpointer, "third").ainvoke(
_input("in-3"), _at(config, base), durability=durability
)
state = await graph.aget_state(config)
_assert_fork_is_clean(state, "in-2")
assert state.values["log"] == [*base.values["log"], "in-3", "third-out"]
def test_fork_off_checkpoint_before_first_input(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Fork off the root checkpoint, which predates any value for ``log``.
``create_checkpoint`` drops a requested snapshot for a channel absent from
``channel_versions``, so the fork's own first checkpoint cannot carry the
blob. The request has to stay queued until a superstep gives the channel a
value, otherwise the root's ``in-1`` write still leaks into the fork.
"""
config = _thread("t")
graph = _build(sync_checkpointer, "first")
graph.invoke(_input("in-1"), config, durability=durability)
root = list(graph.get_state_history(config))[-1]
assert root.values["log"] == []
_build(sync_checkpointer, "third").invoke(
_input("in-9"), _at(config, root), durability=durability
)
state = graph.get_state(config)
_assert_fork_is_clean(state, "in-1")
assert state.values["log"] == ["in-9", "third-out"]
async def test_afork_off_checkpoint_before_first_input(
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Async twin of ``test_fork_off_checkpoint_before_first_input``."""
config = _thread("t")
graph = _build(async_checkpointer, "first")
await graph.ainvoke(_input("in-1"), config, durability=durability)
root = [snapshot async for snapshot in graph.aget_state_history(config)][-1]
assert root.values["log"] == []
await _build(async_checkpointer, "third").ainvoke(
_input("in-9"), _at(config, root), durability=durability
)
state = await graph.aget_state(config)
_assert_fork_is_clean(state, "in-1")
assert state.values["log"] == ["in-9", "third-out"]
def test_fork_by_update_state(sync_checkpointer: BaseCheckpointSaver) -> None:
config = _thread("t")
_build(sync_checkpointer, "first").invoke(_input("in-1"), config)
graph = _build(sync_checkpointer, "second")
graph.invoke(_input("in-2"), config)
base = next(
snapshot
for snapshot in graph.get_state_history(config)
if "in-2" not in snapshot.values["log"]
)
forked = graph.update_state(_at(config, base), _input("patched"))
state = graph.get_state(forked)
_assert_fork_is_clean(state, "in-2")
assert state.values["log"] == [*base.values["log"], "patched"]
async def test_afork_by_update_state(
async_checkpointer: BaseCheckpointSaver,
) -> None:
config = _thread("t")
await _build(async_checkpointer, "first").ainvoke(_input("in-1"), config)
graph = _build(async_checkpointer, "second")
await graph.ainvoke(_input("in-2"), config)
base = await anext(
snapshot
async for snapshot in graph.aget_state_history(config)
if "in-2" not in snapshot.values["log"]
)
forked = await graph.aupdate_state(_at(config, base), _input("patched"))
state = await graph.aget_state(forked)
_assert_fork_is_clean(state, "in-2")
assert state.values["log"] == [*base.values["log"], "patched"]
def test_unaddressed_run_keeps_snapshot_cadence(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""A run with no explicitly addressed checkpoint writes no snapshot blob.
Guards the cost of the fix: the forced snapshot is one per addressed run,
not a change to the normal ``snapshot_frequency`` cadence.
"""
config = _thread("t")
graph = _build(sync_checkpointer, "first")
graph.invoke(_input("in-1"), config, durability=durability)
graph.invoke(_input("in-2"), config, durability=durability)
assert not _snapshotted_checkpoints(sync_checkpointer, config)
def test_fork_before_first_value_when_fork_never_writes_the_channel(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Seal the fork even when the channel has no value to snapshot.
Forking before ``log`` was ever written leaves nothing to copy into the
fork's first checkpoint, so without a minted version and an empty blob the
boundary goes unrecorded and the walk runs into the base. The fork here
never writes ``log`` at all, so no later superstep can seal it either.
"""
config = _thread("t")
graph = _build(sync_checkpointer, "first")
graph.invoke(_input("in-1"), config, durability=durability)
root = list(graph.get_state_history(config))[-1]
assert root.values["log"] == []
_build_without_delta_writes(sync_checkpointer, "third").invoke(
{"other": ["in-9"]}, _at(config, root), durability=durability
)
state = graph.get_state(config)
_assert_fork_is_clean(state, "in-1")
assert state.values["log"] == []
async def test_afork_before_first_value_when_fork_never_writes_the_channel(
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Async twin of ``test_fork_before_first_value_when_fork_never_writes_the_channel``."""
config = _thread("t")
graph = _build(async_checkpointer, "first")
await graph.ainvoke(_input("in-1"), config, durability=durability)
root = [snapshot async for snapshot in graph.aget_state_history(config)][-1]
assert root.values["log"] == []
await _build_without_delta_writes(async_checkpointer, "third").ainvoke(
{"other": ["in-9"]}, _at(config, root), durability=durability
)
state = await graph.aget_state(config)
_assert_fork_is_clean(state, "in-1")
assert state.values["log"] == []
def test_fork_before_first_value_by_bulk_update(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""The fork's first superstep touches another key, the delta key comes later.
The first checkpoint has to seal the boundary on its own. Deferring until
the superstep that finally writes ``log`` is too late, because that
superstep reconstructs through the unsealed checkpoint first.
"""
config = _thread("t")
graph = _build(sync_checkpointer, "first")
graph.invoke(_input("in-1"), config)
root = list(graph.get_state_history(config))[-1]
assert root.values["log"] == []
forked = graph.bulk_update_state(
_at(config, root),
[
[StateUpdate({"other": ["s1"]}, "n")],
[StateUpdate(_input("s2"), "n")],
],
)
state = graph.get_state(forked)
_assert_fork_is_clean(state, "in-1")
assert state.values["log"] == ["s2"]
@pytest.mark.parametrize("first_as_node", [INPUT, END, "__copy__"])
def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan(
sync_checkpointer: BaseCheckpointSaver, first_as_node: str
) -> None:
"""The fork must be sealed by whichever checkpoint the fork writes first.
``as_node`` of INPUT, END or ``__copy__`` writes a checkpoint and returns
before ``create_checkpoint_plan_for_update_state_api`` runs. If that
checkpoint carries no snapshot the branch is still unsealed, so the next
superstep reconstructs through the shared base, picks up the abandoned
writes, and bakes them into whatever it snapshots. Sealing later is too
late: by then the in-memory value is already wrong.
"""
config = _thread("t")
_build(sync_checkpointer, "first").invoke(_input("in-1"), config)
graph = _build(sync_checkpointer, "second")
graph.invoke(_input("in-2"), config)
base = next(
snapshot
for snapshot in graph.get_state_history(config)
if "in-2" not in snapshot.values["log"]
)
first = (
StateUpdate(_input("first-step"), first_as_node)
if first_as_node == INPUT
else StateUpdate(None, first_as_node)
)
forked = graph.bulk_update_state(
_at(config, base),
[[first], [StateUpdate(_input("second-step"), "n")]],
)
state = graph.get_state(forked)
# END legitimately absorbs the base's already-run task writes, so "in-2"
# belongs there; the plain channel is the oracle for which is which.
assert state.values["log"] == state.values["plain"], (
f"delta channel diverged from the plain channel: "
f"{state.values['log']} != {state.values['plain']}"
)
def test_unaddressed_bulk_update_keeps_snapshot_cadence(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""A multi-superstep ``bulk_update_state`` forks at most once, at the head.
``perform_superstep`` returns the config of the checkpoint it just wrote
and the driver feeds that back in, so every superstep after the first
receives a config naming a checkpoint even when the caller addressed none.
Deriving the fork flag from that config snapshots the whole growing value
once per superstep.
"""
config = _thread("t")
graph = _build(sync_checkpointer, "first")
graph.invoke(_input("in-1"), config)
graph.bulk_update_state(
config,
[[StateUpdate(_input(f"u{i}"), "n")] for i in range(4)],
)
assert not _snapshotted_checkpoints(sync_checkpointer, config)