fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk

Four fixes from an independent review of the reconstruction pipeline, plus
a structural cleanup:

1. Ancestor walk excludes the target checkpoint itself (matches pregel:
   writes stored under checkpoint_id=T are pending for the NEXT step and
   applied separately via apply_writes). Memory saver previously included
   them, diverging from Postgres and causing pending writes to be folded
   into the reconstructed snapshot — visible via get_state during
   interrupts and time-travel into a non-leaf checkpoint.

2. Pre-delta blob terminator. When the walk hits an ancestor whose blob
   for the channel is a real value (not DELTA_SENTINEL), bind that blob
   as DeltaChannelWrites.seed and stop. Without this, threads migrated
   from pre-delta storage would replay ancestor writes to the root
   forever AND lose any value that lived only in the old blob
   (e.g. from update_state). Per-ancestor, the blob is checked BEFORE
   its writes — a pre-delta blob subsumes writes at the same checkpoint,
   so including them would double-count.

3. Base-fallback get_channel_writes follows parent_checkpoint_id instead
   of list(before=...). The previous form returned every tuple with
   id<target, including sibling branches on forked threads.

4. seed replaces the Overwrite-wrapping hack for pre-delta values.
   DeltaChannelWrites(writes, seed=SEED_UNSET) makes the saver's
   reconstruction terminator semantically explicit; drops the lazy
   _make_overwrite import dance. User-emitted Overwrite still reset the
   chain via _apply_write as before.

Postgres: recursive CTE enumerates on-path ancestors and joins once
against checkpoint_writes and once against checkpoint_blobs for every
delta channel in the get_tuple — one roundtrip instead of the previous
3 queries × N channels.

Tests added:
- Pre-delta blob seeding (seed binding, no double-counting of ancestor
  writes at the terminator, pending-at-target excluded).
- Root checkpoint returns empty writes.
- Seed-based from_checkpoint replay (three scenarios: with writes,
  seed-only, seed=None distinct from SEED_UNSET).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-30 14:49:05 -04:00
