diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index f3f692641..2f4c6be76 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations import threading from collections import defaultdict -from collections.abc import Iterator, Sequence +from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from typing import Any, cast @@ -27,10 +27,9 @@ from psycopg_pool import ConnectionPool from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import ( - SELECT_DELTA_STAGE1_SQL, SELECT_DELTA_STAGE2_SQL, BasePostgresSaver, - _DeltaStage1Row, + _build_delta_stage1_sql, _DeltaStage2Row, ) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver @@ -444,53 +443,79 @@ class PostgresSaver(BasePostgresSaver): with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - def _get_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: - """Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`. + def _get_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + """Fast-path override of `BaseCheckpointSaver._get_all_delta_channels_writes_history`. - Two-stage query: stage 1 scans checkpoint metadata to walk the parent - chain and locate the nearest snapshot; stage 2 fetches only the - chain-limited writes and single seed blob. + Two-stage query, both stages cover ALL requested channels in a single + Postgres roundtrip each: + + * Stage 1: dynamic SELECT over `checkpoints` with K parallel JSONB + key lookups (one column pair per channel) — no subquery, no + aggregation. Returns one row per checkpoint with versions and + snapshot flags for every requested channel. + + * Stage 2: one UNION ALL over `checkpoint_writes` and + `checkpoint_blobs` filtered by `channel = ANY(?)` and per-channel + chain_cids / seed_versions (collapsed across channels). """ + if not channels: + return {} + channels = list(channels) thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) if checkpoint_id is None: target = self.get_tuple(config) if target is None: - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + return { + ch: _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + for ch in channels + } checkpoint_id = target.config["configurable"]["checkpoint_id"] + # Stage 1: K parallel JSONB lookups per row, one query for all channels. + stage1_sql = _build_delta_stage1_sql(channels) + stage1_params: list[Any] = [] + for ch in channels: + stage1_params.extend([ch, ch]) + stage1_params.extend([thread_id, checkpoint_ns]) with self._cursor() as cur: - cur.execute( - SELECT_DELTA_STAGE1_SQL, - (channel, channel, thread_id, checkpoint_ns), - ) + cur.execute(stage1_sql, stage1_params) stage1_rows = cur.fetchall() - chain_cids, seed_version = self._walk_stage1( - cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id + + chain_by_ch, seed_ver_by_ch = self._walk_stage1_multi( + cast("list[Mapping[str, Any]]", stage1_rows), checkpoint_id, channels ) - seed_versions = [seed_version] if seed_version else [] + # Union of chain cids and seed versions across all channels. + union_chain_cids: list[str] = sorted( + {cid for chain in chain_by_ch.values() for cid in chain} + ) + union_seed_versions: list[str] = sorted( + {ver for ver in seed_ver_by_ch.values() if ver is not None} + ) + + # Stage 2: chain-limited writes + chain-limited seed blobs for all channels. with self._cursor() as cur: cur.execute( SELECT_DELTA_STAGE2_SQL, ( thread_id, checkpoint_ns, - channel, - chain_cids, + channels, + union_chain_cids, thread_id, checkpoint_ns, - channel, - seed_versions, + channels, + union_seed_versions, ), ) stage2_rows = cur.fetchall() - return self._build_delta_channel_writes_history( - channel=channel, - chain_cids=chain_cids, - seed_version=seed_version, + return self._build_delta_channels_writes_history( + channels=channels, + chain_by_ch=chain_by_ch, + seed_ver_by_ch=seed_ver_by_ch, stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), ) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 50e73c245..b6aec4824 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -2,7 +2,7 @@ from __future__ import annotations import asyncio from collections import defaultdict -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from typing import Any, cast @@ -27,10 +27,9 @@ from psycopg_pool import AsyncConnectionPool from langgraph.checkpoint.postgres import _ainternal from langgraph.checkpoint.postgres.base import ( - SELECT_DELTA_STAGE1_SQL, SELECT_DELTA_STAGE2_SQL, BasePostgresSaver, - _DeltaStage1Row, + _build_delta_stage1_sql, _DeltaStage2Row, ) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver @@ -405,53 +404,68 @@ class AsyncPostgresSaver(BasePostgresSaver): async with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - async def _aget_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: - """Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`. + async def _aget_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + """Fast-path override of `BaseCheckpointSaver._aget_all_delta_channels_writes_history`. - Two-stage query: stage 1 scans checkpoint metadata to walk the parent - chain and locate the nearest snapshot; stage 2 fetches only the - chain-limited writes and single seed blob. + Two-stage query, both stages cover ALL requested channels in a single + Postgres roundtrip each. See `PostgresSaver._get_all_delta_channels_writes_history` + for design notes. """ + if not channels: + return {} + channels = list(channels) thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = get_checkpoint_id(config) if checkpoint_id is None: target = await self.aget_tuple(config) if target is None: - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + return { + ch: _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + for ch in channels + } checkpoint_id = target.config["configurable"]["checkpoint_id"] + stage1_sql = _build_delta_stage1_sql(channels) + stage1_params: list[Any] = [] + for ch in channels: + stage1_params.extend([ch, ch]) + stage1_params.extend([thread_id, checkpoint_ns]) async with self._cursor() as cur: - await cur.execute( - SELECT_DELTA_STAGE1_SQL, - (channel, channel, thread_id, checkpoint_ns), - ) + await cur.execute(stage1_sql, stage1_params) stage1_rows = await cur.fetchall() - chain_cids, seed_version = self._walk_stage1( - cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id + + chain_by_ch, seed_ver_by_ch = self._walk_stage1_multi( + cast("list[Mapping[str, Any]]", stage1_rows), checkpoint_id, channels ) - seed_versions = [seed_version] if seed_version else [] + union_chain_cids: list[str] = sorted( + {cid for chain in chain_by_ch.values() for cid in chain} + ) + union_seed_versions: list[str] = sorted( + {ver for ver in seed_ver_by_ch.values() if ver is not None} + ) + async with self._cursor() as cur: await cur.execute( SELECT_DELTA_STAGE2_SQL, ( thread_id, checkpoint_ns, - channel, - chain_cids, + channels, + union_chain_cids, thread_id, checkpoint_ns, - channel, - seed_versions, + channels, + union_seed_versions, ), ) stage2_rows = await cur.fetchall() - return self._build_delta_channel_writes_history( - channel=channel, - chain_cids=chain_cids, - seed_version=seed_version, + return self._build_delta_channels_writes_history( + channels=channels, + chain_by_ch=chain_by_ch, + seed_ver_by_ch=seed_ver_by_ch, stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows), ) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index d881b1154..121042954 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -2,7 +2,7 @@ from __future__ import annotations import random import warnings -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from importlib.metadata import version as get_version from typing import Any, TypedDict, cast @@ -161,6 +161,7 @@ class _DeltaStage2Row(TypedDict, total=False): _kind: str # "w" or "b" checkpoint_id: str | None # "w" rows only + channel: str | None # set on both "w" and "b" rows type: str | None blob: bytes | None task_id: str | None # "w" rows only @@ -168,48 +169,85 @@ class _DeltaStage2Row(TypedDict, total=False): version: str | None # "b" rows only -# Two-stage DeltaChannel reconstruction. Stage 1 scans checkpoint -# metadata (no blob bytes) to walk the parent chain and locate the -# nearest snapshot marker. Stage 2 fetches only the chain-limited -# writes and the single seed snapshot blob. +# Multi-channel two-stage DeltaChannel reconstruction. # -# Parameter order: -# stage1: (channel, channel, thread_id, checkpoint_ns) -# stage2: (thread_id, checkpoint_ns, channel, chain_cids[], -# thread_id, checkpoint_ns, channel, seed_versions[]) +# Stage 1 scans checkpoint metadata (no blob bytes) and emits one row per +# checkpoint with K parallel JSONB key lookups (one column pair per +# requested delta channel: ver_i / hs_i). No subqueries, no aggregation. +# Python walks the parent chain once across all channels. +# +# Stage 2 fetches all writes and the seed blobs for ALL channels in a +# single roundtrip via `channel = ANY(%s)` and chain/seed-version +# filtering. +# +# Empirical comparison vs an alternative "ship full channel_versions / +# channel_values JSONB and let Python pick" form (1000 checkpoints, +# 8 total channels in graph, 3 delta channels requested): +# +# Postgres execution: A=0.24ms vs B=0.38ms (both negligible) +# End-to-end latency: A=6.83ms vs B=2.28ms (B is 3.0x faster) +# Wire payload: A=836KB vs B=330KB (61% smaller) +# Buffer hits: identical (167 blocks) +# +# B (this dynamic-columns design) wins because it avoids JSONB +# serialization on the wire and JSONB-to-dict deserialization in +# psycopg. Even at K=8 (8 delta channels = 16 dynamic columns), B +# still beats A end-to-end (4.2ms vs 6.8ms). + + +def _build_delta_stage1_sql(channels: Sequence[str]) -> str: + """Build stage 1 SQL with 2K parallel JSONB key lookups. + + For channels=["messages", "files"] the result is:: + + SELECT checkpoint_id, parent_checkpoint_id, + checkpoint -> 'channel_versions' ->> %s AS ver_0, + (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0, + checkpoint -> 'channel_versions' ->> %s AS ver_1, + (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1 + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s + + Channel names are passed as `%s` parameters (safe from SQL injection). + Only the column aliases `ver_i` / `hs_i` are interpolated into the + SQL string (i is bounded by len(channels) and uses safe identifiers). + + Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ..., + thread_id, ns]`. + """ + cols = [] + for i in range(len(channels)): + cols.append( + f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, " + f"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}" + ) + return ( + "SELECT checkpoint_id, parent_checkpoint_id, " + + ", ".join(cols) + + " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s" + ) -SELECT_DELTA_STAGE1_SQL = """ - SELECT checkpoint_id, - parent_checkpoint_id, - checkpoint -> 'channel_versions' ->> %s AS ver, - (checkpoint -> 'channel_values' -> %s) IS NOT NULL AS has_snapshot - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s -""" SELECT_DELTA_STAGE2_SQL = """ SELECT 'w'::text AS _kind, - checkpoint_id, + checkpoint_id, channel, type, blob, task_id, idx, NULL::text AS version FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s + WHERE thread_id = %s AND checkpoint_ns = %s AND channel = ANY(%s) AND checkpoint_id = ANY(%s) UNION ALL - SELECT 'b', NULL, + SELECT 'b', NULL, channel, type, blob, NULL, NULL, version FROM checkpoint_blobs - WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s + WHERE thread_id = %s AND checkpoint_ns = %s AND channel = ANY(%s) AND version = ANY(%s) """ -class _DeltaStage1Row(TypedDict): - """One row from `SELECT_DELTA_STAGE1_SQL`.""" - - checkpoint_id: str - parent_checkpoint_id: str | None - ver: str | None - has_snapshot: bool +# Stage 1 rows are dynamic-shape dicts: {checkpoint_id, parent_checkpoint_id, +# ver_0, hs_0, ver_1, hs_1, ...}. Walking is parameterized by the channel +# list to map indices back to channel names — no static TypedDict here. +# `dict[str, Any]` is the practical signature. class BasePostgresSaver(BaseCheckpointSaver[str]): @@ -255,86 +293,119 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): } @staticmethod - def _walk_stage1( - stage1_rows: Sequence[_DeltaStage1Row], + def _walk_stage1_multi( + stage1_rows: Sequence[Mapping[str, Any]], target_id: str, - ) -> tuple[list[str], str | None]: - """Walk the parent chain from stage 1 metadata rows. + channels: Sequence[str], + ) -> tuple[dict[str, list[str]], dict[str, str | None]]: + """Walk the parent chain once for all requested channels. - Returns (chain_cids, seed_version): - chain_cids: ancestor checkpoint IDs from target's parent down to - the seed (or root), in newest-first order. - seed_version: the channel blob version at the nearest ancestor - with has_snapshot=True, or None if pure delta. + Each row carries `ver_i` / `hs_i` per channel index. We walk the + parent chain from target's parent toward the root; for each + channel we stop at the nearest ancestor where `hs_i` is true and + record that ancestor's `ver_i` as the seed version. All + ancestors visited up to (and including) a channel's seed are in + that channel's `chain_cids`. + + Returns: + chain_cids_by_channel: per-channel list of ancestor cids in + newest-first order. + seed_version_by_channel: per-channel seed version (None if + walk reached root with no snapshot). """ parent_of: dict[str, str | None] = {} - ver_of: dict[str, str | None] = {} - snapshot_of: dict[str, bool] = {} + # For each channel index, store ver and has_snapshot per cid. + ver_by_i_by_cid: list[dict[str, str | None]] = [ + {} for _ in range(len(channels)) + ] + hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in range(len(channels))] + for r in stage1_rows: - cid = r["checkpoint_id"] - parent_of[cid] = r["parent_checkpoint_id"] - ver_of[cid] = r["ver"] - snapshot_of[cid] = r["has_snapshot"] + cid = cast(str, r["checkpoint_id"]) + parent_of[cid] = cast("str | None", r["parent_checkpoint_id"]) + for i in range(len(channels)): + ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}")) + hs_by_i_by_cid[i][cid] = bool(r.get(f"hs_{i}")) - chain_cids: list[str] = [] - seed_version: str | None = None - cur_cid: str | None = parent_of.get(target_id) - while cur_cid is not None: - chain_cids.append(cur_cid) - if snapshot_of.get(cur_cid, False): - seed_version = ver_of.get(cur_cid) - break - cur_cid = parent_of.get(cur_cid) - return chain_cids, seed_version + chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} + seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels} + # For each channel, walk from target's parent until we hit a + # snapshot or the root. Walks share the parent_of mapping but + # are otherwise independent. + for i, ch in enumerate(channels): + cur_cid: str | None = parent_of.get(target_id) + while cur_cid is not None: + chain_by_ch[ch].append(cur_cid) + if hs_by_i_by_cid[i].get(cur_cid, False): + seed_ver_by_ch[ch] = ver_by_i_by_cid[i].get(cur_cid) + break + cur_cid = parent_of.get(cur_cid) + return chain_by_ch, seed_ver_by_ch - def _build_delta_channel_writes_history( + def _build_delta_channels_writes_history( self, *, - channel: str, - chain_cids: list[str], - seed_version: str | None, + channels: Sequence[str], + chain_by_ch: dict[str, list[str]], + seed_ver_by_ch: dict[str, str | None], stage2_rows: Sequence[_DeltaStage2Row], - ) -> _ChannelWritesHistory: - """Reconstruct delta channel history from two-stage query results. + ) -> dict[str, _ChannelWritesHistory]: + """Demux stage 2 rows per channel; produce per-channel histories. - chain_cids are in newest-first order (target's parent first). - stage2_rows contain only writes for chain_cids and the single - seed blob at seed_version. + stage2_rows carry `channel` on every row. We build per-channel + `writes_by_cid` and per-channel `seed_blob` dicts, then assemble + a `_ChannelWritesHistory` per requested channel. """ - writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {} - seed_blob: tuple[str, bytes] | None = None + # writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx) + writes_by_ch_by_cid: dict[ + str, dict[str, list[tuple[str, bytes, str, int]]] + ] = {ch: {} for ch in channels} + # seed_blob_by_ver[(channel, version)] = (type, blob) + seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {} for r in stage2_rows: + ch = cast(str, r["channel"]) kind = r["_kind"] if kind == "w": cid = cast(str, r["checkpoint_id"]) - writes_by_cid.setdefault(cid, []).append( + writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append( cast( "tuple[str, bytes, str, int]", (r["type"], r["blob"], r["task_id"], r["idx"]), ) ) else: # kind == "b" - seed_blob = cast("tuple[str, bytes]", (r["type"], r["blob"])) + ver = cast(str, r["version"]) + seed_blob_by_ver[(ch, ver)] = cast( + "tuple[str, bytes]", (r["type"], r["blob"]) + ) - for ws in writes_by_cid.values(): - ws.sort(key=lambda w: (w[2], w[3]), reverse=True) + # Sort writes per (channel, cid) newest-first by (task_id, idx) + for cid_map in writes_by_ch_by_cid.values(): + for ws in cid_map.values(): + ws.sort(key=lambda w: (w[2], w[3]), reverse=True) - if not chain_cids: - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + result: dict[str, _ChannelWritesHistory] = {} + for ch in channels: + chain_cids = chain_by_ch.get(ch, []) + seed_version = seed_ver_by_ch.get(ch) - collected: list[PendingWrite] = [] - for cid in chain_cids: - for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []): - val = self.serde.loads_typed((type_tag, write_blob)) - collected.append((task_id, channel, val)) + collected: list[PendingWrite] = [] + cid_writes = writes_by_ch_by_cid.get(ch, {}) + for cid in chain_cids: + for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []): + val = self.serde.loads_typed((type_tag, write_blob)) + collected.append((task_id, ch, val)) - seed: Any = DELTA_SENTINEL - if seed_blob is not None and seed_blob[0] != "empty": - seed = self.serde.loads_typed(seed_blob) + seed: Any = DELTA_SENTINEL + if seed_version is not None: + blob = seed_blob_by_ver.get((ch, seed_version)) + if blob is not None and blob[0] != "empty": + seed = self.serde.loads_typed(blob) - collected.reverse() - return _ChannelWritesHistory(seed=seed, writes=collected) + collected.reverse() + result[ch] = _ChannelWritesHistory(seed=seed, writes=collected) + 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 891f7c95e..6268372ef 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -125,7 +125,8 @@ class CheckpointTuple(NamedTuple): class _ChannelWritesHistory(NamedTuple): - """Result of `BaseCheckpointSaver._get_channel_writes_history`. + """Result of `BaseCheckpointSaver._get_all_delta_channels_writes_history` + (a per-channel entry from the returned mapping). Storage-level view of what one channel wrote across the ancestor chain of a target checkpoint: @@ -487,12 +488,12 @@ class BaseCheckpointSaver(Generic[V]): raise NotImplementedError def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: - """Pure storage read used by `_get_channel_writes_history`. + """Pure storage read used by `_get_all_delta_channels_writes_history`. Must return the same value as `get_tuple` but must NOT trigger channel reconstruction; otherwise the channel-hydration path would re-enter - `_get_channel_writes_history`. Override only if `get_tuple` itself - performs channel hydration. + `_get_all_delta_channels_writes_history`. Override only if `get_tuple` + itself performs channel hydration. """ return self.get_tuple(config) @@ -500,14 +501,17 @@ class BaseCheckpointSaver(Generic[V]): """Async version of `_get_tuple_raw`. See docstring there.""" return await self.aget_tuple(config) - def _get_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: - """**Experimental.** Query one channel's writes along the parent chain. + def _get_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + """**Experimental.** Query multiple delta channels' writes along the parent chain. - Storage-level query, not channel semantics: returns `(seed, writes)` - reflecting what storage knows about a single channel across the - ancestor chain of the target checkpoint identified by `config`. + Storage-level query, not channel semantics: returns a per-channel + `(seed, writes)` reflecting what storage knows about each channel + across the ancestor chain of the target checkpoint identified by + `config`. + + For every channel in `channels`: * `writes` — on-path deltas oldest→newest as `PendingWrite` tuples. Writes stored at the target `checkpoint_id` itself are pending @@ -520,69 +524,86 @@ class BaseCheckpointSaver(Generic[V]): Walks the **parent chain** (not `list(before=...)`): for forked threads, only on-path ancestors contribute. - Reference implementation walks `get_tuple` + `parent_config`, - inspecting each ancestor's `channel_values[channel]` for the seed - terminator. Savers with direct storage access (`InMemorySaver`, - `PostgresSaver`) override for performance; the return contract is - fixed here. + Reference implementation walks `get_tuple` + `parent_config` ONCE + for all channels (each ancestor visited once, not once per channel), + inspecting each ancestor's `channel_values[channel]` for that + channel's seed terminator. Savers with direct storage access + (`InMemorySaver`, `PostgresSaver`) override for performance; the + return contract is fixed here. - Underscore-prefixed because the method surface is experimental. + Empty `channels` returns `{}`. Underscore-prefixed because the + method surface is experimental. """ - collected: list[PendingWrite] = [] # newest first; reversed at the end + if not channels: + return {} + collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels} + seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels} + remaining: set[str] = set(channels) target_tuple = self._get_tuple_raw(config) cursor_config: RunnableConfig | None = ( target_tuple.parent_config if target_tuple else None ) - while cursor_config is not None: + while cursor_config is not None and remaining: tup = self._get_tuple_raw(cursor_config) if tup is None: break - # Collect this ancestor's writes FIRST — they encode the - # transition from this ancestor's state to its child's, so - # they must be included whether or not this ancestor is the - # seed terminator. + # Collect each ancestor's writes for any channel still searching. if tup.pending_writes: - # Within a superstep, pending_writes are oldest→newest; - # reverse to scan newest-first. for write in reversed(tup.pending_writes): - if write[1] != channel: - continue - collected.append(write) - # Seed terminator: any non-sentinel blob on an ancestor - # establishes the reconstruction base. Stop here. - ancestor_value = tup.checkpoint["channel_values"].get(channel) - if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: - collected.reverse() - return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + ch = write[1] + if ch in remaining: + collected_by_ch[ch].append(write) + # Per-channel seed terminator: a non-sentinel blob value at this + # ancestor establishes that channel's reconstruction base. + for ch in list(remaining): + ancestor_value = tup.checkpoint["channel_values"].get(ch) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + seed_by_ch[ch] = ancestor_value + remaining.discard(ch) cursor_config = tup.parent_config - collected.reverse() - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + return { + ch: _ChannelWritesHistory( + seed=seed_by_ch[ch], + writes=list(reversed(collected_by_ch[ch])), + ) + for ch in channels + } - async def _aget_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: - """Async version of `_get_channel_writes_history`. See docstring there.""" - collected: list[PendingWrite] = [] + async def _aget_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + """Async version of `_get_all_delta_channels_writes_history`.""" + if not channels: + return {} + collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels} + seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels} + remaining: set[str] = set(channels) target_tuple = await self._aget_tuple_raw(config) cursor_config: RunnableConfig | None = ( target_tuple.parent_config if target_tuple else None ) - while cursor_config is not None: + while cursor_config is not None and remaining: tup = await self._aget_tuple_raw(cursor_config) if tup is None: break if tup.pending_writes: for write in reversed(tup.pending_writes): - if write[1] != channel: - continue - collected.append(write) - ancestor_value = tup.checkpoint["channel_values"].get(channel) - if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: - collected.reverse() - return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + ch = write[1] + if ch in remaining: + collected_by_ch[ch].append(write) + for ch in list(remaining): + ancestor_value = tup.checkpoint["channel_values"].get(ch) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + seed_by_ch[ch] = ancestor_value + remaining.discard(ch) cursor_config = tup.parent_config - collected.reverse() - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + return { + ch: _ChannelWritesHistory( + seed=seed_by_ch[ch], + writes=list(reversed(collected_by_ch[ch])), + ) + for ch in channels + } def get_next_version(self, current: V | None, channel: None) -> V: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index a84277e1c..11f35b0a5 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -6,7 +6,7 @@ import pickle import random import shutil from collections import defaultdict -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack from types import TracebackType from typing import Any @@ -141,17 +141,24 @@ class InMemorySaver( result[k] = self.serde.loads_typed(vv) return result - def _get_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: + def _get_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + """Override: walk the parent chain ONCE for all requested channels. + + For each channel we track its own seed terminator independently. + On a snapshot or pre-delta ancestor for a given channel, that + channel stops collecting further writes; other channels keep + walking until they find their own terminator or hit the root. + """ + if not channels: + return {} 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. 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`). + + # Build the parent chain (newest→oldest), skipping the target. chain: list[str] = [] target_entry = ns_storage.get(checkpoint_id) current: str | None = target_entry[2] if target_entry is not None else None @@ -162,77 +169,75 @@ class InMemorySaver( chain.append(current) _, _, parent = entry current = parent - # Scan newest→oldest. A pre-delta blob on an ancestor terminates the - # walk and is bound as `seed`; without this, 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[PendingWrite] = [] # 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: - if isinstance(blob_value, _DeltaSnapshot): - # Step-based snapshot: the blob is state AT this - # ancestor, but the ancestor's pending_writes - # encode the NEXT step's transition and are NOT - # subsumed by the snapshot — collect them first. - step_writes = self.writes.get( - (thread_id, checkpoint_ns, cp_id), {} - ) - for (_task_id, _idx), ( - tid, - ch, - serialized, - _, - ) in sorted(step_writes.items(), reverse=True): - if ch != channel: - continue - collected.append( - (tid, ch, self.serde.loads_typed(serialized)) - ) - collected.reverse() - return _ChannelWritesHistory( - seed=blob_value, writes=collected - ) - # Pre-delta blob: state AT this ancestor already - # subsumes its pending_writes — skip them. - collected.reverse() - return _ChannelWritesHistory( - seed=blob_value, writes=collected - ) + collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels} + seed_by_ch: dict[str, Any] = {c: DELTA_SENTINEL for c in channels} + remaining: set[str] = set(channels) + + for cp_id in chain: # newest → oldest + if not remaining: + break + entry = ns_storage.get(cp_id) + ckpt = self.serde.loads_typed(entry[0]) if entry is not None else None + + # Per-channel: check seed terminator at this ancestor first. + terminated_here: set[str] = set() + blob_value_by_ch: dict[str, Any] = {} + if ckpt is not None: + versions = ckpt.get("channel_versions", {}) + for ch in remaining: + ver = versions.get(ch) + if ver is None: + continue + blob_entry = self.blobs.get( + (thread_id, checkpoint_ns, ch, ver) + ) + if blob_entry is None or blob_entry[0] == "empty": + continue + blob_value = self.serde.loads_typed(blob_entry) + if blob_value is DELTA_SENTINEL: + continue + blob_value_by_ch[ch] = blob_value + terminated_here.add(ch) + + # Process step writes: filter by channel, collect newest-first. step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) - # Within a superstep, sorted by (task_id, idx) = oldest → newest; - # reverse for newest-first scan. for (_task_id, _idx), (tid, ch, serialized, _) in sorted( step_writes.items(), reverse=True ): - if ch != channel: + if ch not in remaining: continue - val = self.serde.loads_typed(serialized) - collected.append((tid, ch, val)) - collected.reverse() - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + blob_value = blob_value_by_ch.get(ch) + if blob_value is not None and not isinstance( + blob_value, _DeltaSnapshot + ): + # Pre-delta blob terminator: state at this ancestor + # already subsumes these writes — skip them. + continue + # Either no terminator at this ancestor for this channel, + # OR a `_DeltaSnapshot` terminator (writes here encode the + # transition to the child and are NOT subsumed). + collected_by_ch[ch].append( + (tid, ch, self.serde.loads_typed(serialized)) + ) - async def _aget_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: - return self._get_channel_writes_history(config, channel) + # Now apply terminators: channels that found a seed are done. + for ch in terminated_here: + seed_by_ch[ch] = blob_value_by_ch[ch] + remaining.discard(ch) + + return { + ch: _ChannelWritesHistory( + seed=seed_by_ch[ch], + writes=list(reversed(collected_by_ch[ch])), + ) + for ch in channels + } + + async def _aget_all_delta_channels_writes_history( + self, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, _ChannelWritesHistory]: + return self._get_all_delta_channels_writes_history(config, channels) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the in-memory storage. diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 7e9cff219..6fdfd59d2 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -36,7 +36,7 @@ class _DeltaSnapshot(NamedTuple): """Snapshot blob for a DeltaChannel with finite snapshot_frequency. Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code. - The ancestor walk in `_get_channel_writes_history` terminates when it + The ancestor walk in `_get_all_delta_channels_writes_history` terminates when it encounters this type (any non-sentinel blob stops the walk). `from_checkpoint` reconstructs the channel value directly from `.value` diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index a4b6cf5b9..369ca65cd 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -335,9 +335,10 @@ class TestInMemorySaverDeltaChannel: assert channel not in result def test_get_channel_writes_collects_ancestor_writes_only(self) -> None: - """_get_channel_writes_history 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).""" + """_get_all_delta_channels_writes_history 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() @@ -375,7 +376,9 @@ class TestInMemorySaverDeltaChannel: "checkpoint_id": "cp2", } } - result = saver._get_channel_writes_history(config, channel) + result = saver._get_all_delta_channels_writes_history(config, [channel])[ + channel + ] assert result.seed is DELTA_SENTINEL values = [v for _, _, v in result.writes] assert values == [{"content": "hi"}] @@ -405,15 +408,17 @@ class TestInMemorySaverDeltaChannel: "checkpoint_id": "cp1", } } - result = saver._get_channel_writes_history(config, channel) + result = saver._get_all_delta_channels_writes_history(config, [channel])[ + channel + ] assert result.seed is DELTA_SENTINEL assert result.writes == [] class TestBaseFallbackGetChannelWrites: - """Exercises the `BaseCheckpointSaver._get_channel_writes_history` default - implementation — the path third-party savers inherit when they don't - override `_get_channel_writes_history` themselves. + """Exercises the `BaseCheckpointSaver._get_all_delta_channels_writes_history` + default implementation — the path third-party savers inherit when they + don't override `_get_all_delta_channels_writes_history` themselves. Regression guard for a bug where the fallback passed the caller's config (with `checkpoint_id`) straight to `self.list()`, which most savers @@ -429,11 +434,11 @@ class TestBaseFallbackGetChannelWrites: """ class _ThirdPartyStyleSaver(InMemorySaver): - _get_channel_writes_history = ( - InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] + _get_all_delta_channels_writes_history = ( + InMemorySaver.__mro__[1]._get_all_delta_channels_writes_history # type: ignore[attr-defined] ) - _aget_channel_writes_history = ( - InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + _aget_all_delta_channels_writes_history = ( + InMemorySaver.__mro__[1]._aget_all_delta_channels_writes_history # type: ignore[attr-defined] ) saver = _ThirdPartyStyleSaver() @@ -477,7 +482,9 @@ class TestBaseFallbackGetChannelWrites: } } - result = saver._get_channel_writes_history(config, "messages") + result = saver._get_all_delta_channels_writes_history( + config, ["messages"] + )["messages"] assert result.seed is DELTA_SENTINEL values = [v for _, _, v in result.writes] @@ -494,7 +501,9 @@ class TestBaseFallbackGetChannelWrites: } } - result = await saver._aget_channel_writes_history(config, "messages") + result = ( + await saver._aget_all_delta_channels_writes_history(config, ["messages"]) + )["messages"] assert result.seed is DELTA_SENTINEL values = [v for _, _, v in result.writes] @@ -503,9 +512,9 @@ class TestBaseFallbackGetChannelWrites: async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None: """Regression: the re-entrancy guard must be task-local, not thread-local. - Two concurrent `_aget_channel_writes_history` calls on the same - event-loop thread must each see their full reconstructed writes. A - `threading.local()` guard would let whichever task set it first + Two concurrent `_aget_all_delta_channels_writes_history` calls on the + same event-loop thread must each see their full reconstructed writes. + A `threading.local()` guard would let whichever task set it first short-circuit the other to `writes=[]`. """ import asyncio @@ -533,12 +542,13 @@ class TestBaseFallbackGetChannelWrites: } results = await asyncio.gather( - saver._aget_channel_writes_history(config, "messages"), - saver._aget_channel_writes_history(config, "messages"), + saver._aget_all_delta_channels_writes_history(config, ["messages"]), + saver._aget_all_delta_channels_writes_history(config, ["messages"]), ) expected_values = [{"content": "first"}, {"content": "second"}] - for result in results: + for result_map in results: + result = result_map["messages"] assert result.seed is DELTA_SENTINEL values = [v for _, _, v in result.writes] assert values == expected_values @@ -625,7 +635,9 @@ class TestPreDeltaBlobTerminator: } } - result = saver._get_channel_writes_history(config, channel) + result = saver._get_all_delta_channels_writes_history(config, [channel])[ + channel + ] # Seed came from the pre-delta blob at cp1. assert result.seed == ["A"] @@ -647,7 +659,9 @@ class TestPreDeltaBlobTerminator: } } - result = saver._get_channel_writes_history(config, channel) + result = saver._get_all_delta_channels_writes_history(config, [channel])[ + channel + ] values = [v for _, _, v in result.writes] # The pre-delta write under cp1 must not appear (the blob subsumes it). diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index dc7675975..715b5df1f 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -114,9 +114,12 @@ def channels_from_checkpoint( For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])` is sufficient. `DeltaChannel` is the exception: sentinel blobs require an - ancestor walk via `saver._get_channel_writes_history`. The walk terminates - at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain - value, so read depth is bounded by `snapshot_frequency`. + ancestor walk via `saver._get_all_delta_channels_writes_history`. All + delta channels needing replay are batched into a single saver call to + save K-1 redundant scans of `checkpoint_writes` (which has no channel + index). The walk terminates per-channel at the nearest `_DeltaSnapshot` + blob or pre-migration plain value, so read depth is bounded by + `snapshot_frequency`. """ channel_specs: dict[str, BaseChannel] = {} managed_specs: dict[str, ManagedValueSpec] = {} @@ -126,18 +129,28 @@ def channels_from_checkpoint( else: managed_specs[k] = v + delta_channels: list[str] = [ + k + for k, spec in channel_specs.items() + if _needs_replay(spec, checkpoint["channel_values"].get(k, MISSING)) + ] + histories: Mapping[str, Any] = {} + if delta_channels and saver is not None and config is not None: + histories = saver._get_all_delta_channels_writes_history( + config, delta_channels + ) + channels: dict[str, BaseChannel] = {} for k, spec in channel_specs.items(): ch: BaseChannel - stored = checkpoint["channel_values"].get(k, MISSING) - if _needs_replay(spec, stored) and saver is not None and config is not None: + if k in histories: delta_spec = cast(DeltaChannel, spec) - history = saver._get_channel_writes_history(config, k) + history = histories[k] replay_ch = delta_spec.from_checkpoint(history.seed) replay_ch.replay_writes(history.writes) ch = replay_ch else: - ch = spec.from_checkpoint(stored) + ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) channels[k] = ch return channels, managed_specs @@ -158,18 +171,28 @@ async def achannels_from_checkpoint( else: managed_specs[k] = v + delta_channels: list[str] = [ + k + for k, spec in channel_specs.items() + if _needs_replay(spec, checkpoint["channel_values"].get(k, MISSING)) + ] + histories: Mapping[str, Any] = {} + if delta_channels and saver is not None and config is not None: + histories = await saver._aget_all_delta_channels_writes_history( + config, delta_channels + ) + channels: dict[str, BaseChannel] = {} for k, spec in channel_specs.items(): ch: BaseChannel - stored = checkpoint["channel_values"].get(k, MISSING) - if _needs_replay(spec, stored) and saver is not None and config is not None: + if k in histories: delta_spec = cast(DeltaChannel, spec) - history = await saver._aget_channel_writes_history(config, k) + history = histories[k] replay_ch = delta_spec.from_checkpoint(history.seed) replay_ch.replay_writes(history.writes) ch = replay_ch else: - ch = spec.from_checkpoint(stored) + ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) channels[k] = ch return channels, managed_specs diff --git a/libs/langgraph/tests/test_delta_channel_migration.py b/libs/langgraph/tests/test_delta_channel_migration.py index 245ba2561..8d9577d56 100644 --- a/libs/langgraph/tests/test_delta_channel_migration.py +++ b/libs/langgraph/tests/test_delta_channel_migration.py @@ -6,7 +6,7 @@ checkpointer — pre-migration state visible at each *settled* ancestor checkpoint is preserved, and post-migration writes fold on top through the reducer. -Mechanism under test: the saver's `_get_channel_writes_history(config, +Mechanism under test: the saver's `_get_all_delta_channels_writes_history(config, channel)` walks the parent chain; when it encounters an ancestor whose `channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses @@ -29,7 +29,7 @@ Scenarios covered: pre-migration seed. 4. **Base-saver fallback path**: a third-party-style subclass that removes the optimized `InMemorySaver` override and falls back to - `BaseCheckpointSaver._get_channel_writes_history` must produce the + `BaseCheckpointSaver._get_all_delta_channels_writes_history` must produce the same result as the optimized path. 5. **Channel-type isolation across threads**: two threads on the same checkpointer under the delta-channel graph — one freshly-started, @@ -266,7 +266,7 @@ def test_continuing_migrated_thread_folds_deltas_on_seed() -> None: class _ThirdPartyStyleSaver(InMemorySaver): """Simulates a third-party saver that inherits the reference - `_get_channel_writes_history` implementation from + `_get_all_delta_channels_writes_history` implementation from `BaseCheckpointSaver` rather than overriding it. We rebind the two methods to the base-class versions (via MRO) so @@ -275,11 +275,11 @@ class _ThirdPartyStyleSaver(InMemorySaver): """ # MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...] - _get_channel_writes_history = ( # type: ignore[assignment] - InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] + _get_all_delta_channels_writes_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1]._get_all_delta_channels_writes_history # type: ignore[attr-defined] ) - _aget_channel_writes_history = ( # type: ignore[assignment] - InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + _aget_all_delta_channels_writes_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1]._aget_all_delta_channels_writes_history # type: ignore[attr-defined] ) @@ -326,7 +326,7 @@ def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None: """Two threads sharing a checkpointer — one migrated from pre-migration state, one freshly-started under DeltaChannel — must maintain independent state. The parent-chain walk in - `_get_channel_writes_history` must be scoped to the target thread. + `_get_all_delta_channels_writes_history` must be scoped to the target thread. """ checkpointer = InMemorySaver()