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