Compare commits

...
Author SHA1 Message Date
Quanzheng LongandClaude Opus 4.8 8e77945f40 fix: preserve updated_channels shape in update_state delta path
The supersteps-based delta history walk requires update_state to write
counters_since_delta_snapshot metadata, but passing updated_channels to
create_checkpoint also overrode checkpoint["updated_channels"] (was None),
breaking deferred-node triggering on resume for non-delta graphs
(test_in_one_fan_out_state_graph_defer_node). Keep updated_channels at its
default so resume/trigger semantics are unchanged; still write the counter
metadata and snapshot delta channels that hit their cadence.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-01 10:59:45 -07:00
Quanzheng Long 728fb748be fix 2026-06-30 21:22:00 -07:00
13 changed files with 1209 additions and 583 deletions
@@ -54,6 +54,12 @@ async def build_delta_chain(
stored: list[RunnableConfig] = [] stored: list[RunnableConfig] = []
parent_cfg: RunnableConfig | None = None parent_cfg: RunnableConfig | None = None
# Track supersteps since the channel's last snapshot, mirroring Pregel's
# `counters_since_delta_snapshot` bookkeeping: every step increments the
# counter; a snapshot step resets it (and carries no counter entry, so
# the walk locates the seed exactly `supersteps` hops back).
supersteps_since_snapshot = 0
for step in range(total_steps): for step in range(total_steps):
config: RunnableConfig = { config: RunnableConfig = {
"configurable": { "configurable": {
@@ -68,11 +74,19 @@ async def build_delta_chain(
channel_values: dict[str, Any] = {} channel_values: dict[str, Any] = {}
channel_versions: dict[str, int] = {} channel_versions: dict[str, int] = {}
supersteps_since_snapshot += 1
counters: dict[str, tuple[int, int]] = {}
if step in snapshot_set: if step in snapshot_set:
channel_values[channel] = _DeltaSnapshot( channel_values[channel] = _DeltaSnapshot(
write_value_fn(step), write_value_fn(step),
) )
channel_versions[channel] = step + 1 channel_versions[channel] = step + 1
supersteps_since_snapshot = 0
else:
counters[channel] = (
supersteps_since_snapshot,
supersteps_since_snapshot,
)
cp = Checkpoint( cp = Checkpoint(
v=1, v=1,
@@ -84,9 +98,10 @@ async def build_delta_chain(
updated_channels=None, updated_channels=None,
) )
new_versions = dict(channel_versions) new_versions = dict(channel_versions)
parent_cfg = await saver.aput( md = generate_metadata(step=step)
config, cp, generate_metadata(step=step), new_versions if counters:
) md["counters_since_delta_snapshot"] = counters
parent_cfg = await saver.aput(config, cp, md, new_versions)
stored.append(parent_cfg) stored.append(parent_cfg)
# Write a pending write for non-snapshot steps so the walk has # Write a pending write for non-snapshot steps so the walk has
@@ -87,6 +87,10 @@ async def test_history_multi_channel(
from langgraph.checkpoint.conformance.test_utils import generate_metadata from langgraph.checkpoint.conformance.test_utils import generate_metadata
# Per-channel supersteps since last snapshot ("a" snapshots at step 1,
# "b" at step 3). Drives the per-channel walk depth from the head.
s_a = 0
s_b = 0
for step in range(5): for step in range(5):
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}} config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
if parent_cfg: if parent_cfg:
@@ -95,12 +99,21 @@ async def test_history_multi_channel(
] ]
cv: dict = {} cv: dict = {}
cvs: dict = {} cvs: dict = {}
s_a += 1
s_b += 1
counters: dict = {}
if step == 1: if step == 1:
cv["a"] = _DeltaSnapshot("snap_a") cv["a"] = _DeltaSnapshot("snap_a")
cvs["a"] = step + 1 cvs["a"] = step + 1
s_a = 0
else:
counters["a"] = (s_a, s_a)
if step == 3: if step == 3:
cv["b"] = _DeltaSnapshot("snap_b") cv["b"] = _DeltaSnapshot("snap_b")
cvs["b"] = step + 1 cvs["b"] = step + 1
s_b = 0
else:
counters["b"] = (s_b, s_b)
cp = Checkpoint( cp = Checkpoint(
v=1, v=1,
id=str(uuid6(clock_seq=-1)), id=str(uuid6(clock_seq=-1)),
@@ -110,7 +123,10 @@ async def test_history_multi_channel(
versions_seen={}, versions_seen={},
updated_channels=None, updated_channels=None,
) )
parent_cfg = await saver.aput(config, cp, generate_metadata(step=step), cvs) md = generate_metadata(step=step)
if counters:
md["counters_since_delta_snapshot"] = counters
parent_cfg = await saver.aput(config, cp, md, cvs)
configs.append(parent_cfg) configs.append(parent_cfg)
await saver.aput_writes(parent_cfg, [("a", step), ("b", step)], str(uuid4())) await saver.aput_writes(parent_cfg, [("a", step), ("b", step)], str(uuid4()))
@@ -148,7 +164,12 @@ async def test_history_walk_to_root_no_seed(
) )
head = configs[-1] head = configs[-1]
result = await saver.aget_delta_channel_history(config=head, channels=["ch"]) result = await saver.aget_delta_channel_history(config=head, channels=["ch"])
# The head's supersteps counter (4) runs one hop past the real root —
# the snapshot checkpoint was never persisted (implicit empty baseline).
# No seed, and the full chain (steps 0..2) replays from empty.
assert "seed" not in result["ch"], f"Expected no seed, got {result['ch']}" assert "seed" not in result["ch"], f"Expected no seed, got {result['ch']}"
values = [w[2] for w in result["ch"]["writes"]]
assert values == [0, 1, 2], f"Expected full-chain writes [0,1,2], got {values}"
async def test_history_migration_plain_value_as_seed( async def test_history_migration_plain_value_as_seed(
@@ -170,6 +191,10 @@ async def test_history_migration_plain_value_as_seed(
configs: list = [] configs: list = []
parent_cfg = None parent_cfg = None
# The channel was a non-delta channel through step 1 (plain value, no
# delta counter), then migrated to DeltaChannel at step 2. Supersteps
# count from the migration boundary: step 2 -> 1, step 3 -> 2, so the
# head (step 3) walks 2 hops back to the plain-value seed at step 1.
for step in range(4): for step in range(4):
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}} config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
if parent_cfg: if parent_cfg:
@@ -182,6 +207,9 @@ async def test_history_migration_plain_value_as_seed(
if step == 1: if step == 1:
cv["ch"] = [10, 20, 30] cv["ch"] = [10, 20, 30]
cvs["ch"] = step + 1 cvs["ch"] = step + 1
md = generate_metadata(step=step)
if step >= 2:
md["counters_since_delta_snapshot"] = {"ch": (step - 1, step - 1)}
cp = Checkpoint( cp = Checkpoint(
v=1, v=1,
id=str(uuid6(clock_seq=-1)), id=str(uuid6(clock_seq=-1)),
@@ -191,7 +219,7 @@ async def test_history_migration_plain_value_as_seed(
versions_seen={}, versions_seen={},
updated_channels=None, updated_channels=None,
) )
parent_cfg = await saver.aput(config, cp, generate_metadata(step=step), cvs) parent_cfg = await saver.aput(config, cp, md, cvs)
configs.append(parent_cfg) configs.append(parent_cfg)
if step != 1: if step != 1:
await saver.aput_writes(parent_cfg, [("ch", step)], str(uuid4())) await saver.aput_writes(parent_cfg, [("ch", step)], str(uuid4()))
@@ -208,6 +236,67 @@ async def test_history_migration_plain_value_as_seed(
assert values == [2], f"Expected [2], got {values}" assert values == [2], f"Expected [2], got {values}"
async def test_history_migration_skips_seed_checkpoint_writes(
saver: BaseCheckpointSaver,
) -> None:
"""A migrated plain-value seed already incorporates its own checkpoint's
writes, so those writes must NOT be re-emitted in the history.
This is the discriminating case for non-additive reducers (e.g. an
even-only filter): re-applying the seed checkpoint's writes on top of
the seed would corrupt the reconstructed value. The saver contract is
to skip them — only writes strictly after the seed are returned.
"""
from langgraph.checkpoint.base import Checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.conformance.test_utils import generate_metadata
tid = str(uuid4())
configs: list = []
parent_cfg = None
for step in range(4):
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
if parent_cfg:
config["configurable"]["checkpoint_id"] = parent_cfg["configurable"][
"checkpoint_id"
]
cv: dict = {}
cvs: dict = {}
# Step 1: pre-delta plain accumulated value AND its own writes.
if step == 1:
cv["ch"] = [2, 4]
cvs["ch"] = step + 1
md = generate_metadata(step=step)
if step >= 2:
md["counters_since_delta_snapshot"] = {"ch": (step - 1, step - 1)}
cp = Checkpoint(
v=1,
id=str(uuid6(clock_seq=-1)),
ts="",
channel_values=cv,
channel_versions=cvs,
versions_seen={},
updated_channels=None,
)
parent_cfg = await saver.aput(config, cp, md, cvs)
configs.append(parent_cfg)
# The seed checkpoint (step 1) carries writes that the plain value
# already subsumes; later steps carry post-migration delta writes.
write_value = 99 if step == 1 else step * 10
await saver.aput_writes(parent_cfg, [("ch", write_value)], str(uuid4()))
head = configs[-1]
result = await saver.aget_delta_channel_history(config=head, channels=["ch"])
assert result["ch"].get("seed") == [2, 4], (
f"Expected plain seed [2, 4], got {result['ch'].get('seed')}"
)
values = [w[2] for w in result["ch"]["writes"]]
# Seed checkpoint's own write (99) is skipped; step 2's write (20) kept.
assert values == [20], f"Expected only post-seed writes [20], got {values}"
ALL_DELTA_CHANNEL_HISTORY_TESTS = [ ALL_DELTA_CHANNEL_HISTORY_TESTS = [
test_history_returns_writes_oldest_first, test_history_returns_writes_oldest_first,
test_history_seed_is_nearest_snapshot, test_history_seed_is_nearest_snapshot,
@@ -216,6 +305,7 @@ ALL_DELTA_CHANNEL_HISTORY_TESTS = [
test_history_empty_channels_returns_empty, test_history_empty_channels_returns_empty,
test_history_walk_to_root_no_seed, test_history_walk_to_root_no_seed,
test_history_migration_plain_value_as_seed, test_history_migration_plain_value_as_seed,
test_history_migration_skips_seed_checkpoint_writes,
] ]
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
CheckpointMetadata, CheckpointMetadata,
CheckpointTuple, CheckpointTuple,
DeltaChannelHistory, DeltaChannelHistory,
_parse_supersteps_since_last_snapshot_by_channel,
get_checkpoint_id, get_checkpoint_id,
get_serializable_checkpoint_metadata, get_serializable_checkpoint_metadata,
) )
@@ -28,9 +29,11 @@ from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import ( from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE, _DELTA_PAGE_SIZE,
BasePostgresSaver, BasePostgresSaver,
_build_delta_stage1_sql, _advance_shared_chain,
_build_delta_stage2_sql, _build_delta_fetch_sql,
_build_delta_walk_sql,
_DeltaStage2Row, _DeltaStage2Row,
_ingest_walk_page,
) )
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -446,107 +449,137 @@ class PostgresSaver(BasePostgresSaver):
) -> Mapping[str, DeltaChannelHistory]: ) -> Mapping[str, DeltaChannelHistory]:
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`. """Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
Two-stage query, both stages cover ALL requested channels: Reconstructs each delta channel's state for the target checkpoint by
walking the parent chain `supersteps` hops (read from the target's
`counters_since_delta_snapshot` metadata) to its seed snapshot, then
collecting the writes between the seed and the target.
* Stage 1 (paged): dynamic SELECT over `checkpoints` with K parallel Two passes (see the `# Multi-channel two-pass` comment in `base.py`):
JSONB key lookups (one column pair per channel) — no subquery, no
aggregation. Pages newest-first by `checkpoint_id` with a cursor;
page size is `_DELTA_PAGE_SIZE`. Stops paging when every channel
has found its seed or the chain is exhausted.
* Stage 2 (per-channel UNION ALL): one branch per channel reading * WALK (paged): dynamic SELECT over `checkpoints` with K parallel
`checkpoint_writes` filtered to that channel's specific JSONB lookups for `channel_versions[ch]` (the seed blob version),
`chain_cids`, plus one branch per channel that has a seed reading following `parent_checkpoint_id` newest-first. Page size is
`checkpoint_blobs` for that channel + version. Avoids the `_DELTA_PAGE_SIZE`; stops once the shared chain reaches the deepest
over-fetch of a single `channel = ANY(channels)` filter when requested `supersteps`, the root is reached, or the chain is
channels have different chain depths. exhausted.
* FETCH (per-channel UNION ALL): one branch per channel reading
`checkpoint_writes` for its `chain_cids`, plus one branch per
channel that located a seed reading `checkpoint_blobs` at the
seed's version.
""" """
if not channels: if not channels:
return {} return {}
channels = list(channels) channels = list(channels)
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
if not thread_id:
raise ValueError("empty thread ID")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
target = self.get_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
checkpoint_id = target.config["configurable"]["checkpoint_id"]
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in # Resolve the target checkpoint id + its metadata (for supersteps).
# Python after each page. Stops as soon as every channel has its seed. checkpoint_id = get_checkpoint_id(config)
stage1_sql = _build_delta_stage1_sql(channels, paged=True) with self._cursor() as cur:
if checkpoint_id is None:
cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s "
"ORDER BY checkpoint_id DESC LIMIT 1",
(thread_id, checkpoint_ns),
)
else:
cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s "
"AND checkpoint_id = %s",
(thread_id, checkpoint_ns, checkpoint_id),
)
target_row = cur.fetchone()
if target_row is None:
return {ch: {"writes": []} for ch in channels}
target_id = cast(str, target_row["checkpoint_id"])
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
target_row["metadata"] or {}, channels
)
max_supersteps = max(supersteps_by_ch.values(), default=0)
# WALK: page the parent chain, bounded by the deepest supersteps.
walk_sql = _build_delta_walk_sql(channels)
parent_of: dict[str, str | None] = {} parent_of: dict[str, str | None] = {}
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels] ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels] shared_cpid_chain: list[str] = []
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} has_reached_root = False
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
walk_cursor_by_ch: dict[str, str | None] = {}
seeded: set[str] = set()
cursor: str | None = None cursor: str | None = None
with self._cursor() as cur: while max_supersteps > 0:
while True: walk_params: list[Any] = [
stage1_params: list[Any] = [] *channels,
for ch in channels: thread_id,
stage1_params.extend([ch, ch]) checkpoint_ns,
stage1_params.extend( cursor,
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE] cursor,
) _DELTA_PAGE_SIZE,
cur.execute(stage1_sql, stage1_params) ]
with self._cursor() as cur:
cur.execute(walk_sql, walk_params)
page = cur.fetchall() page = cur.fetchall()
if not page: if not page:
break break
oldest = self._ingest_stage1_page( oldest = _ingest_walk_page(
cast("list[Mapping[str, Any]]", page), cast("list[Mapping[str, Any]]", page),
channels, channels,
parent_of, parent_of,
ver_by_i_by_cid, ver_by_i_by_cid,
hs_by_i_by_cid, )
) has_reached_root = _advance_shared_chain(
self._try_advance_walks( target_id, parent_of, shared_cpid_chain, max_supersteps
checkpoint_id, )
channels, if (
parent_of, has_reached_root
ver_by_i_by_cid, or len(shared_cpid_chain) >= max_supersteps
hs_by_i_by_cid, or len(page) < _DELTA_PAGE_SIZE
chain_by_ch, ):
seed_ver_by_ch, break
walk_cursor_by_ch, cursor = oldest
seeded,
)
# Stop if every channel is seeded, or the page was short
# (chain exhausted — no more rows to fetch).
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
break
cursor = oldest
# Stage 2: per-channel UNION ALL — one writes branch per channel chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch = (
# with non-empty chain, plus one blob branch per seeded channel. self._resolve_delta_chains(
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] channels,
supersteps_by_ch,
shared_cpid_chain,
ver_by_i_by_cid,
has_reached_root,
thread_id,
)
)
# FETCH: per-channel UNION ALL — writes branch per chained channel,
# blob branch per seeded channel.
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None] channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
stage2_sql = _build_delta_stage2_sql( fetch_sql = _build_delta_fetch_sql(
channels_with_chain=channels_with_chain, channels_with_chain=channels_with_chain,
channels_with_seed=channels_with_seed, channels_with_seed=channels_with_seed,
) )
if fetch_sql:
if stage2_sql: fetch_params: list[Any] = []
stage2_params: list[Any] = []
for ch in channels_with_chain: for ch in channels_with_chain:
stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]]) fetch_params.extend(
[thread_id, checkpoint_ns, ch, chained_cpid_by_ch[ch]]
)
for ch in channels_with_seed: for ch in channels_with_seed:
stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]]) fetch_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
with self._cursor() as cur: with self._cursor() as cur:
cur.execute(stage2_sql, stage2_params) cur.execute(fetch_sql, fetch_params)
stage2_rows = cur.fetchall() fetch_rows = cur.fetchall()
else: else:
stage2_rows = [] fetch_rows = []
return self._build_delta_channels_writes_history( return self._assemble_delta_history(
channels=channels, channels=channels,
chain_by_ch=chain_by_ch, chained_cpid_by_ch=chained_cpid_by_ch,
seed_cpid_by_ch=seed_cpid_by_ch,
seed_ver_by_ch=seed_ver_by_ch, seed_ver_by_ch=seed_ver_by_ch,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), fetch_rows=cast("list[_DeltaStage2Row]", fetch_rows),
) )
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
CheckpointMetadata, CheckpointMetadata,
CheckpointTuple, CheckpointTuple,
DeltaChannelHistory, DeltaChannelHistory,
_parse_supersteps_since_last_snapshot_by_channel,
get_checkpoint_id, get_checkpoint_id,
get_serializable_checkpoint_metadata, get_serializable_checkpoint_metadata,
) )
@@ -28,9 +29,11 @@ from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import ( from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE, _DELTA_PAGE_SIZE,
BasePostgresSaver, BasePostgresSaver,
_build_delta_stage1_sql, _advance_shared_chain,
_build_delta_stage2_sql, _build_delta_fetch_sql,
_build_delta_walk_sql,
_DeltaStage2Row, _DeltaStage2Row,
_ingest_walk_page,
) )
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -408,89 +411,122 @@ class AsyncPostgresSaver(BasePostgresSaver):
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`. """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
See `PostgresSaver.get_delta_channel_history` for design notes; this is See `PostgresSaver.get_delta_channel_history` for design notes; this is
the async equivalent with internal stage-1 paging and per-channel the async equivalent: paged supersteps-bounded WALK + per-channel
UNION ALL stage-2. UNION ALL FETCH.
""" """
if not channels: if not channels:
return {} return {}
channels = list(channels) channels = list(channels)
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
if not thread_id:
raise ValueError("empty thread ID")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
target = await self.aget_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
checkpoint_id = target.config["configurable"]["checkpoint_id"]
stage1_sql = _build_delta_stage1_sql(channels, paged=True) # Resolve the target checkpoint id + its metadata (for supersteps).
checkpoint_id = get_checkpoint_id(config)
async with self._cursor() as cur:
if checkpoint_id is None:
await cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s "
"ORDER BY checkpoint_id DESC LIMIT 1",
(thread_id, checkpoint_ns),
)
else:
await cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s "
"AND checkpoint_id = %s",
(thread_id, checkpoint_ns, checkpoint_id),
)
target_row = await cur.fetchone()
if target_row is None:
return {ch: {"writes": []} for ch in channels}
target_id = cast(str, target_row["checkpoint_id"])
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
target_row["metadata"] or {}, channels
)
max_supersteps = max(supersteps_by_ch.values(), default=0)
# WALK: page the parent chain, bounded by the deepest supersteps.
walk_sql = _build_delta_walk_sql(channels)
parent_of: dict[str, str | None] = {} parent_of: dict[str, str | None] = {}
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels] ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels] shared_cpid_chain: list[str] = []
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} has_reached_root = False
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
walk_cursor_by_ch: dict[str, str | None] = {}
seeded: set[str] = set()
cursor: str | None = None cursor: str | None = None
async with self._cursor() as cur: while max_supersteps > 0:
while True: walk_params: list[Any] = [
stage1_params: list[Any] = [] *channels,
for ch in channels: thread_id,
stage1_params.extend([ch, ch]) checkpoint_ns,
stage1_params.extend( cursor,
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE] cursor,
) _DELTA_PAGE_SIZE,
await cur.execute(stage1_sql, stage1_params) ]
async with self._cursor() as cur:
await cur.execute(walk_sql, walk_params)
page = await cur.fetchall() page = await cur.fetchall()
if not page: if not page:
break break
oldest = self._ingest_stage1_page( oldest = _ingest_walk_page(
cast("list[Mapping[str, Any]]", page), cast("list[Mapping[str, Any]]", page),
channels, channels,
parent_of, parent_of,
ver_by_i_by_cid, ver_by_i_by_cid,
hs_by_i_by_cid, )
) has_reached_root = _advance_shared_chain(
self._try_advance_walks( target_id, parent_of, shared_cpid_chain, max_supersteps
checkpoint_id, )
channels, if (
parent_of, has_reached_root
ver_by_i_by_cid, or len(shared_cpid_chain) >= max_supersteps
hs_by_i_by_cid, or len(page) < _DELTA_PAGE_SIZE
chain_by_ch, ):
seed_ver_by_ch, break
walk_cursor_by_ch, cursor = oldest
seeded,
)
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
break
cursor = oldest
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch = (
self._resolve_delta_chains(
channels,
supersteps_by_ch,
shared_cpid_chain,
ver_by_i_by_cid,
has_reached_root,
thread_id,
)
)
# FETCH: per-channel UNION ALL — writes branch per chained channel,
# blob branch per seeded channel.
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None] channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
stage2_sql = _build_delta_stage2_sql( fetch_sql = _build_delta_fetch_sql(
channels_with_chain=channels_with_chain, channels_with_chain=channels_with_chain,
channels_with_seed=channels_with_seed, channels_with_seed=channels_with_seed,
) )
if stage2_sql: if fetch_sql:
stage2_params: list[Any] = [] fetch_params: list[Any] = []
for ch in channels_with_chain: for ch in channels_with_chain:
stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]]) fetch_params.extend(
[thread_id, checkpoint_ns, ch, chained_cpid_by_ch[ch]]
)
for ch in channels_with_seed: for ch in channels_with_seed:
stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]]) fetch_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
async with self._cursor() as cur: async with self._cursor() as cur:
await cur.execute(stage2_sql, stage2_params) await cur.execute(fetch_sql, fetch_params)
stage2_rows = await cur.fetchall() fetch_rows = await cur.fetchall()
else: else:
stage2_rows = [] fetch_rows = []
return self._build_delta_channels_writes_history( return self._assemble_delta_history(
channels=channels, channels=channels,
chain_by_ch=chain_by_ch, chained_cpid_by_ch=chained_cpid_by_ch,
seed_cpid_by_ch=seed_cpid_by_ch,
seed_ver_by_ch=seed_ver_by_ch, seed_ver_by_ch=seed_ver_by_ch,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), fetch_rows=cast("list[_DeltaStage2Row]", fetch_rows),
) )
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import logging
import random import random
import warnings import warnings
from collections.abc import Mapping, Sequence from collections.abc import Mapping, Sequence
@@ -15,10 +16,12 @@ from langgraph.checkpoint.base import (
PendingWrite, PendingWrite,
get_checkpoint_id, get_checkpoint_id,
) )
from langgraph.checkpoint.serde.types import TASKS from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
from psycopg.types.json import Jsonb from psycopg.types.json import Jsonb
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal logger = logging.getLogger(__name__)
# Page size for the paged WALK scan in `get_delta_channel_history`. Internal
# constant — exposing this as a kwarg is left as a follow-up. # constant — exposing this as a kwarg is left as a follow-up.
_DELTA_PAGE_SIZE = 1024 _DELTA_PAGE_SIZE = 1024
@@ -160,7 +163,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """
class _DeltaStage2Row(TypedDict, total=False): class _DeltaStage2Row(TypedDict, total=False):
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs).""" """One row from `_build_delta_fetch_sql` (a UNION ALL of writes and blobs)."""
_kind: str # "w" or "b" _kind: str # "w" or "b"
checkpoint_id: str | None # "w" rows only checkpoint_id: str | None # "w" rows only
@@ -172,42 +175,58 @@ class _DeltaStage2Row(TypedDict, total=False):
version: str | None # "b" rows only version: str | None # "b" rows only
# Multi-channel two-stage DeltaChannel reconstruction. # Multi-channel two-pass DeltaChannel reconstruction.
# #
# Stage 1 scans checkpoint metadata (no blob bytes) and emits one row per # A `DeltaChannel` does not store its full value at every checkpoint — it
# checkpoint with K parallel JSONB key lookups (one column pair per # stores periodic full-value *snapshots* and accumulates intermediate
# requested delta channel: ver_i / hs_i). No subqueries, no aggregation. # *writes* between snapshots. To rebuild a channel's value at a target
# Python walks the parent chain once across all channels. # checkpoint we need:
# #
# Stage 2 fetches all writes and the seed blobs for ALL channels in a # - the **seed** — the most recent snapshot at-or-before the target
# single roundtrip via `channel = ANY(%s)` and chain/seed-version # (a single blob row in `checkpoint_blobs`); and
# filtering. # - the **chain writes** — every write for this channel committed
# between that snapshot and the target, in order
# (rows in `checkpoint_writes`).
# #
# Empirical comparison vs an alternative "ship full channel_versions / # Two passes, in order:
# channel_values JSONB and let Python pick" form (1000 checkpoints,
# 8 total channels in graph, 3 delta channels requested):
# #
# Postgres execution: A=0.24ms vs B=0.38ms (both negligible) # 1. WALK — scan checkpoint metadata only (no blob bytes). For each
# End-to-end latency: A=6.83ms vs B=2.28ms (B is 3.0x faster) # requested channel we read `channel_versions[ch]` (the seed's blob
# Wire payload: A=836KB vs B=330KB (61% smaller) # version pointer). Then follow `parent_checkpoint_id` from the
# Buffer hits: identical (167 blocks) # target backwards in pages of `_DELTA_PAGE_SIZE` rows.
# #
# B (this dynamic-columns design) wins because it avoids JSONB # Walk depth is driven by the *supersteps since last snapshot*
# serialization on the wire and JSONB-to-dict deserialization in # counter — `metadata.counters_since_delta_snapshot[ch][1]` — read
# psycopg. Even at K=8 (8 delta channels = 16 dynamic columns), B # from the target checkpoint. A channel's seed snapshot sits exactly
# still beats A end-to-end (4.2ms vs 6.8ms). # `supersteps` hops back along the parent chain; walking by the
# counter (rather than scanning `channel_values` for the snapshot
# marker) is the only reliable way to locate seeds that aren't a
# `_DeltaSnapshot` sentinel — e.g. legacy plain-value blobs left by a
# thread that migrated from a non-delta channel, which `put` stores
# out of the inline `channel_values` map.
#
# 2. FETCH — given each channel's chained checkpoint ids and seed
# version from WALK, pull only the rows we need: writes for those
# exact checkpoint_ids and the seed blob at that exact version. One
# roundtrip, per-channel UNION ALL — no over-fetch.
#
# Walking the *parent chain* (not `list(before=...)`) matters: forked
# threads have multiple branches, and only on-path ancestors contribute.
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: def _build_delta_walk_sql(channels: Sequence[str]) -> str:
"""Build stage 1 SQL with 2K parallel JSONB key lookups. """Build the paged WALK SQL — scans checkpoint metadata only.
For channels=["messages", "files"] (with `paged=True`) the result is:: Emits one row per checkpoint with K parallel JSONB key lookups (one
`ver_i` column per requested channel — the channel's blob version, the
pointer we dereference in FETCH if this checkpoint is the seed). No
blob bytes; the result set fits a paged `LIMIT` cleanly.
For channels=["messages", "files"] the result is::
SELECT checkpoint_id, parent_checkpoint_id, SELECT checkpoint_id, parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver_0, checkpoint -> 'channel_versions' ->> %s AS ver_0,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0, checkpoint -> 'channel_versions' ->> %s AS ver_1
checkpoint -> 'channel_versions' ->> %s AS ver_1,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1
FROM checkpoints FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s WHERE thread_id = %s AND checkpoint_ns = %s
AND (%s::text IS NULL OR checkpoint_id < %s) AND (%s::text IS NULL OR checkpoint_id < %s)
@@ -215,45 +234,38 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
LIMIT %s LIMIT %s
Channel names are passed as `%s` parameters (safe from SQL injection). Channel names are passed as `%s` parameters (safe from SQL injection).
Only the column aliases `ver_i` / `hs_i` are interpolated into the Only the column aliases `ver_i` are interpolated into the SQL string
SQL string (i is bounded by len(channels) and uses safe identifiers). (i is bounded by len(channels) and uses safe identifiers).
Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ..., Caller must extend params with `[ch_0, ch_1, ..., thread_id, ns,
thread_id, ns, cursor, cursor, page_size]` when `paged=True`. cursor, cursor, page_size]`. The `cursor` is the smallest
`checkpoint_id` from the previous page (or `None` on the first page);
When `paged=False`, the WHERE has no cursor predicate and there's no `(%s::text IS NULL OR ...)` makes the first-page `WHERE` a no-op.
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
""" """
cols = [] cols = [
for i in range(len(channels)): f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}"
cols.append( for i in range(len(channels))
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, " ]
f"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}" return (
)
sql = (
"SELECT checkpoint_id, parent_checkpoint_id, " "SELECT checkpoint_id, parent_checkpoint_id, "
+ ", ".join(cols) + ", ".join(cols)
+ " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s" + " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
" AND (%s::text IS NULL OR checkpoint_id < %s)"
" ORDER BY checkpoint_id DESC LIMIT %s"
) )
if paged:
sql += (
" AND (%s::text IS NULL OR checkpoint_id < %s)"
" ORDER BY checkpoint_id DESC LIMIT %s"
)
return sql
def _build_delta_stage2_sql( def _build_delta_fetch_sql(
*, *,
channels_with_chain: Sequence[str], channels_with_chain: Sequence[str],
channels_with_seed: Sequence[str], channels_with_seed: Sequence[str],
) -> str: ) -> str:
"""Build stage 2 SQL as a per-channel UNION ALL. """Build the FETCH SQL as a per-channel UNION ALL.
For each channel with a non-empty chain, emit one branch reading For each channel with a non-empty chain, emit one branch reading
`checkpoint_writes` for that specific channel + chain_cids. For each `checkpoint_writes` for that specific channel + chain_cids. For each
channel with a seed_version, emit one branch reading `checkpoint_blobs` channel with a seed_version, emit one branch reading `checkpoint_blobs`
for that channel + version. This avoids the over-fetch of the prior for that channel + version. This avoids the over-fetch of a single
`channel = ANY(channels) AND checkpoint_id = ANY(union)` form when `channel = ANY(channels) AND checkpoint_id = ANY(union)` form when
channels have different chain depths. channels have different chain depths.
@@ -269,6 +281,7 @@ def _build_delta_stage2_sql(
""" """
branches: list[str] = [] branches: list[str] = []
for _ in channels_with_chain: for _ in channels_with_chain:
# NOTE: no ORDER BY on this branch — writes are sorted in assembly.
branches.append( branches.append(
"SELECT 'w'::text AS _kind, " "SELECT 'w'::text AS _kind, "
"checkpoint_id, channel, " "checkpoint_id, channel, "
@@ -288,10 +301,59 @@ def _build_delta_stage2_sql(
return " UNION ALL ".join(branches) return " UNION ALL ".join(branches)
# Stage 1 rows are dynamic-shape dicts: {checkpoint_id, parent_checkpoint_id, def _ingest_walk_page(
# ver_0, hs_0, ver_1, hs_1, ...}. Walking is parameterized by the channel page_rows: Sequence[Mapping[str, Any]],
# list to map indices back to channel names — no static TypedDict here. channels: Sequence[str],
# `dict[str, Any]` is the practical signature. parent_of: dict[str, str | None],
ver_by_i_by_cid: list[dict[str, str | None]],
) -> str | None:
"""Fold one WALK page into `parent_of` + per-channel `ver_by_cid`.
Returns the oldest checkpoint_id seen on this page (smallest, since
pages come back DESC). Caller uses it as the cursor for the next page
(`AND checkpoint_id < cursor`).
"""
oldest: str | None = None
for r in page_rows:
cid = cast(str, r["checkpoint_id"])
parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
for i in range(len(channels)):
ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
# Rows are DESC; the last one is the smallest cid in the page.
oldest = cid
return oldest
def _advance_shared_chain(
target_id: str,
parent_of: Mapping[str, str | None],
shared_cpid_chain: list[str],
max_supersteps: int,
) -> bool:
"""Extend the shared parent chain as far as ingested pages allow.
The chain holds ancestors of the target, newest first: `chain[0]` is
the target's parent, `chain[1]` its grandparent, and so on. A single
chain is shared across all channels and grown to the maximum requested
depth; each channel later slices `chain[:supersteps]`.
Stops when:
- the chain reaches `max_supersteps` hops, OR
- the root is reached (parent is None) — returns True, OR
- the next ancestor's cid isn't in `parent_of` yet (waits for the
next page).
Returns True iff the root was reached.
"""
while len(shared_cpid_chain) < max_supersteps:
top_of_chain = shared_cpid_chain[-1] if shared_cpid_chain else target_id
if top_of_chain not in parent_of:
return False # wait for the next page
parent = parent_of[top_of_chain]
if parent is None:
return True # hit the root
shared_cpid_chain.append(parent)
return False
class BasePostgresSaver(BaseCheckpointSaver[str]): class BasePostgresSaver(BaseCheckpointSaver[str]):
@@ -336,107 +398,92 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
if t.decode() != "empty" if t.decode() != "empty"
} }
@staticmethod def _resolve_delta_chains(
def _ingest_stage1_page( self,
stage1_rows: Sequence[Mapping[str, Any]],
channels: Sequence[str], channels: Sequence[str],
parent_of: dict[str, str | None], supersteps_by_ch: Mapping[str, int],
ver_by_i_by_cid: list[dict[str, str | None]], shared_cpid_chain: Sequence[str],
hs_by_i_by_cid: list[dict[str, bool]],
) -> str | None:
"""Fold one stage-1 page into the running walk-state mappings.
Returns the oldest checkpoint_id seen on this page (smallest, since
pages come back DESC). Caller uses it as the cursor for the next
page (`AND checkpoint_id < cursor`).
"""
oldest: str | None = None
for r in stage1_rows:
cid = cast(str, r["checkpoint_id"])
parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
for i in range(len(channels)):
ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
hs_by_i_by_cid[i][cid] = bool(r.get(f"hs_{i}"))
# Rows are DESC; the last one is the smallest cid in the page.
oldest = cid
return oldest
@staticmethod
def _try_advance_walks(
target_id: str,
channels: Sequence[str],
parent_of: Mapping[str, str | None],
ver_by_i_by_cid: Sequence[Mapping[str, str | None]], ver_by_i_by_cid: Sequence[Mapping[str, str | None]],
hs_by_i_by_cid: Sequence[Mapping[str, bool]], has_reached_root: bool,
chain_by_ch: dict[str, list[str]], thread_id: str,
seed_ver_by_ch: dict[str, str | None], ) -> tuple[
walk_cursor_by_ch: dict[str, str | None], dict[str, list[str]],
seeded: set[str], dict[str, str | None],
) -> None: dict[str, str | None],
"""Advance each not-yet-seeded channel's walk as far as possible. ]:
"""Slice the shared parent chain into per-channel chain/seed mappings.
Uses the partial `parent_of` map accumulated so far. A walk stops For each channel the seed snapshot sits `supersteps` hops back, so
either because: the seed checkpoint is `shared_cpid_chain[supersteps - 1]` and the
(a) it found a snapshot for its channel (channel becomes seeded), chain is `shared_cpid_chain[:supersteps]` (newest first).
(b) it reached a real root (parent_of[cid] is None — fully
materialized at this point), or
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
a later page; the cursor stays put).
Mutates `chain_by_ch`, `seed_ver_by_ch`, `walk_cursor_by_ch`, and When the chain is shorter than `supersteps` but the walk reached the
`seeded` in place. root, the persisted chain is "compressed" relative to the logical
superstep count — either because intermediate supersteps were never
persisted (`durability="exit"`) or because the thread never produced
a snapshot at all. In both cases the seed candidate is the oldest
persisted checkpoint (`shared_cpid_chain[-1]`): FETCH loads its blob,
and assembly keeps it only if non-empty (a real snapshot or migrated
value) — otherwise it omits `seed` and replays the full chain on an
empty baseline.
""" """
chained_cpid_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
seed_cpid_by_ch: dict[str, str | None] = {ch: None for ch in channels}
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
for i, ch in enumerate(channels): for i, ch in enumerate(channels):
if ch in seeded: bound = supersteps_by_ch.get(ch, 0)
if bound <= 0:
continue continue
# First-time entry: cursor starts at the target's parent. if len(shared_cpid_chain) >= bound:
if ch not in walk_cursor_by_ch: seed_depth = bound
walk_cursor_by_ch[ch] = parent_of.get(target_id) elif has_reached_root:
cur_cid = walk_cursor_by_ch[ch] seed_depth = len(shared_cpid_chain)
ch_chain = chain_by_ch[ch] else:
hs_i = hs_by_i_by_cid[i] logger.warning(
ver_i = ver_by_i_by_cid[i] "cannot find seed snapshot for delta channel "
while cur_cid is not None: "(thread_id=%s, channel=%s)",
if cur_cid not in parent_of: thread_id,
# Need more pages to continue this walk. ch,
break )
ch_chain.append(cur_cid) continue
if hs_i.get(cur_cid, False): if seed_depth <= 0:
seed_ver_by_ch[ch] = ver_i.get(cur_cid) continue
seeded.add(ch) chained_cpid_by_ch[ch] = list(shared_cpid_chain[:seed_depth])
cur_cid = None seed_cpid_by_ch[ch] = shared_cpid_chain[seed_depth - 1]
break seed_ver_by_ch[ch] = ver_by_i_by_cid[i].get(seed_cpid_by_ch[ch])
cur_cid = parent_of[cur_cid] return chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch
walk_cursor_by_ch[ch] = cur_cid
def _build_delta_channels_writes_history( def _assemble_delta_history(
self, self,
*, *,
channels: Sequence[str], channels: Sequence[str],
chain_by_ch: Mapping[str, list[str]], chained_cpid_by_ch: Mapping[str, Sequence[str]],
seed_cpid_by_ch: Mapping[str, str | None],
seed_ver_by_ch: Mapping[str, str | None], seed_ver_by_ch: Mapping[str, str | None],
stage2_rows: Sequence[_DeltaStage2Row], fetch_rows: Sequence[_DeltaStage2Row],
) -> dict[str, DeltaChannelHistory]: ) -> dict[str, DeltaChannelHistory]:
"""Demux stage 2 rows per channel; produce per-channel histories. """Demux FETCH rows per channel and produce per-channel histories.
stage2_rows carry `channel` on every row. We build per-channel `fetch_rows` carry `channel` on every row. Write rows (`_kind = 'w'`)
`writes_by_cid` and per-channel `seed_blob` dicts, then assemble are bucketed per channel per checkpoint; seed-blob rows
a `DeltaChannelHistory` per requested channel. The `seed` key is omitted (`_kind = 'b'`) give each channel its snapshot value.
when the walk reached root with no snapshot found, or when the
seed blob is sentinel "empty" — in both cases the consumer treats The seed checkpoint's own writes are replayed on top of a
absence as "start empty". `_DeltaSnapshot` seed (the snapshot is the value *prior* to its own
writes), but skipped for a migrated plain-value seed (a legacy
non-delta blob already incorporates those writes). The `seed` key is
omitted when no seed was located or the blob is the "empty"
tombstone — the consumer treats absence as "start empty".
""" """
# writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx) # writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx)
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = { writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
ch: {} for ch in channels ch: {} for ch in channels
} }
# seed_blob_by_ver[(channel, version)] = (type, blob) seed_blob_by_ch: dict[str, tuple[str, bytes]] = {}
seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {}
for r in stage2_rows: for r in fetch_rows:
ch = cast(str, r["channel"]) ch = cast(str, r["channel"])
kind = r["_kind"] if r["_kind"] == "w":
if kind == "w":
cid = cast(str, r["checkpoint_id"]) cid = cast(str, r["checkpoint_id"])
writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append( writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
cast( cast(
@@ -444,35 +491,39 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
(r["type"], r["blob"], r["task_id"], r["idx"]), (r["type"], r["blob"], r["task_id"], r["idx"]),
) )
) )
else: # kind == "b" else: # _kind == "b" — the seed blob for this channel.
ver = cast(str, r["version"]) seed_blob_by_ch[ch] = cast("tuple[str, bytes]", (r["type"], r["blob"]))
seed_blob_by_ver[(ch, ver)] = cast(
"tuple[str, bytes]", (r["type"], r["blob"])
)
# Sort writes per (channel, cid) newest-first by (task_id, idx) # Within a checkpoint, writes apply oldest→newest by (task_id, idx).
for cid_map in writes_by_ch_by_cid.values(): for cid_map in writes_by_ch_by_cid.values():
for ws in cid_map.values(): for ws in cid_map.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True) ws.sort(key=lambda w: (w[2], w[3]))
result: dict[str, DeltaChannelHistory] = {} result: dict[str, DeltaChannelHistory] = {}
for ch in channels: for ch in channels:
chain_cids = chain_by_ch.get(ch, []) entry: DeltaChannelHistory = {"writes": []}
seed_version = seed_ver_by_ch.get(ch)
skip_seed_checkpoint_writes = False
seed_blob = seed_blob_by_ch.get(ch)
if seed_blob is not None and seed_blob[0] != "empty":
seed_value = self.serde.loads_typed(seed_blob)
entry["seed"] = seed_value
# A migrated (non-delta) seed already includes the writes on
# its own checkpoint; a `_DeltaSnapshot` does not.
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
collected: list[PendingWrite] = []
cid_writes = writes_by_ch_by_cid.get(ch, {}) cid_writes = writes_by_ch_by_cid.get(ch, {})
for cid in chain_cids: if cid_writes:
for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []): collected: list[PendingWrite] = []
val = self.serde.loads_typed((type_tag, write_blob)) seed_cpid = seed_cpid_by_ch.get(ch)
collected.append((task_id, ch, val)) # Chain is newest→oldest; replay oldest→newest.
collected.reverse() for cid in reversed(chained_cpid_by_ch.get(ch, [])):
if skip_seed_checkpoint_writes and cid == seed_cpid:
entry: DeltaChannelHistory = {"writes": collected} continue
if seed_version is not None: for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []):
blob = seed_blob_by_ver.get((ch, seed_version)) val = self.serde.loads_typed((type_tag, write_blob))
if blob is not None and blob[0] != "empty": collected.append((task_id, ch, val))
entry["seed"] = self.serde.loads_typed(blob) entry["writes"] = collected
result[ch] = entry result[ch] = entry
return result return result
@@ -24,10 +24,12 @@ from langgraph.checkpoint.base import (
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.sqlite._delta import ( from langgraph.checkpoint.sqlite._delta import (
DELTA_STAGE1_SQL, DELTA_WALK_SQL,
build_delta_channels_writes_history, build_delta_channels_writes_history,
build_delta_stage2_sql, build_delta_writes_fetch_sql,
step_walk_with_row, parse_supersteps_since_last_snapshot_by_channel,
resolve_delta_chains,
step_walk_supersteps,
) )
from langgraph.checkpoint.sqlite.utils import search_where from langgraph.checkpoint.sqlite.utils import search_where
@@ -505,68 +507,106 @@ class SqliteSaver(BaseCheckpointSaver[str]):
) -> Mapping[str, DeltaChannelHistory]: ) -> Mapping[str, DeltaChannelHistory]:
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`. """Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
Two-stage query: Walks the parent chain `supersteps` hops (read from the target's
`counters_since_delta_snapshot` metadata) to each channel's seed
snapshot, then collects the writes between the seed and the target.
* Stage 1 (paged): newest-first slice of `checkpoints` returning * WALK: stream a newest-first slice of `checkpoints` returning
`(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per `(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per
ancestor. Sqlite has no JSONB, so we ship the full serialized ancestor. Sqlite has no JSONB, so we deserialize only the seed
checkpoint blob and inspect `channel_values` in Python. Pages checkpoints to read their inline `channel_values`. Stops once the
newest-first by `checkpoint_id` with a `< cursor` predicate; shared chain reaches the deepest requested `supersteps` or the
page size is `DELTA_PAGE_SIZE`. Stops paging when every channel root is reached.
has found its seed or the chain is exhausted.
* Stage 2 (per-channel UNION ALL): one branch per channel reading * FETCH (per-channel UNION ALL): one branch per channel reading
`writes` filtered to that channel's specific `chain_cids`. No `writes` filtered to that channel's `chain_cids`. No separate
separate seed-blob fetch — sqlite stores `channel_values` inline seed-blob fetch — sqlite stores `channel_values` inline.
in the checkpoint blob, so seeds come back from stage 1.
""" """
if not channels: if not channels:
return {} return {}
channels = list(channels) channels = list(channels)
thread_id = str(config["configurable"]["thread_id"]) thread_id = str(config["configurable"]["thread_id"])
if not thread_id:
raise ValueError("empty thread ID")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
target = self.get_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
checkpoint_id = target.config["configurable"]["checkpoint_id"]
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} # Resolve the target checkpoint id + its metadata (for supersteps).
seed_val_by_ch: dict[str, Any] = {} checkpoint_id = get_checkpoint_id(config)
with self.cursor(transaction=False) as cur:
if checkpoint_id is None:
cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? "
"ORDER BY checkpoint_id DESC LIMIT 1",
(thread_id, checkpoint_ns),
)
else:
cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? "
"AND checkpoint_id = ?",
(thread_id, checkpoint_ns, checkpoint_id),
)
target_row = cur.fetchone()
if target_row is None:
return {ch: {"writes": []} for ch in channels}
target_id = str(target_row[0])
metadata = json.loads(target_row[1]) if target_row[1] is not None else {}
supersteps_by_ch = parse_supersteps_since_last_snapshot_by_channel(
metadata, channels
)
max_supersteps = max(supersteps_by_ch.values(), default=0)
needed_depths = set(supersteps_by_ch.values())
shared_cpid_chain: list[str] = []
walk_state: dict[str, Any] = {} walk_state: dict[str, Any] = {}
seeded: set[str] = set() seed_values_by_depth: dict[int, dict[str, Any]] = {}
with self.cursor(transaction=False) as cur: with self.cursor(transaction=False) as cur:
cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)) if max_supersteps > 0:
for row in cur: cur.execute(DELTA_WALK_SQL, (thread_id, checkpoint_ns, target_id))
cid, parent_cid, type_tag, blob = row for row in cur:
if step_walk_with_row( cid, parent_cid, type_tag, blob = row
cid=cid, if step_walk_supersteps(
parent_cid=parent_cid, cid=cid,
type_tag=type_tag, parent_cid=parent_cid,
blob=blob, type_tag=type_tag,
target_id=checkpoint_id, blob=blob,
serde=self.serde, target_id=target_id,
chain_by_ch=chain_by_ch, serde=self.serde,
seed_val_by_ch=seed_val_by_ch, shared_cpid_chain=shared_cpid_chain,
walk_state=walk_state, walk_state=walk_state,
seeded=seeded, max_supersteps=max_supersteps,
channels=channels, needed_depths=needed_depths,
): seed_values_by_depth=seed_values_by_depth,
break channels=channels,
):
break
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] (
stage2_sql = build_delta_stage2_sql( chained_cpid_by_ch,
chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain], seed_cpid_by_ch,
seed_value_by_ch,
) = resolve_delta_chains(
channels=channels,
supersteps_by_ch=supersteps_by_ch,
shared_cpid_chain=shared_cpid_chain,
seed_values_by_depth=seed_values_by_depth,
has_reached_root=bool(walk_state.get("reached_root")),
thread_id=thread_id,
) )
if stage2_sql:
stage2_params: list[Any] = [] channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
fetch_sql = build_delta_writes_fetch_sql(
chain_lens=[len(chained_cpid_by_ch[ch]) for ch in channels_with_chain],
)
if fetch_sql:
fetch_params: list[Any] = []
for ch in channels_with_chain: for ch in channels_with_chain:
stage2_params.extend( fetch_params.extend(
[thread_id, checkpoint_ns, ch, *chain_by_ch[ch]] [thread_id, checkpoint_ns, ch, *chained_cpid_by_ch[ch]]
) )
cur.execute(stage2_sql, stage2_params) cur.execute(fetch_sql, fetch_params)
stage2_rows = cast( stage2_rows = cast(
"list[tuple[str, str, str, int, str, bytes]]", cur.fetchall() "list[tuple[str, str, str, int, str, bytes]]", cur.fetchall()
) )
@@ -575,9 +615,9 @@ class SqliteSaver(BaseCheckpointSaver[str]):
return build_delta_channels_writes_history( return build_delta_channels_writes_history(
channels=channels, channels=channels,
chain_by_ch=chain_by_ch, chained_cpid_by_ch=chained_cpid_by_ch,
seed_val_by_ch=seed_val_by_ch, seed_cpid_by_ch=seed_cpid_by_ch,
seeded=seeded, seed_value_by_ch=seed_value_by_ch,
stage2_rows=stage2_rows, stage2_rows=stage2_rows,
serde=self.serde, serde=self.serde,
) )
@@ -1,37 +1,50 @@
"""Shared helpers for `get_delta_channel_history` on sqlite savers. """Shared helpers for `get_delta_channel_history` on sqlite savers.
Mirrors the two-stage shape of `BasePostgresSaver` (ancestor walk + Mirrors the supersteps-based two-pass shape of `BasePostgresSaver`
per-channel UNION ALL writes fetch), but adapted for sqlite's (ancestor walk bounded by `counters_since_delta_snapshot` + per-channel
constraints. The structural differences: UNION ALL writes fetch), adapted for sqlite's constraints:
* No JSONB — to inspect `channel_values` for a checkpoint we must * No JSONB — to inspect `channel_values` for a checkpoint we must
deserialize the full blob. Stage 1 streams the cursor row-by-row and deserialize the full blob. The WALK streams the cursor row-by-row and
deserializes only the rows the merged walk visits, freeing each blob deserializes only the seed checkpoints (the ones at a channel's
before advancing. `supersteps` depth), freeing each blob before advancing.
* No separate blob table — `channel_values` lives inline in the * No separate blob table — `channel_values` lives inline in the
checkpoint, so seeds come back from stage 1 with no second fetch. checkpoint, so seeds come back from the WALK with no second fetch.
* Single merged walk (not K independent walks): each visited cid is * Single shared parent-chain walk: each requested channel slices the
deserialized exactly once, regardless of how many channels are still same chain to its own `supersteps` depth.
seeking their seed.
The streaming design keeps peak in-flight memory at roughly one Walk depth is driven by the *supersteps since last snapshot* counter,
deserialized checkpoint at a time, instead of holding the entire not by scanning `channel_values` for the snapshot marker — the only
ancestor chain's worth of raw blobs as a `fetchall()`-materialized list. reliable way to locate seeds that aren't a `_DeltaSnapshot` sentinel
(e.g. legacy plain-value blobs from a thread migrated off a non-delta
channel).
""" """
from __future__ import annotations from __future__ import annotations
import logging
from collections.abc import Mapping, Sequence from collections.abc import Mapping, Sequence
from typing import Any from typing import Any
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite from langgraph.checkpoint.base import (
DeltaChannelHistory,
PendingWrite,
_parse_supersteps_since_last_snapshot_by_channel,
)
from langgraph.checkpoint.serde.types import _DeltaSnapshot
# Stage 1 streams ancestors of `target_cid` newest-first. The `<=` logger = logging.getLogger(__name__)
# Re-exported under the package-local name used by the sqlite savers.
parse_supersteps_since_last_snapshot_by_channel = (
_parse_supersteps_since_last_snapshot_by_channel
)
# The WALK streams ancestors of `target_id` newest-first. The `<=`
# predicate keeps target itself in the stream so we can read its # predicate keeps target itself in the stream so we can read its
# `parent_checkpoint_id` from the first row without a separate lookup; # `parent_checkpoint_id` from the first matching row without a separate
# the caller skips target's own writes/seed (matches the # lookup; target's own writes/seed are not part of the contract.
# `BaseCheckpointSaver` contract). DELTA_WALK_SQL = (
DELTA_STAGE1_SQL = (
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint " "SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
"FROM checkpoints " "FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? " "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
@@ -39,12 +52,12 @@ DELTA_STAGE1_SQL = (
) )
def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str: def build_delta_writes_fetch_sql(*, chain_lens: Sequence[int]) -> str:
"""Stage-2 per-channel UNION ALL fetching writes from `writes`. """Per-channel UNION ALL fetching writes from `writes`.
One branch per channel with a non-empty chain. Each branch inlines its One branch per channel with a non-empty chain. Each branch inlines its
own `IN (?, ?, ...)` placeholder list because sqlite has no array-bind own `IN (?, ?, ...)` placeholder list because sqlite has no array-bind
equivalent of postgres's `= ANY(%s)`. Caller passes parameters in equivalent of postgres's `= ANY(?)`. Caller passes parameters in
matching order: `[thread_id, checkpoint_ns, channel, *chain_cids]` per matching order: `[thread_id, checkpoint_ns, channel, *chain_cids]` per
branch. branch.
@@ -65,7 +78,7 @@ def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str:
return " UNION ALL ".join(branches) return " UNION ALL ".join(branches)
def step_walk_with_row( def step_walk_supersteps(
*, *,
cid: str, cid: str,
parent_cid: str | None, parent_cid: str | None,
@@ -73,75 +86,142 @@ def step_walk_with_row(
blob: bytes, blob: bytes,
target_id: str, target_id: str,
serde: Any, serde: Any,
chain_by_ch: dict[str, list[str]], shared_cpid_chain: list[str],
seed_val_by_ch: dict[str, Any],
walk_state: dict[str, Any], walk_state: dict[str, Any],
seeded: set[str], max_supersteps: int,
needed_depths: set[int],
seed_values_by_depth: dict[int, dict[str, Any]],
channels: Sequence[str], channels: Sequence[str],
) -> bool: ) -> bool:
"""Process one streamed stage-1 row in the merged ancestor walk. """Process one streamed WALK row, extending the shared parent chain.
The cursor returns (cid, parent_cid, type, blob) rows in The cursor returns `(cid, parent_cid, type, checkpoint)` rows in
`checkpoint_id` DESC order starting at target. The first row is `checkpoint_id` DESC order starting at target. The first row is target
target itself; we read its parent_cid to seed the walk and otherwise itself; we read its `parent_cid` to seed the walk and skip it (target's
skip it (target's own writes/seed are not part of the contract). own writes/seed are not part of the contract). Off-path rows (a sibling
branch on the same thread) advance the cursor without doing work.
For each subsequent row, if `cid` matches the walk's current For each on-path ancestor we append its cid to `shared_cpid_chain`
position, we deserialize the blob, append the cid to every (newest first). When the ancestor sits at a depth some channel needs as
not-yet-seeded channel's chain, and check `channel_values` for its seed (`len(chain)` ∈ `needed_depths`), we deserialize it once and
seeds. The deserialized checkpoint is dropped before advancing — no record its `channel_values` for the requested channels. The
cross-row cache, so peak in-flight is one deserialized checkpoint. deserialized checkpoint is dropped immediately — peak in-flight is one
deserialized checkpoint.
Off-path rows (different branch on the same thread) advance the Sets `walk_state["reached_root"]` when an ancestor has no parent.
cursor without doing any work. Returns True when the walk can stop: chain reached `max_supersteps`, or
the root was reached.
Returns True when every requested channel is seeded — the caller
can stop iterating and close the cursor.
""" """
if "started" not in walk_state: if "started" not in walk_state:
if cid == target_id: if cid == target_id:
walk_state["started"] = True walk_state["started"] = True
walk_state["cur_cid"] = parent_cid walk_state["cur_cid"] = parent_cid
walk_state["active"] = {ch for ch in channels if ch not in seeded} if parent_cid is None:
walk_state["reached_root"] = True
return True
# Not target yet (or target not present): keep streaming. # Not target yet (or target not present): keep streaming.
return False return False
active: set[str] = walk_state["active"] if len(shared_cpid_chain) >= max_supersteps:
if not active:
return True return True
if cid != walk_state["cur_cid"]: if cid != walk_state["cur_cid"]:
# Off-path row from a sibling branch — skip without deserializing. # Off-path row from a sibling branch — skip without deserializing.
return False return False
for ch in active: shared_cpid_chain.append(cid)
chain_by_ch[ch].append(cid) depth = len(shared_cpid_chain) # 1-indexed position along the chain
ckpt = serde.loads_typed((type_tag, blob)) # Capture channel_values at any depth a channel may use as its seed: the
channel_values: Mapping[str, Any] = ckpt.get("channel_values") or {} # exact `supersteps` depths, plus the root-most checkpoint (the seed
for ch in [ch for ch in active if ch in channel_values]: # candidate when the chain is shorter than `supersteps`).
seed_val_by_ch[ch] = channel_values[ch] if depth in needed_depths or parent_cid is None:
seeded.add(ch) ckpt = serde.loads_typed((type_tag, blob))
active.discard(ch) channel_values: Mapping[str, Any] = ckpt.get("channel_values") or {}
del ckpt, channel_values seed_values_by_depth[depth] = {
ch: channel_values[ch] for ch in channels if ch in channel_values
}
del ckpt, channel_values
if parent_cid is None:
walk_state["reached_root"] = True
return True
walk_state["cur_cid"] = parent_cid walk_state["cur_cid"] = parent_cid
return not active return len(shared_cpid_chain) >= max_supersteps
def resolve_delta_chains(
*,
channels: Sequence[str],
supersteps_by_ch: Mapping[str, int],
shared_cpid_chain: Sequence[str],
seed_values_by_depth: Mapping[int, Mapping[str, Any]],
has_reached_root: bool,
thread_id: str,
) -> tuple[
dict[str, list[str]],
dict[str, str | None],
dict[str, Any],
]:
"""Slice the shared parent chain into per-channel chain/seed mappings.
For each channel the seed snapshot sits `supersteps` hops back: the seed
checkpoint is `shared_cpid_chain[supersteps - 1]` and the chain is
`shared_cpid_chain[:supersteps]` (newest first), with the seed's inline
`channel_values[ch]` captured during the walk.
When the chain is shorter than `supersteps` but the walk reached the
root, the persisted chain is "compressed" relative to the logical
superstep count (`durability="exit"`, or a thread that never
snapshotted). The seed candidate is then the oldest persisted checkpoint
(`shared_cpid_chain[-1]`); if its `channel_values[ch]` is absent the
channel has no seed and replays the full chain on an empty baseline.
"""
chained_cpid_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
seed_cpid_by_ch: dict[str, str | None] = {ch: None for ch in channels}
seed_value_by_ch: dict[str, Any] = {}
for ch in channels:
bound = supersteps_by_ch.get(ch, 0)
if bound <= 0:
continue
if len(shared_cpid_chain) >= bound:
seed_depth = bound
elif has_reached_root:
seed_depth = len(shared_cpid_chain)
else:
logger.warning(
"cannot find seed snapshot for delta channel "
"(thread_id=%s, channel=%s)",
thread_id,
ch,
)
continue
if seed_depth <= 0:
continue
chained_cpid_by_ch[ch] = list(shared_cpid_chain[:seed_depth])
seed_cpid_by_ch[ch] = shared_cpid_chain[seed_depth - 1]
seed_vals = seed_values_by_depth.get(seed_depth, {})
if ch in seed_vals:
seed_value_by_ch[ch] = seed_vals[ch]
return chained_cpid_by_ch, seed_cpid_by_ch, seed_value_by_ch
def build_delta_channels_writes_history( def build_delta_channels_writes_history(
*, *,
channels: Sequence[str], channels: Sequence[str],
chain_by_ch: Mapping[str, list[str]], chained_cpid_by_ch: Mapping[str, Sequence[str]],
seed_val_by_ch: Mapping[str, Any], seed_cpid_by_ch: Mapping[str, str | None],
seeded: set[str], seed_value_by_ch: Mapping[str, Any],
stage2_rows: Sequence[tuple[str, str, str, int, str, bytes]], stage2_rows: Sequence[tuple[str, str, str, int, str, bytes]],
serde: Any, serde: Any,
) -> dict[str, DeltaChannelHistory]: ) -> dict[str, DeltaChannelHistory]:
"""Demux stage-2 rows per channel; produce per-channel histories. """Demux writes rows per channel; produce per-channel histories.
Stage-2 rows are `(checkpoint_id, channel, task_id, idx, type, value)`. `stage2_rows` are `(checkpoint_id, channel, task_id, idx, type, value)`.
Final write order is oldest→newest globally and `(task_id, idx)` within Final write order is oldest→newest globally and `(task_id, idx)` within
a checkpoint, matching the contract on `DeltaChannelHistory.writes`. a checkpoint, matching the contract on `DeltaChannelHistory.writes`.
`seed` is omitted when the walk reached a true root with no snapshot The seed checkpoint's own writes are replayed on top of a
found (channel never entered `seeded`); consumers treat absence as `_DeltaSnapshot` seed (the snapshot is the value *prior* to its own
"start empty". writes), but skipped for a migrated plain-value seed (a legacy
non-delta blob already incorporates those writes). `seed` is omitted
when no seed was located (implicit empty baseline); consumers treat
absence as "start empty".
""" """
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = { writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
ch: {} for ch in channels ch: {} for ch in channels
@@ -156,17 +236,26 @@ def build_delta_channels_writes_history(
result: dict[str, DeltaChannelHistory] = {} result: dict[str, DeltaChannelHistory] = {}
for ch in channels: for ch in channels:
chain_cids = chain_by_ch.get(ch, []) entry: DeltaChannelHistory = {"writes": []}
skip_seed_checkpoint_writes = False
if ch in seed_value_by_ch:
seed_value = seed_value_by_ch[ch]
entry["seed"] = seed_value
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
cid_writes = writes_by_ch_by_cid.get(ch, {}) cid_writes = writes_by_ch_by_cid.get(ch, {})
collected: list[PendingWrite] = [] if cid_writes:
# Chain is newest-first; iterate oldest-first for the public order. collected: list[PendingWrite] = []
for cid in reversed(chain_cids): seed_cpid = seed_cpid_by_ch.get(ch)
for type_tag, value_blob, task_id, _idx in cid_writes.get(cid, []): # Chain is newest→oldest; replay oldest→newest.
collected.append( for cid in reversed(list(chained_cpid_by_ch.get(ch, []))):
(task_id, ch, serde.loads_typed((type_tag, value_blob))) if skip_seed_checkpoint_writes and cid == seed_cpid:
) continue
entry: DeltaChannelHistory = {"writes": collected} for type_tag, value_blob, task_id, _idx in cid_writes.get(cid, []):
if ch in seeded: collected.append(
entry["seed"] = seed_val_by_ch[ch] (task_id, ch, serde.loads_typed((type_tag, value_blob)))
)
entry["writes"] = collected
result[ch] = entry result[ch] = entry
return result return result
@@ -25,10 +25,12 @@ from langgraph.checkpoint.base import (
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.sqlite._delta import ( from langgraph.checkpoint.sqlite._delta import (
DELTA_STAGE1_SQL, DELTA_WALK_SQL,
build_delta_channels_writes_history, build_delta_channels_writes_history,
build_delta_stage2_sql, build_delta_writes_fetch_sql,
step_walk_with_row, parse_supersteps_since_last_snapshot_by_channel,
resolve_delta_chains,
step_walk_supersteps,
) )
from langgraph.checkpoint.sqlite.utils import search_where from langgraph.checkpoint.sqlite.utils import search_where
@@ -625,61 +627,98 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`. """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
See `SqliteSaver.get_delta_channel_history` for design notes; this See `SqliteSaver.get_delta_channel_history` for design notes; this
is the async equivalent using `aiosqlite` cursors. Stage 1 pages is the async equivalent using `aiosqlite` cursors. The WALK streams
the parent chain newest-first and Python-deserializes each the parent chain newest-first, bounded by the target's
checkpoint blob to find per-channel snapshots; stage 2 fetches `counters_since_delta_snapshot` supersteps, and deserializes only
only the relevant writes via per-channel UNION ALL. seed checkpoints; FETCH pulls the relevant writes via per-channel
UNION ALL.
""" """
if not channels: if not channels:
return {} return {}
channels = list(channels) channels = list(channels)
await self.setup() await self.setup()
thread_id = str(config["configurable"]["thread_id"]) thread_id = str(config["configurable"]["thread_id"])
if not thread_id:
raise ValueError("empty thread ID")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
target = await self.aget_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
checkpoint_id = target.config["configurable"]["checkpoint_id"]
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} # Resolve the target checkpoint id + its metadata (for supersteps).
seed_val_by_ch: dict[str, Any] = {} checkpoint_id = get_checkpoint_id(config)
async with self.lock, self.conn.cursor() as cur:
if checkpoint_id is None:
await cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? "
"ORDER BY checkpoint_id DESC LIMIT 1",
(thread_id, checkpoint_ns),
)
else:
await cur.execute(
"SELECT checkpoint_id, metadata FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? "
"AND checkpoint_id = ?",
(thread_id, checkpoint_ns, checkpoint_id),
)
target_row = await cur.fetchone()
if target_row is None:
return {ch: {"writes": []} for ch in channels}
target_id = str(target_row[0])
metadata = json.loads(target_row[1]) if target_row[1] is not None else {}
supersteps_by_ch = parse_supersteps_since_last_snapshot_by_channel(
metadata, channels
)
max_supersteps = max(supersteps_by_ch.values(), default=0)
needed_depths = set(supersteps_by_ch.values())
shared_cpid_chain: list[str] = []
walk_state: dict[str, Any] = {} walk_state: dict[str, Any] = {}
seeded: set[str] = set() seed_values_by_depth: dict[int, dict[str, Any]] = {}
async with self.lock, self.conn.cursor() as cur: async with self.lock, self.conn.cursor() as cur:
await cur.execute( if max_supersteps > 0:
DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id) await cur.execute(DELTA_WALK_SQL, (thread_id, checkpoint_ns, target_id))
) async for row in cur:
async for row in cur: cid, parent_cid, type_tag, blob = row
cid, parent_cid, type_tag, blob = row if step_walk_supersteps(
if step_walk_with_row( cid=cid,
cid=cid, parent_cid=parent_cid,
parent_cid=parent_cid, type_tag=type_tag,
type_tag=type_tag, blob=blob,
blob=blob, target_id=target_id,
target_id=checkpoint_id, serde=self.serde,
serde=self.serde, shared_cpid_chain=shared_cpid_chain,
chain_by_ch=chain_by_ch, walk_state=walk_state,
seed_val_by_ch=seed_val_by_ch, max_supersteps=max_supersteps,
walk_state=walk_state, needed_depths=needed_depths,
seeded=seeded, seed_values_by_depth=seed_values_by_depth,
channels=channels, channels=channels,
): ):
break break
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] (
stage2_sql = build_delta_stage2_sql( chained_cpid_by_ch,
chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain], seed_cpid_by_ch,
seed_value_by_ch,
) = resolve_delta_chains(
channels=channels,
supersteps_by_ch=supersteps_by_ch,
shared_cpid_chain=shared_cpid_chain,
seed_values_by_depth=seed_values_by_depth,
has_reached_root=bool(walk_state.get("reached_root")),
thread_id=thread_id,
) )
if stage2_sql:
stage2_params: list[Any] = [] channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
fetch_sql = build_delta_writes_fetch_sql(
chain_lens=[len(chained_cpid_by_ch[ch]) for ch in channels_with_chain],
)
if fetch_sql:
fetch_params: list[Any] = []
for ch in channels_with_chain: for ch in channels_with_chain:
stage2_params.extend( fetch_params.extend(
[thread_id, checkpoint_ns, ch, *chain_by_ch[ch]] [thread_id, checkpoint_ns, ch, *chained_cpid_by_ch[ch]]
) )
await cur.execute(stage2_sql, stage2_params) await cur.execute(fetch_sql, fetch_params)
stage2_rows = cast( stage2_rows = cast(
"list[tuple[str, str, str, int, str, bytes]]", "list[tuple[str, str, str, int, str, bytes]]",
await cur.fetchall(), await cur.fetchall(),
@@ -689,9 +728,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
return build_delta_channels_writes_history( return build_delta_channels_writes_history(
channels=channels, channels=channels,
chain_by_ch=chain_by_ch, chained_cpid_by_ch=chained_cpid_by_ch,
seed_val_by_ch=seed_val_by_ch, seed_cpid_by_ch=seed_cpid_by_ch,
seeded=seeded, seed_value_by_ch=seed_value_by_ch,
stage2_rows=stage2_rows, stage2_rows=stage2_rows,
serde=self.serde, serde=self.serde,
) )
@@ -34,6 +34,37 @@ PendingWrite = tuple[str, str, Any]
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def _parse_supersteps_since_last_snapshot_by_channel(
metadata: Mapping[str, Any],
channels: Sequence[str],
) -> dict[str, int]:
"""Per-channel supersteps-since-last-snapshot, parsed from metadata.
Reads `metadata.counters_since_delta_snapshot[ch]`, a `(updates,
supersteps)` pair, and returns `{ch: supersteps}` for channels whose
supersteps count is positive. Channels with no counter entry (just
snapshotted, or never written) are omitted — their seed, if any, is the
target checkpoint itself and needs no ancestor walk.
Used to drive the ancestor-walk depth in `get_delta_channel_history`:
a channel's seed snapshot sits exactly `supersteps` hops back along the
parent chain.
"""
counters = metadata.get("counters_since_delta_snapshot") or {}
result: dict[str, int] = {}
for ch in channels:
entry = counters.get(ch)
if not isinstance(entry, (list, tuple)) or len(entry) < 2:
continue
try:
supersteps = int(entry[1])
except (TypeError, ValueError):
continue
if supersteps > 0:
result[ch] = supersteps
return result
# Marked as total=False to allow for future expansion. # Marked as total=False to allow for future expansion.
class CheckpointMetadata(TypedDict, total=False): class CheckpointMetadata(TypedDict, total=False):
"""Metadata associated with a checkpoint.""" """Metadata associated with a checkpoint."""
@@ -619,34 +650,30 @@ class BaseCheckpointSaver(Generic[V]):
""" """
if not channels: if not channels:
return {} return {}
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels} channels = list(channels)
seed_by_ch: dict[str, Any] = {}
remaining: set[str] = set(channels)
target_tuple = self.get_tuple(config) target_tuple = self.get_tuple(config)
cursor_config: RunnableConfig | None = ( if target_tuple is None:
target_tuple.parent_config if target_tuple else None return {ch: {"writes": []} for ch in channels}
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
target_tuple.metadata or {}, channels
) )
while cursor_config is not None and remaining: max_supersteps = max(supersteps_by_ch.values(), default=0)
chain: list[CheckpointTuple] = []
has_reached_root = False
cursor_config: RunnableConfig | None = target_tuple.parent_config
while cursor_config is not None and len(chain) < max_supersteps:
tup = self.get_tuple(cursor_config) tup = self.get_tuple(cursor_config)
if tup is None: if tup is None:
break break
if tup.pending_writes: chain.append(tup)
for write in reversed(tup.pending_writes): if tup.parent_config is None:
ch = write[1] has_reached_root = True
if ch in remaining: break
collected_by_ch[ch].append(write)
for ch in list(remaining):
if ch in tup.checkpoint["channel_values"]:
seed_by_ch[ch] = tup.checkpoint["channel_values"][ch]
remaining.discard(ch)
cursor_config = tup.parent_config cursor_config = tup.parent_config
result: dict[str, DeltaChannelHistory] = {} return self._assemble_default_delta_history(
for ch in channels: channels, supersteps_by_ch, chain, has_reached_root, config
entry: DeltaChannelHistory = {"writes": list(reversed(collected_by_ch[ch]))} )
if ch in seed_by_ch:
entry["seed"] = seed_by_ch[ch]
result[ch] = entry
return result
async def aget_delta_channel_history( async def aget_delta_channel_history(
self, *, config: RunnableConfig, channels: Sequence[str] self, *, config: RunnableConfig, channels: Sequence[str]
@@ -660,32 +687,106 @@ class BaseCheckpointSaver(Generic[V]):
""" """
if not channels: if not channels:
return {} return {}
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels} channels = list(channels)
seed_by_ch: dict[str, Any] = {}
remaining: set[str] = set(channels)
target_tuple = await self.aget_tuple(config) target_tuple = await self.aget_tuple(config)
cursor_config: RunnableConfig | None = ( if target_tuple is None:
target_tuple.parent_config if target_tuple else None return {ch: {"writes": []} for ch in channels}
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
target_tuple.metadata or {}, channels
) )
while cursor_config is not None and remaining: max_supersteps = max(supersteps_by_ch.values(), default=0)
chain: list[CheckpointTuple] = []
has_reached_root = False
cursor_config: RunnableConfig | None = target_tuple.parent_config
while cursor_config is not None and len(chain) < max_supersteps:
tup = await self.aget_tuple(cursor_config) tup = await self.aget_tuple(cursor_config)
if tup is None: if tup is None:
break break
if tup.pending_writes: chain.append(tup)
for write in reversed(tup.pending_writes): if tup.parent_config is None:
ch = write[1] has_reached_root = True
if ch in remaining: break
collected_by_ch[ch].append(write)
for ch in list(remaining):
if ch in tup.checkpoint["channel_values"]:
seed_by_ch[ch] = tup.checkpoint["channel_values"][ch]
remaining.discard(ch)
cursor_config = tup.parent_config cursor_config = tup.parent_config
return self._assemble_default_delta_history(
channels, supersteps_by_ch, chain, has_reached_root, config
)
def _assemble_default_delta_history(
self,
channels: Sequence[str],
supersteps_by_ch: Mapping[str, int],
chain: Sequence[CheckpointTuple],
has_reached_root: bool,
config: RunnableConfig,
) -> dict[str, DeltaChannelHistory]:
"""Slice the walked parent chain into per-channel histories.
For each channel the seed snapshot sits `supersteps` hops back along
`chain` (newest first): the seed checkpoint is `chain[supersteps - 1]`
and its `channel_values[ch]` is the seed value. The seed checkpoint's
own writes are replayed on top of a `_DeltaSnapshot` seed but skipped
for a migrated plain-value seed (the legacy non-delta blob already
incorporates them).
When the chain is shorter than `supersteps` but the root was reached,
the persisted chain is "compressed" relative to the logical superstep
count (`durability="exit"`, or a thread that never snapshotted). The
seed candidate is then the oldest checkpoint (`chain[-1]`); if its
`channel_values[ch]` is absent the channel has no seed and replays
the full chain on an empty baseline.
"""
# Imported lazily to avoid a hard checkpoint→serde-types coupling at
# module import; only the delta surface needs the runtime check.
from langgraph.checkpoint.serde.types import _DeltaSnapshot
thread_id = config["configurable"].get("thread_id")
result: dict[str, DeltaChannelHistory] = {} result: dict[str, DeltaChannelHistory] = {}
for ch in channels: for ch in channels:
entry: DeltaChannelHistory = {"writes": list(reversed(collected_by_ch[ch]))} entry: DeltaChannelHistory = {"writes": []}
if ch in seed_by_ch: bound = supersteps_by_ch.get(ch, 0)
entry["seed"] = seed_by_ch[ch] if bound <= 0:
result[ch] = entry
continue
if len(chain) >= bound:
seed_depth = bound
elif has_reached_root:
seed_depth = len(chain)
else:
logger.warning(
"cannot find seed snapshot for delta channel "
"(thread_id=%s, channel=%s)",
thread_id,
ch,
)
result[ch] = entry
continue
if seed_depth <= 0:
result[ch] = entry
continue
chain_slice: Sequence[CheckpointTuple] = chain[:seed_depth]
seed_tuple: CheckpointTuple | None = chain[seed_depth - 1]
skip_seed_checkpoint_writes = False
channel_values = seed_tuple.checkpoint["channel_values"]
if ch in channel_values:
seed_value = channel_values[ch]
entry["seed"] = seed_value
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
else:
# No stored value at the oldest checkpoint → empty baseline,
# so the chain's own writes are all replayed from empty.
seed_tuple = None
collected: list[PendingWrite] = []
# Chain is newest→oldest; replay oldest→newest.
for tup in reversed(list(chain_slice)):
if skip_seed_checkpoint_writes and tup is seed_tuple:
continue
for write in tup.pending_writes or []:
if write[1] == ch:
collected.append(write)
entry["writes"] = collected
result[ch] = entry result[ch] = entry
return result return result
@@ -23,6 +23,7 @@ from langgraph.checkpoint.base import (
DeltaChannelHistory, DeltaChannelHistory,
PendingWrite, PendingWrite,
SerializerProtocol, SerializerProtocol,
_parse_supersteps_since_last_snapshot_by_channel,
get_checkpoint_id, get_checkpoint_id,
get_checkpoint_metadata, get_checkpoint_metadata,
) )
@@ -144,14 +145,20 @@ class InMemorySaver(
) -> Mapping[str, DeltaChannelHistory]: ) -> Mapping[str, DeltaChannelHistory]:
"""Override: walk the parent chain ONCE for all requested channels. """Override: walk the parent chain ONCE for all requested channels.
Each channel terminates independently at the nearest ancestor Walk depth is driven by the target checkpoint's
whose stored blob is non-empty. Other channels keep walking until `counters_since_delta_snapshot[ch]` supersteps counter: each
they find their own terminator or hit the root. channel's seed snapshot sits exactly `supersteps` hops back along
the parent chain. This locates seeds reliably even when they aren't
a `_DeltaSnapshot` sentinel — e.g. a legacy plain-value blob from a
thread migrated off a non-delta channel.
Pre-delta plain-value blobs subsume their ancestor's pending Pre-delta plain-value seeds subsume their own checkpoint's pending
writes (the value already includes them); `_DeltaSnapshot` blobs writes (the value already includes them), so those are skipped;
do not (snapshot is the value AT that ancestor, prior to its own `_DeltaSnapshot` seeds do not (the snapshot is the value AT that
pending writes that produce the child). ancestor, prior to its own pending writes), so they are replayed.
When the chain reaches the root short of `supersteps`, the snapshot
was never persisted (implicit empty baseline) — replay the full
chain with no seed.
""" """
if not channels: if not channels:
return {} return {}
@@ -159,72 +166,98 @@ class InMemorySaver(
# module import; only this override needs the runtime check. # module import; only this override needs the runtime check.
from langgraph.checkpoint.serde.types import _DeltaSnapshot from langgraph.checkpoint.serde.types import _DeltaSnapshot
channels = list(channels)
thread_id = config["configurable"]["thread_id"] thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
# Resolve the target checkpoint id + its metadata (for supersteps).
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
checkpoint_id = next(reversed(ns_storage), None)
target_entry = ns_storage.get(checkpoint_id) if checkpoint_id else None
if target_entry is None:
return {ch: {"writes": []} for ch in channels}
metadata = self.serde.loads_typed(target_entry[1])
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
metadata or {}, channels
)
max_supersteps = max(supersteps_by_ch.values(), default=0)
# Walk the parent chain (newest first) up to the deepest supersteps.
chain: list[str] = [] chain: list[str] = []
target_entry = ns_storage.get(checkpoint_id) has_reached_root = False
current: str | None = target_entry[2] if target_entry is not None else None current: str | None = target_entry[2]
while current is not None: while current is not None and len(chain) < max_supersteps:
entry = ns_storage.get(current) entry = ns_storage.get(current)
if entry is None: if entry is None:
break break
chain.append(current) chain.append(current)
_, _, parent = entry parent = entry[2]
current = parent if parent is None:
has_reached_root = True
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
seed_by_ch: dict[str, Any] = {}
remaining: set[str] = set(channels)
for cp_id in chain:
if not remaining:
break break
entry = ns_storage.get(cp_id) current = parent
ckpt = self.serde.loads_typed(entry[0]) if entry is not None else None
terminated_here: set[str] = set()
blob_value_by_ch: dict[str, Any] = {}
if ckpt is not None:
versions = ckpt.get("channel_versions", {})
for ch in remaining:
ver = versions.get(ch)
if ver is None:
continue
blob_entry = self.blobs.get((thread_id, checkpoint_ns, ch, ver))
if blob_entry is None or blob_entry[0] == "empty":
continue
blob_value_by_ch[ch] = self.serde.loads_typed(blob_entry)
terminated_here.add(ch)
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
step_writes.items(), reverse=True
):
if ch not in remaining:
continue
blob_value = blob_value_by_ch.get(ch)
if blob_value is not None and not isinstance(
blob_value, _DeltaSnapshot
):
continue
collected_by_ch[ch].append(
(tid, ch, self.serde.loads_typed(serialized))
)
for ch in terminated_here:
seed_by_ch[ch] = blob_value_by_ch[ch]
remaining.discard(ch)
result: dict[str, DeltaChannelHistory] = {} result: dict[str, DeltaChannelHistory] = {}
for ch in channels: for ch in channels:
entry_h: DeltaChannelHistory = { entry_h: DeltaChannelHistory = {"writes": []}
"writes": list(reversed(collected_by_ch[ch])) bound = supersteps_by_ch.get(ch, 0)
} if bound <= 0:
if ch in seed_by_ch: result[ch] = entry_h
entry_h["seed"] = seed_by_ch[ch] continue
if len(chain) >= bound:
seed_depth = bound
elif has_reached_root:
seed_depth = len(chain)
else:
logger.warning(
"cannot find seed snapshot for delta channel "
"(thread_id=%s, channel=%s)",
thread_id,
ch,
)
result[ch] = entry_h
continue
if seed_depth <= 0:
result[ch] = entry_h
continue
chain_slice: list[str] = chain[:seed_depth]
seed_cpid: str | None = chain[seed_depth - 1]
skip_seed_checkpoint_writes = False
seed_entry = ns_storage.get(seed_cpid)
if seed_entry is not None:
ckpt = self.serde.loads_typed(seed_entry[0])
ver = (ckpt.get("channel_versions") or {}).get(ch)
blob_entry = (
self.blobs.get((thread_id, checkpoint_ns, ch, ver))
if ver is not None
else None
)
if blob_entry is not None and blob_entry[0] != "empty":
seed_value = self.serde.loads_typed(blob_entry)
entry_h["seed"] = seed_value
skip_seed_checkpoint_writes = not isinstance(
seed_value, _DeltaSnapshot
)
else:
# No stored value at the oldest checkpoint → empty
# baseline; the chain's own writes replay from empty.
seed_cpid = None
collected: list[PendingWrite] = []
# Chain is newest→oldest; replay oldest→newest.
for cp_id in reversed(chain_slice):
if skip_seed_checkpoint_writes and cp_id == seed_cpid:
continue
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
for (_task_id, _idx), (tid, w_ch, serialized, _) in sorted(
step_writes.items()
):
if w_ch == ch:
collected.append((tid, ch, self.serde.loads_typed(serialized)))
entry_h["writes"] = collected
result[ch] = entry_h result[ch] = entry_h
return result return result
+15 -3
View File
@@ -351,9 +351,12 @@ class TestInMemorySaverDeltaChannel:
cp1["id"] = "cp1" cp1["id"] = "cp1"
cp2 = empty_checkpoint() cp2 = empty_checkpoint()
cp2["id"] = "cp2" cp2["id"] = "cp2"
# Target (cp2) carries the supersteps counter; no snapshot was ever
# taken, so the walk runs back to the root → no seed, empty baseline.
cp2_md = {"counters_since_delta_snapshot": {channel: [1, 2]}}
saver.storage[thread_id][ns] = { saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), "cp2": (serde.dumps_typed(cp2), serde.dumps_typed(cp2_md), "cp1"),
} }
# Writes stored at cp1 produced the cp1 snapshot; part of history. # Writes stored at cp1 produced the cp1 snapshot; part of history.
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = ( saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
@@ -456,10 +459,15 @@ class TestBaseFallbackGetChannelWrites:
cp1["id"] = "00000000000000000000000000000002.0000000000000000" cp1["id"] = "00000000000000000000000000000002.0000000000000000"
cp2 = empty_checkpoint() cp2 = empty_checkpoint()
cp2["id"] = "00000000000000000000000000000003.0000000000000000" cp2["id"] = "00000000000000000000000000000003.0000000000000000"
# The target (cp2) carries `counters_since_delta_snapshot` — the
# supersteps-since-last-snapshot driving the ancestor-walk depth. No
# snapshot was ever taken, so the count runs back past the oldest
# persisted checkpoint (root) → implicit empty baseline, no seed.
cp2_md = {"counters_since_delta_snapshot": {channel: [2, 3]}}
saver.storage[thread_id][ns] = { saver.storage[thread_id][ns] = {
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None), cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]), cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]), cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed(cp2_md), cp1["id"]),
} }
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's. # Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = ( saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
@@ -607,10 +615,14 @@ class TestPreDeltaBlobTerminator:
cp3["id"] = "cp3" cp3["id"] = "cp3"
cp3["channel_versions"][channel] = v3 cp3["channel_versions"][channel] = v3
# cp1 is the pre-delta (migration) seed; the channel became a
# DeltaChannel at cp2, so its supersteps counter starts there: cp2 -> 1,
# cp3 -> 2. The target (cp3) walks 2 hops back to the cp1 seed blob.
cp3_md = {"counters_since_delta_snapshot": {channel: [2, 2]}}
saver.storage[thread_id][ns] = { saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), "cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"), "cp3": (serde.dumps_typed(cp3), serde.dumps_typed(cp3_md), "cp2"),
} }
# Write under cp1 would be from the pre-delta era and MUST be ignored # Write under cp1 would be from the pre-delta era and MUST be ignored
# (the blob already captures it). We add one and assert it is not # (the blob already captures it). We add one and assert it is not
@@ -70,6 +70,42 @@ def delta_channels_to_snapshot(
return result return result
def advance_delta_counters(
channels: Mapping[str, BaseChannel],
prev_metadata: Mapping[str, Any] | None,
updated_channels: set[str] | None,
) -> tuple[dict[str, tuple[int, int]], set[str]]:
"""Advance per-delta-channel `(updates, supersteps)` counters one step.
Mirrors the counter bookkeeping in `PregelLoop._put_checkpoint`: every
superstep bumps `supersteps` for all delta channels and `updates` for
those written this step, then channels that hit their snapshot cadence
are reset to `(0, 0)`.
Used by `update_state` (which creates a delta checkpoint outside the
main loop) so the written checkpoint carries a correct
`counters_since_delta_snapshot` — without it, the supersteps-based
`get_delta_channel_history` walk can't locate the seed. Returns
`(new_counters, channels_to_snapshot)`; callers drop `(0, 0)` entries
before writing the metadata field.
"""
prev = dict((prev_metadata or {}).get("counters_since_delta_snapshot") or {})
updated = updated_channels or set()
new_counters: dict[str, tuple[int, int]] = {}
for name, ch in channels.items():
if not isinstance(ch, DeltaChannel):
continue
u, s = prev.get(name, (0, 0))
s += 1
if name in updated:
u += 1
new_counters[name] = (u, s)
channels_to_snapshot = delta_channels_to_snapshot(channels, new_counters)
for k in channels_to_snapshot:
new_counters[k] = (0, 0)
return new_counters, channels_to_snapshot
def create_checkpoint( def create_checkpoint(
checkpoint: Checkpoint, checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel] | None, channels: Mapping[str, BaseChannel] | None,
+63 -12
View File
@@ -130,6 +130,7 @@ from langgraph.pregel._algo import (
from langgraph.pregel._call import identifier from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import ( from langgraph.pregel._checkpoint import (
achannels_from_checkpoint, achannels_from_checkpoint,
advance_delta_counters,
channels_from_checkpoint, channels_from_checkpoint,
copy_checkpoint, copy_checkpoint,
create_checkpoint, create_checkpoint,
@@ -2032,15 +2033,40 @@ class Pregel(
checkpointer.get_next_version, checkpointer.get_next_version,
self.trigger_to_nodes, self.trigger_to_nodes,
) )
checkpoint = create_checkpoint(checkpoint, channels, step + 1) # Advance delta-channel snapshot counters for this manual super-
# step so the written checkpoint carries a correct
# `counters_since_delta_snapshot` (mirrors the main loop). Without
# it, the supersteps-based delta history walk can't find the seed.
updated_channel_names = {
c for task in run_tasks for c, _ in task.writes if c != PUSH
}
new_counters, channels_to_snapshot = advance_delta_counters(
channels,
saved.metadata if saved else None,
updated_channel_names,
)
# Keep `updated_channels` at its default (None) — matching the
# historical checkpoint shape here — so resume/trigger semantics
# (e.g. deferred nodes) are unaffected. `get_next_version` /
# `channels_to_snapshot` are no-ops unless a delta channel snapshots.
checkpoint = create_checkpoint(
checkpoint,
channels,
step + 1,
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
)
update_metadata: dict[str, Any] = {
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
}
if non_zero := {k: v for k, v in new_counters.items() if v != (0, 0)}:
update_metadata["counters_since_delta_snapshot"] = non_zero
next_config = checkpointer.put( next_config = checkpointer.put(
checkpoint_config, checkpoint_config,
checkpoint, checkpoint,
{ update_metadata,
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
},
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"] checkpoint_previous_versions, checkpoint["channel_versions"]
), ),
@@ -2500,16 +2526,41 @@ class Pregel(
checkpointer.get_next_version, checkpointer.get_next_version,
self.trigger_to_nodes, self.trigger_to_nodes,
) )
checkpoint = create_checkpoint(checkpoint, channels, step + 1) # Advance delta-channel snapshot counters for this manual super-
# step so the written checkpoint carries a correct
# `counters_since_delta_snapshot` (mirrors the main loop). Without
# it, the supersteps-based delta history walk can't find the seed.
updated_channel_names = {
c for task in run_tasks for c, _ in task.writes if c != PUSH
}
new_counters, channels_to_snapshot = advance_delta_counters(
channels,
saved.metadata if saved else None,
updated_channel_names,
)
# Keep `updated_channels` at its default (None) — matching the
# historical checkpoint shape here — so resume/trigger semantics
# (e.g. deferred nodes) are unaffected. `get_next_version` /
# `channels_to_snapshot` are no-ops unless a delta channel snapshots.
checkpoint = create_checkpoint(
checkpoint,
channels,
step + 1,
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
)
update_metadata: dict[str, Any] = {
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
}
if non_zero := {k: v for k, v in new_counters.items() if v != (0, 0)}:
update_metadata["counters_since_delta_snapshot"] = non_zero
# save checkpoint, after applying writes # save checkpoint, after applying writes
next_config = await checkpointer.aput( next_config = await checkpointer.aput(
checkpoint_config, checkpoint_config,
checkpoint, checkpoint,
{ update_metadata,
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
},
get_new_channel_versions( get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"] checkpoint_previous_versions, checkpoint["channel_versions"]
), ),