From ffacba950a2dc262f8883f568c68fd4e640d75ff Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 22 Apr 2026 22:09:08 -0400 Subject: [PATCH] fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../langgraph/checkpoint/postgres/aio.py | 198 ++++++++++++---- .../langgraph/checkpoint/postgres/base.py | 220 +++++++++++++----- .../langgraph/checkpoint/base/__init__.py | 120 ++++++---- .../langgraph/checkpoint/memory/__init__.py | 61 +++-- .../langgraph/checkpoint/serde/types.py | 23 ++ libs/checkpoint/tests/test_memory.py | 185 +++++++++++++-- libs/langgraph/langgraph/channels/delta.py | 15 +- libs/langgraph/tests/test_channels.py | 50 ++++ 8 files changed, 684 insertions(+), 188 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 686fca289..7dbefcf51 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -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( { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 5c06d5648..c5837ca61 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -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, diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index a007f0e71..5ec709dee 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -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 diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index d9558c8f6..05633418f 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -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: diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index c8df72960..29da8f67b 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -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) diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 2bc18ccd4..ed129838c 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -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 diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 6ab6d0914..56cc8123f 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -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 diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 63a7d986c..2f049f953 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -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