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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+151 -47
View File
@@ -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(
{
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+165 -55
View File
@@ -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,
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+75 -45
View File
@@ -4,10 +4,9 @@ import copy
import logging
import threading
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import ( # noqa: UP035
from typing import (
Any,
Generic,
List,
Literal,
NamedTuple,
TypedDict,
@@ -30,6 +29,9 @@ from langgraph.checkpoint.serde.types import (
SCHEDULED,
ChannelProtocol,
)
from langgraph.checkpoint.serde.types import (
SEED_UNSET as SEED_UNSET,
)
from langgraph.checkpoint.serde.types import (
DeltaChannelWrites as DeltaChannelWrites,
)
@@ -504,71 +506,99 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006
"""Collect writes for `channel` across this checkpoint's ancestry, oldest→newest.
def get_channel_writes(
self, config: RunnableConfig, channel: str
) -> DeltaChannelWrites:
"""Reconstruct a `DeltaChannel`'s write history at this checkpoint.
Scans newest→oldest and stops at the first `Overwrite` (either from
`snapshot_every` or user code), so reconstruction cost is bounded.
Default implementation walks the full thread history via `list()`; savers
can override with a more efficient query (InMemorySaver and PostgresSaver
do this).
Returns a `DeltaChannelWrites` carrying:
`List` is used instead of `list` to avoid mypy confusing it with the
saver's own `list` method.
* `writes` — per-step deltas from ancestors, oldest→newest, ready
to be replayed through the reducer in `DeltaChannel.from_checkpoint`.
* `seed` — when the ancestor walk hits a pre-delta blob (a value
stored before `DeltaChannel` was enabled for this field), replay
starts from that snapshot instead of the channel's empty value.
Default `SEED_UNSET` means no seed.
Walks the **parent chain** (not `list(before=...)`): for a thread with
forks, only on-path ancestors contribute. Scans newest→oldest and
stops at the first `Overwrite`, so reconstruction cost is bounded.
Writes stored at the target `checkpoint_id` itself are pending writes
for the next step and are excluded — pregel applies them separately
via `apply_writes`.
The base implementation uses `get_tuple` and `pending_writes`; it
never sees blobs, so it never sets `seed`. Savers that can read the
blob table directly (`InMemorySaver`, `PostgresSaver`) override this
method to set `seed` when appropriate, which both shortens the walk
and recovers state from pre-delta threads after migration.
"""
# Guard against re-entrant calls: when list() triggers reconstruction
# which calls list() again, the inner call returns tuples with
# DELTA_SENTINEL in channel_values (which get_channel_writes ignores —
# it only reads pending_writes). This breaks the recursion safely.
# Guard against re-entrant calls: when get_tuple() triggers
# reconstruction which calls get_tuple() again, the inner call
# returns tuples with DELTA_SENTINEL in channel_values (which this
# method ignores — it only reads pending_writes).
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
return DeltaChannelWrites(writes=[])
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = [] # newest first
for tup in self.list(list_config, before=before_config): # newest → oldest
if not tup.pending_writes:
continue
# Within a superstep, pending_writes are oldest→newest; reverse
# to scan newest-first.
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return collected
target_tuple = self.get_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = self.get_tuple(cursor_config)
if tup is None:
break
if tup.pending_writes:
# Within a superstep, pending_writes are oldest→newest;
# reverse to scan newest-first.
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return DeltaChannelWrites(writes=collected)
cursor_config = tup.parent_config
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
finally:
_DELTA_RECONSTRUCTION.active = False
async def aget_channel_writes(
self, config: RunnableConfig, channel: str
) -> List[Any]: # noqa: UP006
"""Async version of get_channel_writes."""
) -> DeltaChannelWrites:
"""Async version of `get_channel_writes`. See docstring there."""
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
return DeltaChannelWrites(writes=[])
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = []
async for tup in self.alist(list_config, before=before_config):
if not tup.pending_writes:
continue
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return collected
target_tuple = await self.aget_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = await self.aget_tuple(cursor_config)
if tup is None:
break
if tup.pending_writes:
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return DeltaChannelWrites(writes=collected)
cursor_config = tup.parent_config
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
finally:
_DELTA_RECONSTRUCTION.active = False
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+45 -16
View File
@@ -146,22 +146,26 @@ class InMemorySaver(
channel_values: dict[str, Any],
) -> None:
"""Replace DELTA_SENTINEL entries with DeltaChannelWrites so
DeltaChannel.from_checkpoint can distinguish reconstructed writes from
a pre-DeltaChannel accumulated list."""
`DeltaChannel.from_checkpoint` can distinguish reconstructed writes
from a pre-DeltaChannel accumulated value."""
for channel, value in channel_values.items():
if value is DELTA_SENTINEL:
channel_values[channel] = DeltaChannelWrites(
self.get_channel_writes(config, channel)
)
channel_values[channel] = self.get_channel_writes(config, channel)
def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]:
def get_channel_writes(
self, config: RunnableConfig, channel: str
) -> DeltaChannelWrites:
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
# Walk the parent chain newest→oldest collecting checkpoint IDs.
# Walk the parent chain newest→oldest. Skip the target itself —
# writes stored AT `checkpoint_id` are pending for the next step
# (pregel applies them via `apply_writes`; they aren't part of the
# snapshot value AT `checkpoint_id`).
chain: list[str] = []
current: str | None = checkpoint_id
target_entry = ns_storage.get(checkpoint_id)
current: str | None = target_entry[2] if target_entry is not None else None
while current is not None:
entry = ns_storage.get(current)
if entry is None:
@@ -171,14 +175,39 @@ class InMemorySaver(
current = parent
overwrite_types = _overwrite_types()
# Scan writes newest→oldest. Stop at the first `Overwrite` — it
# dominates all older history. Either from `snapshot_every` or from
# user code: the bound applies the same way.
# Scan newest→oldest. Two terminators stop the walk:
# 1. a user-emitted `Overwrite` in writes — replaces prior history;
# 2. a pre-delta blob on an ancestor — bind it as `seed`.
# Without (2), a thread migrated from pre-delta storage would replay
# ancestor writes all the way to the root AND miss any value that
# lived only in the old blob (e.g. from `update_state`).
#
# At each ancestor, check the blob BEFORE processing its pending
# writes: a pre-delta blob represents the state AT that ancestor,
# which already subsumes any writes stored under it. Processing
# those writes first would fold them into the reconstructed value
# twice (once via the blob, once via replay).
collected: list[Any] = [] # newest first
for cp_id in chain: # newest → oldest
entry = ns_storage.get(cp_id)
if entry is not None:
ckpt = self.serde.loads_typed(entry[0])
ver = ckpt.get("channel_versions", {}).get(channel)
if ver is not None:
blob_entry = self.blobs.get(
(thread_id, checkpoint_ns, channel, ver)
)
if blob_entry is not None and blob_entry[0] != "empty":
blob_value = self.serde.loads_typed(blob_entry)
if blob_value is not DELTA_SENTINEL:
# Pre-delta snapshot terminator. Skip this
# ancestor's writes — the blob subsumes them.
collected.reverse()
return DeltaChannelWrites(writes=collected, seed=blob_value)
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
# within a superstep, sorted by (task_id, idx) = oldest → newest,
# so reverse to get newest-first scan.
# Within a superstep, sorted by (task_id, idx) = oldest → newest;
# reverse for newest-first scan.
for (_task_id, _idx), (_, ch, serialized, _) in sorted(
step_writes.items(), reverse=True
):
@@ -188,13 +217,13 @@ class InMemorySaver(
collected.append(val)
if isinstance(val, overwrite_types):
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
async def aget_channel_writes(
self, config: RunnableConfig, channel: str
) -> list[Any]:
) -> DeltaChannelWrites:
return self.get_channel_writes(config, channel)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
libs/checkpoint/langgraph/checkpoint/serde/types.py
+23
View File
@@ -34,14 +34,37 @@ class _DeltaSentinel:
DELTA_SENTINEL = _DeltaSentinel()
class _SeedUnset:
"""Marker used as the default for `DeltaChannelWrites.seed`.
Distinct from `None`, which is a legitimate pre-delta value
(e.g. an `Optional` field whose accumulated value really was `None`).
"""
__slots__ = ()
def __repr__(self) -> str:
return "SEED_UNSET"
SEED_UNSET = _SeedUnset()
@dataclasses.dataclass
class DeltaChannelWrites:
"""In-memory wrapper around per-step writes reconstructed by a saver.
Consumed by `DeltaChannel.from_checkpoint`. Never serialized — if this
reaches the wire, something upstream forgot to unwrap it.
`seed` is the value from which chain replay should begin. When the saver
encounters a pre-delta blob during the ancestor walk, it uses that blob
as the seed and stops walking further back (the older chain is
represented by the seed). `SEED_UNSET` means "no seed — replay from the
channel's empty value".
"""
writes: list[Any]
seed: Any = SEED_UNSET
Value = TypeVar("Value", covariant=True)
libs/checkpoint/tests/test_memory.py
+165 -20
View File
@@ -6,8 +6,11 @@ from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
SEED_UNSET,
Checkpoint,
CheckpointMetadata,
DeltaChannelWrites,
create_checkpoint,
empty_checkpoint,
)
@@ -208,8 +211,6 @@ class TestMemorySaver:
async def test_memory_saver() -> None:
from langgraph.checkpoint.memory import InMemorySaver
memory_saver = InMemorySaver()
assert isinstance(memory_saver, InMemorySaver)
@@ -325,11 +326,6 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
class TestInMemorySaverDeltaChannel:
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
empty_checkpoint,
)
saver = InMemorySaver()
serde = JsonPlusSerializer()
@@ -349,10 +345,10 @@ class TestInMemorySaverDeltaChannel:
assert channel in result
assert result[channel] is DELTA_SENTINEL
def test_get_channel_writes_collects_writes(self) -> None:
"""get_channel_writes collects per-step writes oldest→newest."""
from langgraph.checkpoint.base import empty_checkpoint
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
"""get_channel_writes collects ancestor writes oldest→newest, and
excludes writes stored at the target checkpoint itself (those are
pending writes for the next step, applied separately by pregel)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
@@ -366,18 +362,20 @@ class TestInMemorySaverDeltaChannel:
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
}
# cp1 has a write for channel
# Writes stored at cp1 produced the cp1 snapshot; part of history.
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "hi"}),
"",
)
# cp2 has a write for channel
# Writes stored at cp2 are pending — they will produce cp3 when the
# step that loaded cp2 completes. They MUST NOT appear in the
# reconstructed snapshot value at cp2.
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed({"content": "bye"}),
serde.dumps_typed({"content": "pending"}),
"",
)
@@ -389,7 +387,37 @@ class TestInMemorySaverDeltaChannel:
}
}
result = saver.get_channel_writes(config, channel)
assert result == [{"content": "hi"}, {"content": "bye"}]
assert result == DeltaChannelWrites(writes=[{"content": "hi"}])
assert result.seed is SEED_UNSET
def test_get_channel_writes_at_root_returns_empty(self) -> None:
"""Reconstructing the root checkpoint's state: no ancestors → []."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
}
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "pending"}),
"",
)
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": "cp1",
}
}
assert saver.get_channel_writes(config, channel) == DeltaChannelWrites(
writes=[]
)
class TestBaseFallbackGetChannelWrites:
@@ -461,7 +489,10 @@ class TestBaseFallbackGetChannelWrites:
result = saver.get_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
assert result == DeltaChannelWrites(
writes=[{"content": "first"}, {"content": "second"}]
)
assert result.seed is SEED_UNSET
async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None:
saver, thread_id, ns = self._build_saver_with_chain()
@@ -476,7 +507,10 @@ class TestBaseFallbackGetChannelWrites:
result = await saver.aget_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
assert result == DeltaChannelWrites(
writes=[{"content": "first"}, {"content": "second"}]
)
assert result.seed is SEED_UNSET
def test_fallback_stops_at_first_overwrite(self) -> None:
"""An `Overwrite` dominates older history: scan newest→oldest stops at
@@ -509,6 +543,117 @@ class TestBaseFallbackGetChannelWrites:
result = saver.get_channel_writes(config, "messages")
assert len(result) == 1
assert isinstance(result[0], Overwrite)
assert result[0].value == [{"content": "reset"}]
assert len(result.writes) == 1
assert isinstance(result.writes[0], Overwrite)
assert result.writes[0].value == [{"content": "reset"}]
assert result.seed is SEED_UNSET
class TestPreDeltaBlobTerminator:
"""Verify the pre-delta blob terminator: when the ancestor walk hits a
checkpoint whose blob for the channel is a real value (not
DELTA_SENTINEL), reconstruction seeds from it and stops. This guards
* back-compat: a thread written by pre-delta code, then extended under
delta — reconstruction must return the correct value without walking
past the last pre-delta ancestor;
* perf: without the terminator, every reconstruct-after-migration would
walk all the way to the thread root.
"""
def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]:
"""Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta,
write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield
seed=[A] + writes=[B, C].
Returns `(saver, thread_id, ns, channel, cp3_id)`.
"""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0"
v2 = "00000000000000000000000000000002.0"
v3 = "00000000000000000000000000000003.0"
# Pre-delta: cp1 stored a real blob for the channel.
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
cp2 = empty_checkpoint()
cp2["id"] = "cp2"
cp2["channel_versions"][channel] = v2
cp3 = empty_checkpoint()
cp3["id"] = "cp3"
cp3["channel_versions"][channel] = v3
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
}
# 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
# folded into the reconstructed result.
saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = (
"task0",
channel,
serde.dumps_typed("PRE-DELTA-WRITE"),
"",
)
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed("B"),
"",
)
saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = (
"task3",
channel,
serde.dumps_typed("PENDING-AT-TARGET"),
"",
)
return saver, thread_id, ns, channel, "cp3"
def test_seed_from_pre_delta_ancestor_blob(self) -> None:
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver.get_channel_writes(config, channel)
# Seed came from the pre-delta blob at cp1.
assert result.seed == ["A"]
# Delta-era writes from cp2 replay through the reducer on top of seed.
# cp3 is the target — its own write is pending for the NEXT step and
# must be excluded.
assert result.writes == ["B"]
def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None:
"""Writes stored at the pre-delta ancestor itself must not be replayed
(the blob subsumes them)."""
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver.get_channel_writes(config, channel)
# The pre-delta write under cp1 must not appear (the blob subsumes it).
assert "PRE-DELTA-WRITE" not in result.writes
# And the pending write at the target is never folded in.
assert "PENDING-AT-TARGET" not in result.writes
libs/langgraph/langgraph/channels/delta.py
+10 -5
View File
@@ -4,7 +4,7 @@ import copy as _copy
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DELTA_SENTINEL, DeltaChannelWrites
from langgraph.checkpoint.base import DELTA_SENTINEL, SEED_UNSET, DeltaChannelWrites
from typing_extensions import Self
from langgraph._internal._typing import MISSING
@@ -124,9 +124,14 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
new._writes_since_snapshot = 0
elif isinstance(checkpoint, DeltaChannelWrites):
# Saver reconstructed per-step writes; replay through the operator.
# Counter tracks writes since the last Overwrite so snapshot cadence
# stays accurate across reloads.
value: Any = _empty(new.typ)
# `seed` (if set) is a pre-delta accumulated value that terminates
# the ancestor walk on the saver side: replay starts from it
# instead of the channel's empty value. Counter tracks writes
# since the last Overwrite so snapshot cadence stays accurate
# across reloads.
value: Any = (
_empty(new.typ) if checkpoint.seed is SEED_UNSET else checkpoint.seed
)
counter = 0
for write in checkpoint.writes:
value, counter = new._apply_write(value, write, counter)
@@ -134,7 +139,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
new._writes_since_snapshot = counter
else:
# Backward compat: a pre-DeltaChannel thread stored the accumulated
# value directly. Trust it as-is.
# value directly (no saver-side reconstruction happened). Trust it.
new.value = checkpoint
new._writes_since_snapshot = 0
return new
libs/langgraph/tests/test_channels.py
+50
View File
@@ -2,13 +2,17 @@ import operator
from collections.abc import Sequence
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import SEED_UNSET, DeltaChannelWrites
from langgraph._internal._typing import MISSING
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.graph.message import add_messages
pytestmark = pytest.mark.anyio
@@ -664,6 +668,52 @@ def test_delta_channel_user_overwrite_resets_counter() -> None:
assert ch.should_snapshot()
def test_delta_channel_from_checkpoint_honors_seed() -> None:
"""DeltaChannelWrites(seed=...) starts replay from that snapshot.
Guards the pre-delta migration path: when the saver's ancestor walk hits
a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs
the post-migration state correctly rather than replaying from empty.
"""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="pre-delta", id="p1")]
writes = DeltaChannelWrites(
writes=[
AIMessage(content="delta-1", id="d1"),
HumanMessage(content="delta-2", id="d2"),
],
seed=seed,
)
ch = spec.from_checkpoint(writes)
msgs = ch.get()
assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"]
def test_delta_channel_from_checkpoint_seed_without_writes() -> None:
"""Reconstruction at a pre-delta ancestor with no newer deltas returns
just the seed — the saver's terminator fired immediately."""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="only-snap", id="s1")]
ch = spec.from_checkpoint(DeltaChannelWrites(writes=[], seed=seed))
assert ch.get() == seed
def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_unset() -> None:
"""`seed=None` must start replay from None, not from the channel's empty
value. `SEED_UNSET` is the sentinel meaning 'no seed'."""
def replace(left, right):
return right
spec = DeltaChannel(replace)
ch = spec.from_checkpoint(DeltaChannelWrites(writes=["after"], seed=None))
# Reducer replaces; seed=None → first write produces "after".
assert ch.get() == "after"
# And the default (unset) is distinct.
unset = DeltaChannelWrites(writes=["after"])
assert unset.seed is SEED_UNSET
def test_delta_channel_replay_tracks_counter_across_overwrite() -> None:
"""Counter reloaded from writes reflects writes-since-last-Overwrite."""
from langchain_core.messages import HumanMessage