co-authored by Claude Opus 4.7
parent 5f1e946b1a
commit ffacba950a
8 changed files with 684 additions and 188 deletions
@@ -394,64 +394,162 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
async def _aget_channel_writes_cur(
async def _areconstruct_delta_channels_cur(
self,
*,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
channels: Sequence[str],
cur: Any,
) -> list[Any]:
"""Async version of _get_channel_writes_cur — see sync version for rationale."""
) -> dict[str, DeltaChannelWrites]:
"""Async mirror of `_reconstruct_delta_channels_cur`.
Single recursive CTE enumerates on-path ancestors, LEFT-joined once
against `checkpoint_writes` and once against `checkpoint_blobs`. Per
channel the walk stops at the first terminator — a user `Overwrite`
in writes or a pre-delta blob (captured as `seed`). See the sync
docstring for the full rationale.
"""
if not channels:
return {}
overwrite_types = _overwrite_types()
channels_list = list(channels)
await cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s",
(thread_id, checkpoint_ns),
"""
WITH RECURSIVE ancestors(cid, parent, depth) AS (
SELECT parent_checkpoint_id, NULL::text, 0
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
AND checkpoint_id = %s AND parent_checkpoint_id IS NOT NULL
UNION ALL
SELECT c.parent_checkpoint_id, a.cid, a.depth + 1
FROM checkpoints c
JOIN ancestors a ON c.checkpoint_id = a.cid
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
AND c.parent_checkpoint_id IS NOT NULL
),
walk AS (
SELECT a.cid, a.depth, c.checkpoint
FROM ancestors a
JOIN checkpoints c
ON c.thread_id = %s AND c.checkpoint_ns = %s
AND c.checkpoint_id = a.cid
)
SELECT w.cid, w.depth,
cw.channel AS write_channel, cw.type AS write_type,
cw.blob AS write_blob, cw.task_id, cw.idx,
bl.channel AS blob_channel, bl.type AS blob_type,
bl.blob AS blob_blob
FROM walk w
LEFT JOIN checkpoint_writes cw
ON cw.thread_id = %s AND cw.checkpoint_ns = %s
AND cw.checkpoint_id = w.cid AND cw.channel = ANY(%s)
LEFT JOIN checkpoint_blobs bl
ON bl.thread_id = %s AND bl.checkpoint_ns = %s
AND bl.channel = ANY(%s)
AND bl.version = (w.checkpoint->'channel_versions'->>bl.channel)
ORDER BY w.depth ASC, cw.task_id DESC, cw.idx DESC
""",
(
thread_id,
checkpoint_ns,
checkpoint_id,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channels_list,
thread_id,
checkpoint_ns,
channels_list,
),
)
parent_map: dict[str, str | None] = {
row["checkpoint_id"]: row["parent_checkpoint_id"]
for row in await cur.fetchall()
}
ancestor_ids: list[str] = []
cid: str | None = parent_map.get(checkpoint_id)
while cid is not None:
ancestor_ids.append(cid)
cid = parent_map.get(cid)
if not ancestor_ids:
return []
await cur.execute(
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
" AND checkpoint_id = ANY(%s) "
"ORDER BY task_id DESC, idx DESC",
(thread_id, checkpoint_ns, channel, ancestor_ids),
)
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
rows_by_cid: dict[str, dict[str, Any]] = {}
cid_order: list[str] = []
seen_blob: set[tuple[str, str]] = set()
seen_write: set[tuple[str, str, str, int]] = set()
for row in await cur.fetchall():
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
collected: list[Any] = []
for cid in ancestor_ids:
for type_tag, blob in writes_by_cp.get(cid, []):
val = self.serde.loads_typed((type_tag, blob))
collected.append(val)
if isinstance(val, overwrite_types):
collected.reverse()
return collected
collected.reverse()
return collected
cid = row["cid"]
if cid not in rows_by_cid:
rows_by_cid[cid] = {"writes_per_channel": {}, "blob_per_channel": {}}
cid_order.append(cid)
ch_w = row["write_channel"]
if ch_w is not None:
key = (cid, ch_w, row["task_id"], row["idx"])
if key not in seen_write:
seen_write.add(key)
rows_by_cid[cid]["writes_per_channel"].setdefault(ch_w, []).append(
(row["write_type"], row["write_blob"])
)
ch_b = row["blob_channel"]
if ch_b is not None and (cid, ch_b) not in seen_blob:
seen_blob.add((cid, ch_b))
rows_by_cid[cid]["blob_per_channel"][ch_b] = (
row["blob_type"],
row["blob_blob"],
)
collected: dict[str, list[Any]] = {ch: [] for ch in channels_list}
done: set[str] = set()
seeds: dict[str, Any] = {}
for cid in cid_order:
bucket = rows_by_cid[cid]
# Pre-delta blob check first — subsumes any writes at this cp.
for ch in list(channels_list):
if ch in done:
continue
blob = bucket["blob_per_channel"].get(ch)
if blob is None or blob[0] == "empty":
continue
blob_value = self.serde.loads_typed(blob)
if blob_value is DELTA_SENTINEL:
continue
seeds[ch] = blob_value
done.add(ch)
for ch in list(channels_list):
if ch in done:
continue
for type_tag, blob in bucket["writes_per_channel"].get(ch, []):
val = self.serde.loads_typed((type_tag, blob))
collected[ch].append(val)
if isinstance(val, overwrite_types):
done.add(ch)
break
if len(done) == len(channels_list):
break
result: dict[str, DeltaChannelWrites] = {}
for ch in channels_list:
ch_writes = collected[ch]
ch_writes.reverse()
if ch in seeds:
result[ch] = DeltaChannelWrites(writes=ch_writes, seed=seeds[ch])
else:
result[ch] = DeltaChannelWrites(writes=ch_writes)
return result
async def aget_channel_writes(
self, config: RunnableConfig, channel: str
) -> list[Any]:
) -> DeltaChannelWrites:
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"]["checkpoint_id"]
async with self._cursor() as cur:
return await self._aget_channel_writes_cur(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
result = await self._areconstruct_delta_channels_cur(
thread_id=thread_id,
checkpoint_ns=checkpoint_ns,
checkpoint_id=checkpoint_id,
channels=[channel],
cur=cur,
)
return result.get(channel, DeltaChannelWrites(writes=[]))
async def _load_checkpoint_tuple(
self, value: DictRow, cur: AsyncCursor[DictRow]
@@ -478,13 +576,19 @@ class AsyncPostgresSaver(BasePostgresSaver):
channel_values: dict[str, Any] = {}
if blob_values:
channel_values = self._load_blobs(blob_values)
for channel, v in channel_values.items():
if v is DELTA_SENTINEL:
channel_values[channel] = DeltaChannelWrites(
await self._aget_channel_writes_cur(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
)
)
delta_channels = [
ch for ch, v in channel_values.items() if v is DELTA_SENTINEL
]
if delta_channels:
reconstructed = await self._areconstruct_delta_channels_cur(
thread_id=thread_id,
checkpoint_ns=checkpoint_ns,
checkpoint_id=checkpoint_id,
channels=delta_channels,
cur=cur,
)
for ch, writes in reconstructed.items():
channel_values[ch] = writes
return CheckpointTuple(
{
@@ -2,7 +2,6 @@ from __future__ import annotations
import random
import warnings
from collections import defaultdict
from collections.abc import Sequence
from importlib.metadata import version as get_version
from typing import Any, cast
@@ -207,73 +206,184 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
channel_values: dict[str, Any],
cur: Any,
) -> None:
for channel, value in channel_values.items():
if value is DELTA_SENTINEL:
channel_values[channel] = DeltaChannelWrites(
self._get_channel_writes_cur(
config["configurable"]["thread_id"],
config["configurable"].get("checkpoint_ns", ""),
config["configurable"]["checkpoint_id"],
channel,
cur,
)
)
delta_channels = [ch for ch, v in channel_values.items() if v is DELTA_SENTINEL]
if not delta_channels:
return
reconstructed = self._reconstruct_delta_channels_cur(
thread_id=config["configurable"]["thread_id"],
checkpoint_ns=config["configurable"].get("checkpoint_ns", ""),
checkpoint_id=config["configurable"]["checkpoint_id"],
channels=delta_channels,
cur=cur,
)
for ch, writes in reconstructed.items():
channel_values[ch] = writes
def _get_channel_writes_cur(
def _reconstruct_delta_channels_cur(
self,
*,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
channels: Sequence[str],
cur: Any,
) -> list[Any]:
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.
) -> dict[str, DeltaChannelWrites]:
"""Reconstruct `DeltaChannelWrites` for every `channel` in `channels`.
Scans newest→oldest and stops at the first `Overwrite` — a snapshot
marker (either from `snapshot_every` or user code) dominates all older
writes. Two queries:
A single recursive CTE enumerates on-path ancestors (never siblings),
left-joined once against `checkpoint_writes` and once against
`checkpoint_blobs`. One roundtrip covers every delta channel in the
`get_tuple` — avoids the N-channels × 3-queries blowup and fits the
intent of the schema (`parent_checkpoint_id` is an indexed ancestor
pointer; postgres's planner handles this CTE shape well for the
ancestor depths seen in practice).
1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread.
2. Walk ancestry in Python, then fetch writes with a plain ANY() filter.
The walk newest→oldest stops per channel at the first terminator:
* a user-emitted `Overwrite` in `checkpoint_writes` — replaces
prior history;
* a non-sentinel blob in `checkpoint_blobs` — a pre-delta snapshot;
bound as `DeltaChannelWrites.seed` so replay starts from it.
Writes stored AT `checkpoint_id` are pending for the next step and
excluded; the recursion starts from the target's parent.
"""
overwrite_types = _overwrite_types()
if not channels:
return {}
overwrite_types = _overwrite_types()
channels_list = list(channels)
# depth=0 → target's parent (we never read writes or blob for the
# target itself). A NULL parent stops the recursion.
cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s",
(thread_id, checkpoint_ns),
"""
WITH RECURSIVE ancestors(cid, parent, depth) AS (
SELECT parent_checkpoint_id, NULL::text, 0
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
AND checkpoint_id = %s AND parent_checkpoint_id IS NOT NULL
UNION ALL
SELECT c.parent_checkpoint_id, a.cid, a.depth + 1
FROM checkpoints c
JOIN ancestors a ON c.checkpoint_id = a.cid
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
AND c.parent_checkpoint_id IS NOT NULL
),
walk AS (
-- Each ancestor plus the channel_versions mapping for blob join.
SELECT a.cid, a.depth, c.checkpoint
FROM ancestors a
JOIN checkpoints c
ON c.thread_id = %s AND c.checkpoint_ns = %s
AND c.checkpoint_id = a.cid
)
SELECT w.cid, w.depth,
cw.channel AS write_channel, cw.type AS write_type,
cw.blob AS write_blob, cw.task_id, cw.idx,
bl.channel AS blob_channel, bl.type AS blob_type,
bl.blob AS blob_blob
FROM walk w
LEFT JOIN checkpoint_writes cw
ON cw.thread_id = %s AND cw.checkpoint_ns = %s
AND cw.checkpoint_id = w.cid AND cw.channel = ANY(%s)
LEFT JOIN checkpoint_blobs bl
ON bl.thread_id = %s AND bl.checkpoint_ns = %s
AND bl.channel = ANY(%s)
AND bl.version = (w.checkpoint->'channel_versions'->>bl.channel)
ORDER BY w.depth ASC, cw.task_id DESC, cw.idx DESC
""",
(
thread_id,
checkpoint_ns,
checkpoint_id, # anchor
thread_id,
checkpoint_ns, # recursion
thread_id,
checkpoint_ns, # walk join
thread_id,
checkpoint_ns,
channels_list, # writes join
thread_id,
checkpoint_ns,
channels_list, # blobs join
),
)
parent_map: dict[str, str | None] = {
row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall()
}
ancestor_ids: list[str] = []
cid: str | None = parent_map.get(checkpoint_id)
while cid is not None:
ancestor_ids.append(cid)
cid = parent_map.get(cid)
if not ancestor_ids:
return []
# Order newest→oldest so we can stop at the first Overwrite.
cur.execute(
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
" AND checkpoint_id = ANY(%s) "
"ORDER BY task_id DESC, idx DESC",
(thread_id, checkpoint_ns, channel, ancestor_ids),
)
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
# Group incoming rows by `cid` so we can, per ancestor, decide whether
# to terminate (pre-delta blob) BEFORE processing writes at that
# ancestor. Rows within a cid arrive in (task_id DESC, idx DESC)
# order; we preserve that when building per-cid writes lists.
rows_by_cid: dict[str, dict[str, Any]] = {}
cid_order: list[str] = [] # newest → oldest
seen_blob: set[tuple[str, str]] = set()
seen_write: set[tuple[str, str, str, int]] = set()
for row in cur.fetchall():
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
collected: list[Any] = [] # newest first
for cid in ancestor_ids: # newest → oldest
for type_tag, blob in writes_by_cp.get(cid, []):
val = self.serde.loads_typed((type_tag, blob))
collected.append(val)
if isinstance(val, overwrite_types):
collected.reverse()
return collected
collected.reverse()
return collected
cid = row["cid"]
if cid not in rows_by_cid:
rows_by_cid[cid] = {"writes_per_channel": {}, "blob_per_channel": {}}
cid_order.append(cid)
# Writes: dedupe via (cid, channel, task_id, idx).
ch_w = row["write_channel"]
if ch_w is not None:
key = (cid, ch_w, row["task_id"], row["idx"])
if key not in seen_write:
seen_write.add(key)
rows_by_cid[cid]["writes_per_channel"].setdefault(ch_w, []).append(
(row["write_type"], row["write_blob"])
)
# Blobs: dedupe via (cid, channel).
ch_b = row["blob_channel"]
if ch_b is not None and (cid, ch_b) not in seen_blob:
seen_blob.add((cid, ch_b))
rows_by_cid[cid]["blob_per_channel"][ch_b] = (
row["blob_type"],
row["blob_blob"],
)
# Per-channel state. `collected[ch]` is newest→oldest during the walk.
collected: dict[str, list[Any]] = {ch: [] for ch in channels_list}
done: set[str] = set()
seeds: dict[str, Any] = {}
for cid in cid_order: # newest → oldest
bucket = rows_by_cid[cid]
# At each ancestor, check the blob FIRST — a pre-delta blob
# subsumes any writes stored under the same checkpoint, so we
# must not fold those writes in before terminating.
for ch in list(channels_list):
if ch in done:
continue
blob = bucket["blob_per_channel"].get(ch)
if blob is None or blob[0] == "empty":
continue
blob_value = self.serde.loads_typed(blob)
if blob_value is DELTA_SENTINEL:
continue
seeds[ch] = blob_value
done.add(ch)
# Then process per-channel writes for any channel still live.
for ch in list(channels_list):
if ch in done:
continue
for type_tag, blob in bucket["writes_per_channel"].get(ch, []):
val = self.serde.loads_typed((type_tag, blob))
collected[ch].append(val)
if isinstance(val, overwrite_types):
done.add(ch)
break
if len(done) == len(channels_list):
break
result: dict[str, DeltaChannelWrites] = {}
for ch in channels_list:
ch_writes = collected[ch]
ch_writes.reverse() # oldest → newest
if ch in seeds:
result[ch] = DeltaChannelWrites(writes=ch_writes, seed=seeds[ch])
else:
result[ch] = DeltaChannelWrites(writes=ch_writes)
return result
def _dump_blobs(
self,