diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index f3f692641..18186d89b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -2,19 +2,18 @@ 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 from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( - DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, - _ChannelWritesHistory, + DeltaChannelHistory, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -27,10 +26,10 @@ 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, + _DELTA_PAGE_SIZE, BasePostgresSaver, - _DeltaStage1Row, + _build_delta_stage1_sql, + _build_delta_stage2_sql, _DeltaStage2Row, ) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver @@ -311,9 +310,7 @@ class PostgresSaver(BasePostgresSaver): # others are stored in blobs table blob_values = {} for k, v in checkpoint["channel_values"].items(): - if v is DELTA_SENTINEL: - copy["channel_values"].pop(k) - elif isinstance(v, _DeltaSnapshot): + if isinstance(v, _DeltaSnapshot): blob_values[k] = copy["channel_values"].pop(k) copy["channel_values"][k] = True elif v is None or isinstance(v, (str, int, float, bool)): @@ -444,53 +441,111 @@ 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_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Fast-path override of `BaseCheckpointSaver.get_delta_channel_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: + + * 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. + + * 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. """ + 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: {"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) + 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() + cursor: str | None = None + with self._cursor() as cur: - cur.execute( - SELECT_DELTA_STAGE1_SQL, - (channel, channel, thread_id, checkpoint_ns), - ) - stage1_rows = cur.fetchall() - chain_cids, seed_version = self._walk_stage1( - cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id + 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) + 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 + + # 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]] + channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None] + stage2_sql = _build_delta_stage2_sql( + channels_with_chain=channels_with_chain, + channels_with_seed=channels_with_seed, ) - seed_versions = [seed_version] if seed_version else [] - with self._cursor() as cur: - cur.execute( - SELECT_DELTA_STAGE2_SQL, - ( - thread_id, - checkpoint_ns, - channel, - chain_cids, - thread_id, - checkpoint_ns, - channel, - seed_versions, - ), - ) - stage2_rows = cur.fetchall() - return self._build_delta_channel_writes_history( - channel=channel, - chain_cids=chain_cids, - seed_version=seed_version, + + if stage2_sql: + stage2_params: list[Any] = [] + for ch in channels_with_chain: + stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]]) + for ch in channels_with_seed: + stage2_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() + else: + stage2_rows = [] + + 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..09fb964d5 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -2,19 +2,18 @@ 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 from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( - DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, - _ChannelWritesHistory, + DeltaChannelHistory, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -27,10 +26,10 @@ 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, + _DELTA_PAGE_SIZE, BasePostgresSaver, - _DeltaStage1Row, + _build_delta_stage1_sql, + _build_delta_stage2_sql, _DeltaStage2Row, ) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver @@ -270,9 +269,7 @@ class AsyncPostgresSaver(BasePostgresSaver): # others are stored in blobs table blob_values = {} for k, v in checkpoint["channel_values"].items(): - if v is DELTA_SENTINEL: - copy["channel_values"].pop(k) - elif isinstance(v, _DeltaSnapshot): + if isinstance(v, _DeltaSnapshot): blob_values[k] = copy["channel_values"].pop(k) copy["channel_values"][k] = True elif v is None or isinstance(v, (str, int, float, bool)): @@ -405,53 +402,94 @@ 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_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_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. + 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. """ + 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: {"writes": []} for ch in channels} checkpoint_id = target.config["configurable"]["checkpoint_id"] + stage1_sql = _build_delta_stage1_sql(channels, paged=True) + 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() + cursor: str | None = None + async with self._cursor() as cur: - await cur.execute( - SELECT_DELTA_STAGE1_SQL, - (channel, channel, thread_id, checkpoint_ns), - ) - stage1_rows = await cur.fetchall() - chain_cids, seed_version = self._walk_stage1( - cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id + 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) + 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 + + channels_with_chain = [ch for ch in channels if chain_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( + channels_with_chain=channels_with_chain, + channels_with_seed=channels_with_seed, ) - seed_versions = [seed_version] if seed_version else [] - async with self._cursor() as cur: - await cur.execute( - SELECT_DELTA_STAGE2_SQL, - ( - thread_id, - checkpoint_ns, - channel, - chain_cids, - thread_id, - checkpoint_ns, - channel, - seed_versions, - ), - ) - stage2_rows = await cur.fetchall() - return self._build_delta_channel_writes_history( - channel=channel, - chain_cids=chain_cids, - seed_version=seed_version, + + if stage2_sql: + stage2_params: list[Any] = [] + for ch in channels_with_chain: + stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]]) + for ch in channels_with_seed: + stage2_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() + else: + stage2_rows = [] + + 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..bbdff8cf8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -2,23 +2,26 @@ 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 from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( - DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, + DeltaChannelHistory, PendingWrite, - _ChannelWritesHistory, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS from psycopg.types.json import Jsonb +# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal +# constant — exposing this as a kwarg is left as a follow-up. +_DELTA_PAGE_SIZE = 1024 + MetadataInput = dict[str, Any] | None try: @@ -157,10 +160,11 @@ INSERT_CHECKPOINT_WRITES_SQL = """ class _DeltaStage2Row(TypedDict, total=False): - """One row from `SELECT_DELTA_STAGE2_SQL` (a UNION ALL of writes and blobs).""" + """One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs).""" _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 +172,126 @@ 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[]) - -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, - type, blob, task_id, idx, NULL::text AS version - FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s - AND checkpoint_id = ANY(%s) - UNION ALL - SELECT 'b', NULL, - type, blob, NULL, NULL, version - FROM checkpoint_blobs - WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s - AND version = ANY(%s) -""" +# 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). -class _DeltaStage1Row(TypedDict): - """One row from `SELECT_DELTA_STAGE1_SQL`.""" +def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: + """Build stage 1 SQL with 2K parallel JSONB key lookups. - checkpoint_id: str - parent_checkpoint_id: str | None - ver: str | None - has_snapshot: bool + For channels=["messages", "files"] (with `paged=True`) 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 + AND (%s::text IS NULL OR checkpoint_id < %s) + ORDER BY checkpoint_id DESC + 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). + + 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. + """ + 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 = ( + "SELECT checkpoint_id, parent_checkpoint_id, " + + ", ".join(cols) + + " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %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( + *, + channels_with_chain: Sequence[str], + channels_with_seed: Sequence[str], +) -> str: + """Build stage 2 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 + `channel = ANY(channels) AND checkpoint_id = ANY(union)` form when + channels have different chain depths. + + The caller must pass parameters in matching order: + + for ch in channels_with_chain: + params += [thread_id, checkpoint_ns, ch, chain_cids[ch]] + for ch in channels_with_seed: + params += [thread_id, checkpoint_ns, ch, seed_version[ch]] + + Returns an empty SQL string if both channel lists are empty (caller + must skip executing in that case). + """ + branches: list[str] = [] + for _ in channels_with_chain: + branches.append( + "SELECT 'w'::text AS _kind, " + "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 " + "AND checkpoint_id = ANY(%s)" + ) + for _ in channels_with_seed: + branches.append( + "SELECT 'b'::text, NULL, channel, " + "type, blob, NULL, NULL, version " + "FROM checkpoint_blobs " + "WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s " + "AND version = %s" + ) + 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. class BasePostgresSaver(BaseCheckpointSaver[str]): @@ -255,86 +337,144 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): } @staticmethod - def _walk_stage1( - stage1_rows: Sequence[_DeltaStage1Row], - target_id: str, - ) -> tuple[list[str], str | None]: - """Walk the parent chain from stage 1 metadata rows. + def _ingest_stage1_page( + stage1_rows: Sequence[Mapping[str, Any]], + 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 (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. + 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`). """ - parent_of: dict[str, str | None] = {} - ver_of: dict[str, str | None] = {} - snapshot_of: dict[str, bool] = {} + oldest: str | None = None 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}")) + # Rows are DESC; the last one is the smallest cid in the page. + oldest = cid + return oldest - 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 + @staticmethod + def _try_advance_walks( + target_id: str, + channels: Sequence[str], + parent_of: Mapping[str, str | None], + 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. - def _build_delta_channel_writes_history( + 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). + + Mutates `chain_by_ch`, `seed_ver_by_ch`, `walk_cursor_by_ch`, and + `seeded` in place. + """ + for i, ch in enumerate(channels): + if ch in seeded: + 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 + + def _build_delta_channels_writes_history( self, *, - channel: str, - chain_cids: list[str], - seed_version: str | None, + channels: Sequence[str], + chain_by_ch: Mapping[str, list[str]], + seed_ver_by_ch: Mapping[str, str | None], stage2_rows: Sequence[_DeltaStage2Row], - ) -> _ChannelWritesHistory: - """Reconstruct delta channel history from two-stage query results. + ) -> dict[str, DeltaChannelHistory]: + """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 `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". """ - 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, DeltaChannelHistory] = {} + 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)) + collected.reverse() - seed: Any = DELTA_SENTINEL - if seed_blob is not None and seed_blob[0] != "empty": - seed = self.serde.loads_typed(seed_blob) - - collected.reverse() - return _ChannelWritesHistory(seed=seed, writes=collected) + 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) + result[ch] = entry + return result def _dump_blobs( self, diff --git a/libs/checkpoint-postgres/uv.lock b/libs/checkpoint-postgres/uv.lock index a508d02d8..aa0a31e00 100644 --- a/libs/checkpoint-postgres/uv.lock +++ b/libs/checkpoint-postgres/uv.lock @@ -244,10 +244,11 @@ wheels = [ [[package]] name = "langchain-core" -version = "1.2.28" +version = "1.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jsonpatch" }, + { name = "langchain-protocol" }, { name = "langsmith" }, { name = "packaging" }, { name = "pydantic" }, @@ -256,9 +257,21 @@ dependencies = [ { name = "typing-extensions" }, { name = "uuid-utils" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f8/a4/317a1a3ac1df33a64adb3670bf88bbe3b3d5baa274db6863a979db472897/langchain_core-1.2.28.tar.gz", hash = "sha256:271a3d8bd618f795fdeba112b0753980457fc90537c46a0c11998516a74dc2cb", size = 846119, upload-time = "2026-04-08T18:19:34.867Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/92/32f785f077c7e898da97064f113c73fbd9ad55d1e2169cf3a391b183dedb/langchain_core-1.2.28-py3-none-any.whl", hash = "sha256:80764232581eaf8057bcefa71dbf8adc1f6a28d257ebd8b95ba9b8b452e8c6ac", size = 508727, upload-time = "2026-04-08T18:19:32.823Z" }, + { url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" }, +] + +[[package]] +name = "langchain-protocol" +version = "0.0.15" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4f/24/9777489d6fbbee64af0c8f96d4f840239c408cf694f3394672807dafc490/langchain_protocol-0.0.15.tar.gz", hash = "sha256:9ab2d11ee73944754f10e037e717098d3a6796f0e58afa9cadda6154e7655ade", size = 5862, upload-time = "2026-05-01T22:30:04.748Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/7a/9c97a7b9cbe4c5dc6a44cdb1545450c28f0c8ce89b9c1f0ee7fbad896263/langchain_protocol-0.0.15-py3-none-any.whl", hash = "sha256:461eb794358f83d5e42635a5797799ffec7b4702314e34edf73ac21e75d3ef79", size = 6982, upload-time = "2026-05-01T22:30:03.877Z" }, ] [[package]] diff --git a/libs/checkpoint-sqlite/tests/test_get_delta_channel_history.py b/libs/checkpoint-sqlite/tests/test_get_delta_channel_history.py new file mode 100644 index 000000000..8e19c8b85 --- /dev/null +++ b/libs/checkpoint-sqlite/tests/test_get_delta_channel_history.py @@ -0,0 +1,271 @@ +"""Smoke tests for `BaseCheckpointSaver.get_delta_channel_history` on sqlite. + +`SqliteSaver` (and `AsyncSqliteSaver`) deliberately don't override the +default implementation in `BaseCheckpointSaver` — these tests pin the +default impl to behave correctly end-to-end against a real persistent +saver and a real `DeltaChannel`-backed graph. + +Scenarios covered: + +* Empty `channels` argument returns an empty mapping (no I/O). +* On a non-trivial multi-checkpoint thread, per-channel writes come back + oldest→newest. +* When the walk reaches the root without ever finding a stored value, + `seed` is omitted from the entry (consumer treats absence as "start + empty"). +* When a `_DeltaSnapshot` blob is present at an ancestor, it is returned + as the `seed`. +* The async saver returns the same shape via `aget_delta_channel_history`. +""" + +from __future__ import annotations + +import operator +from typing import Annotated, Any + +import pytest +from langchain_core.runnables import RunnableConfig + +# `langgraph` is not a dep of `langgraph-checkpoint-sqlite`. When tests run +# in the sqlite lib's standalone CI environment without it installed, skip +# the whole module rather than failing at import. +pytest.importorskip("langgraph.channels.delta", reason="langgraph core not installed") +pytest.importorskip("langgraph.graph", reason="langgraph core not installed") + +from langgraph.channels.delta import DeltaChannel # type: ignore[import-untyped] # noqa: E402,I001 +from langgraph.checkpoint.serde.types import _DeltaSnapshot # noqa: E402 +from langgraph.graph import END, START, StateGraph # type: ignore[import-untyped] # noqa: E402 +from typing_extensions import TypedDict # noqa: E402 + +from langgraph.checkpoint.sqlite import SqliteSaver # noqa: E402 +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver # noqa: E402 + +pytestmark = pytest.mark.anyio + + +# --------------------------------------------------------------------------- +# Graph helpers +# --------------------------------------------------------------------------- + + +def _noop(_state: Any) -> dict[str, Any]: + return {} + + +class _DeltaState(TypedDict): + items: Annotated[list, DeltaChannel(operator.add)] + + +def _delta_graph(checkpointer: Any) -> Any: + return ( + StateGraph(_DeltaState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _drive(graph: Any, config: RunnableConfig, n: int) -> None: + for i in range(n): + graph.invoke({"items": [f"v{i}"]}, config) + + +async def _adrive(graph: Any, config: RunnableConfig, n: int) -> None: + for i in range(n): + await graph.ainvoke({"items": [f"v{i}"]}, config) + + +def _pick_non_root(saver: Any, config: RunnableConfig) -> RunnableConfig: + """Return a config pointing at a checkpoint that has at least one ancestor. + + `get_delta_channel_history` walks the parent chain — calling it on the root + checkpoint produces `writes=[]` and no `seed`, which is uninteresting + for the multi-step assertions below. + """ + history = list(saver.list(config)) + assert history, "expected non-empty history" + # `list` yields newest→oldest; the second entry has the first entry + # as its parent, so its parent_config is non-None. + for tup in history: + if tup.parent_config is not None: + return tup.config + raise AssertionError("no checkpoint with a parent in history") + + +async def _apick_non_root(saver: Any, config: RunnableConfig) -> RunnableConfig: + history = [tup async for tup in saver.alist(config)] + assert history, "expected non-empty history" + for tup in history: + if tup.parent_config is not None: + return tup.config + raise AssertionError("no checkpoint with a parent in history") + + +# --------------------------------------------------------------------------- +# Sync: SqliteSaver +# --------------------------------------------------------------------------- + + +def test_empty_channels_returns_empty_mapping_sync() -> None: + """Empty `channels` short-circuits to `{}` without touching storage.""" + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "empty"}} + assert saver.get_delta_channel_history(config=config, channels=[]) == {} + + +def test_writes_history_oldest_to_newest_sync() -> None: + """Per-channel writes accumulated across the walk come back oldest→newest.""" + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "history-sync"}} + graph = _delta_graph(saver) + _drive(graph, config, 3) + + target_cfg = _pick_non_root(saver, config) + result = saver.get_delta_channel_history(config=target_cfg, channels=["items"]) + + assert "items" in result + entry = result["items"] + assert isinstance(entry["writes"], list) + + # If any writes were collected, their values should be in oldest→newest + # order — i.e. tagged 'v0', 'v1', ... matching invoke order. + write_values: list[Any] = [] + for _task_id, channel, value in entry["writes"]: + assert channel == "items" + write_values.extend(value if isinstance(value, list) else [value]) + + # `_drive` invokes with payloads ['v0'], ['v1'], ['v2']. Whatever + # subset shows up in the chain must be a contiguous prefix in order. + for idx, val in enumerate(write_values): + assert val == f"v{idx}", ( + f"writes not in oldest→newest order: {write_values}" + ) + + +def test_seed_present_when_snapshot_in_ancestor_sync() -> None: + """Inserting a `_DeltaSnapshot` blob at an ancestor → walk returns it as `seed`.""" + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "seed-sync"}} + graph = _delta_graph(saver) + _drive(graph, config, 2) + + # Find the oldest non-root checkpoint, then walk to its parent and + # rewrite that parent's `channel_values["items"]` to a real + # `_DeltaSnapshot`. After this surgery, calling `get_delta_channel_history` + # at the leaf must return the snapshot value as `seed`. + history = list(saver.list(config)) + assert len(history) >= 2 + leaf_tup = history[0] + # Walk to an ancestor with a parent_config (any non-root will do). + ancestor_tup = next( + (tup for tup in history if tup.parent_config is not None), None + ) + assert ancestor_tup is not None + parent_cfg = ancestor_tup.parent_config + assert parent_cfg is not None + parent_tup = saver.get_tuple(parent_cfg) + assert parent_tup is not None + + snapshot_value = ["seeded", "items"] + parent_tup.checkpoint["channel_values"]["items"] = _DeltaSnapshot( + snapshot_value + ) + # Make sure the channel has a version so the optimized blob lookup + # in any future override has something to hit. + parent_tup.checkpoint["channel_versions"].setdefault("items", 1) + saver.put( + parent_tup.parent_config or {"configurable": parent_cfg["configurable"]}, + parent_tup.checkpoint, + parent_tup.metadata, + {}, + ) + + result = saver.get_delta_channel_history( + config=leaf_tup.config, channels=["items"] + ) + entry = result["items"] + assert "seed" in entry, f"expected seed to be present, got {entry}" + seed = entry["seed"] + assert isinstance(seed, _DeltaSnapshot), ( + f"expected _DeltaSnapshot, got {seed!r}" + ) + assert seed.value == snapshot_value + + +def test_seed_omitted_when_walk_reaches_root_sync() -> None: + """`get_delta_channel_history` on the root checkpoint → no `seed` key, no writes.""" + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "root-sync"}} + graph = _delta_graph(saver) + _drive(graph, config, 1) + + history = list(saver.list(config)) + # Root is the oldest checkpoint (no parent_config). + root_tup = history[-1] + assert root_tup.parent_config is None + + result = saver.get_delta_channel_history( + config=root_tup.config, channels=["items"] + ) + entry = result["items"] + assert "seed" not in entry, f"root-walk should have no seed, got {entry}" + assert entry["writes"] == [] + + +# --------------------------------------------------------------------------- +# Async: AsyncSqliteSaver +# --------------------------------------------------------------------------- + + +async def test_empty_channels_returns_empty_mapping_async() -> None: + """Async equivalent of the empty-channels short-circuit.""" + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "empty-async"}} + assert await saver.aget_delta_channel_history(config=config, channels=[]) == {} + + +async def test_writes_history_oldest_to_newest_async() -> None: + """Async equivalent of the oldest→newest ordering check.""" + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "history-async"}} + graph = _delta_graph(saver) + await _adrive(graph, config, 3) + + target_cfg = await _apick_non_root(saver, config) + result = await saver.aget_delta_channel_history( + config=target_cfg, channels=["items"] + ) + + assert "items" in result + entry = result["items"] + assert isinstance(entry["writes"], list) + + write_values: list[Any] = [] + for _task_id, channel, value in entry["writes"]: + assert channel == "items" + write_values.extend(value if isinstance(value, list) else [value]) + + for idx, val in enumerate(write_values): + assert val == f"v{idx}", ( + f"writes not in oldest→newest order: {write_values}" + ) + + +async def test_seed_omitted_when_walk_reaches_root_async() -> None: + """Async equivalent of the root-walk seed-absence check.""" + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "root-async"}} + graph = _delta_graph(saver) + await _adrive(graph, config, 1) + + history = [tup async for tup in saver.alist(config)] + root_tup = history[-1] + assert root_tup.parent_config is None + + result = await saver.aget_delta_channel_history( + config=root_tup.config, channels=["items"] + ) + entry = result["items"] + assert "seed" not in entry, f"root-walk should have no seed, got {entry}" + assert entry["writes"] == [] diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 891f7c95e..b34509f00 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -13,14 +13,12 @@ from typing import ( ) from langchain_core.runnables import RunnableConfig +from typing_extensions import NotRequired from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.serde.types import ( - DELTA_SENTINEL as DELTA_SENTINEL, -) from langgraph.checkpoint.serde.types import ( ERROR, INTERRUPT, @@ -62,6 +60,16 @@ class CheckpointMetadata(TypedDict, total=False): """ run_id: str """The ID of the run that created this checkpoint.""" + delta_updates_since_snapshot: dict[str, int] + """Per-channel update count since the last `_DeltaSnapshot` was written. + + Maps channel name → number of supersteps that wrote to this channel + since its last snapshot blob. Used by `pregel.create_checkpoint` to + decide when to write the next snapshot (when the count reaches the + channel's `snapshot_frequency`, snapshot fires and the count resets + to 0). Absent on threads that don't use delta channels. Version-format + independent — works for int, float, and string version schemes. + """ ChannelVersions = dict[str, str | int | float] @@ -124,28 +132,26 @@ class CheckpointTuple(NamedTuple): pending_writes: list[PendingWrite] | None = None -class _ChannelWritesHistory(NamedTuple): - """Result of `BaseCheckpointSaver._get_channel_writes_history`. +class DeltaChannelHistory(TypedDict): + """Per-channel result entry from `BaseCheckpointSaver.get_delta_channel_history`. - Storage-level view of what one channel wrote across the ancestor chain - of a target checkpoint: + Storage-level view of what one channel contributed across the ancestor + chain of a target checkpoint: - * `seed` — the nearest ancestor's stored blob value for this channel, - or `DELTA_SENTINEL` if the walk reached the root without finding a - stored value. A non-sentinel seed typically indicates a pre-delta - snapshot preserved across a channel-type migration (e.g. - `BinaryOperatorAggregate` storage extended under `DeltaChannel`). - * `writes` — on-path deltas oldest→newest, one `PendingWrite` per - step that wrote to this channel. Writes stored at the target - checkpoint itself are pending for the next super-step and are - excluded. - - Experimental: method surface may change; the NamedTuple shape is the - contract. + * `writes` — on-path deltas oldest→newest as `PendingWrite` tuples. + Always present; possibly empty. Already filtered to one channel. + Writes stored at the target checkpoint itself are pending for the + next super-step and are excluded. + * `seed` — the stored value at the nearest ancestor whose + `channel_values[ch]` is populated. Omitted if the walk reached the + root without finding any stored value (consumer treats absence as + "start empty"). Typically a `_DeltaSnapshot` for delta channels with + finite snapshot frequency, or a plain value for threads migrated + from a pre-delta channel type. """ - seed: Any writes: list[PendingWrite] + seed: NotRequired[Any] class BaseCheckpointSaver(Generic[V]): @@ -486,103 +492,101 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: - """Pure storage read used by `_get_channel_writes_history`. + def get_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Walk the parent chain returning per-channel writes + seed. - 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. - """ - return self.get_tuple(config) - - async def _aget_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: - """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. - - 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`. - - * `writes` — on-path deltas oldest→newest as `PendingWrite` tuples. - Writes stored at the target `checkpoint_id` itself are pending - for the next super-step and are excluded. - * `seed` — the nearest ancestor's stored blob value for this - channel; `DELTA_SENTINEL` if the walk reached the root without - finding a stored value. A non-sentinel seed typically indicates - a pre-delta snapshot preserved across a channel-type migration. + For each requested channel, walks ancestors of the checkpoint + identified by `config` (following `parent_config`) and accumulates + `pending_writes` for that channel. The walk terminates per-channel + at the nearest ancestor whose `channel_values[ch]` is populated; + that value is returned as `seed`. If the walk reaches the root + without finding a stored value, `seed` is omitted from that + channel's entry — the consumer treats the absence as "start + empty." 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`, + The default implementation walks `get_tuple` + `parent_config` + once for all channels — each ancestor visited once, not once per + channel. Savers with direct storage access (`InMemorySaver`, `PostgresSaver`) override for performance; the return contract is fixed here. - Underscore-prefixed because the method surface is experimental. - """ - collected: list[PendingWrite] = [] # newest first; reversed at the end - 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: - 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. - 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) - cursor_config = tup.parent_config - collected.reverse() - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + Args: + config: Configuration identifying the target checkpoint. + channels: Channel names to walk for. Empty → empty mapping. - 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] = [] - target_tuple = await self._aget_tuple_raw(config) + Returns: + Per-channel `DeltaChannelHistory` for every name in `channels`. + """ + 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) + target_tuple = self.get_tuple(config) cursor_config: RunnableConfig | None = ( target_tuple.parent_config if target_tuple else None ) - while cursor_config is not None: - tup = await self._aget_tuple_raw(cursor_config) + while cursor_config is not None and remaining: + tup = self.get_tuple(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): + if ch in tup.checkpoint["channel_values"]: + seed_by_ch[ch] = tup.checkpoint["channel_values"][ch] + remaining.discard(ch) cursor_config = tup.parent_config - collected.reverse() - return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + 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 + + async def aget_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Async version of `get_delta_channel_history`.""" + 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) + target_tuple = await self.aget_tuple(config) + cursor_config: RunnableConfig | None = ( + target_tuple.parent_config if target_tuple else None + ) + while cursor_config is not None and remaining: + 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) + 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 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..80043c710 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 @@ -14,20 +14,18 @@ from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( - DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, + DeltaChannelHistory, PendingWrite, SerializerProtocol, - _ChannelWritesHistory, get_checkpoint_id, get_checkpoint_metadata, ) -from langgraph.checkpoint.serde.types import _DeltaSnapshot logger = logging.getLogger(__name__) @@ -141,17 +139,31 @@ class InMemorySaver( result[k] = self.serde.loads_typed(vv) return result - def _get_channel_writes_history( - self, config: RunnableConfig, channel: str - ) -> _ChannelWritesHistory: + def get_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> 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. + + 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). + """ + if not channels: + return {} + # Imported lazily to avoid a hard checkpoint→serde-types coupling at + # module import; only this override needs the runtime check. + from langgraph.checkpoint.serde.types import _DeltaSnapshot + 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`). + 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 +174,64 @@ 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 + + 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: + break 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 - ) + 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), {}) - # 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 + ): + continue + 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) + for ch in terminated_here: + seed_by_ch[ch] = blob_value_by_ch[ch] + remaining.discard(ch) + + 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] + result[ch] = entry_h + return result + + async def aget_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + return self.get_delta_channel_history(config=config, channels=channels) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the in-memory storage. @@ -452,9 +451,7 @@ class InMemorySaver( values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc] for k, v in new_versions.items(): self.blobs[(thread_id, checkpoint_ns, k, v)] = ( - self.serde.dumps_typed(values[k]) - if k in values and values[k] is not DELTA_SENTINEL - else ("empty", b"") + self.serde.dumps_typed(values[k]) if k in values else ("empty", b"") ) self.storage[thread_id][checkpoint_ns].update( { diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 7e9cff219..72e372064 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -16,28 +16,13 @@ RESUME = "__resume__" TASKS = "__pregel_tasks" -class _DeltaSentinel: - """In-memory marker for a DeltaChannel field with no snapshot. - - Never serialized to storage — checkpointers strip it before writing. - Compare with `is DELTA_SENTINEL`; always the same module-level instance. - """ - - __slots__ = () - - def __repr__(self) -> str: - return "DELTA_SENTINEL" - - -DELTA_SENTINEL = _DeltaSentinel() - - 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 - encounters this type (any non-sentinel blob stops the walk). + The ancestor walk in `BaseCheckpointSaver.get_delta_channel_history` terminates + when it encounters this type (any non-empty channel_values entry stops + the walk for that channel). `from_checkpoint` reconstructs the channel value directly from `.value` without replaying writes — the snapshot IS the accumulated state. diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index a4b6cf5b9..70e22e0d8 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -6,7 +6,6 @@ from langchain_core.runnables import RunnableConfig from pydantic import BaseModel from langgraph.checkpoint.base import ( - DELTA_SENTINEL, Checkpoint, CheckpointMetadata, create_checkpoint, @@ -335,9 +334,14 @@ 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_delta_channel_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). + + When the walk reaches the root without finding any stored value for + the channel, the per-channel entry has no `seed` key (TypedDict + absence indicates "start empty"). + """ saver = InMemorySaver() serde = JsonPlusSerializer() @@ -375,9 +379,12 @@ class TestInMemorySaverDeltaChannel: "checkpoint_id": "cp2", } } - result = saver._get_channel_writes_history(config, channel) - assert result.seed is DELTA_SENTINEL - values = [v for _, _, v in result.writes] + result = saver.get_delta_channel_history(config=config, channels=[channel])[ + channel + ] + # Walk reached the root without finding a stored value → no `seed` key. + assert "seed" not in result + values = [v for _, _, v in result["writes"]] assert values == [{"content": "hi"}] def test_get_channel_writes_at_root_returns_empty(self) -> None: @@ -405,15 +412,18 @@ class TestInMemorySaverDeltaChannel: "checkpoint_id": "cp1", } } - result = saver._get_channel_writes_history(config, channel) - assert result.seed is DELTA_SENTINEL - assert result.writes == [] + result = saver.get_delta_channel_history(config=config, channels=[channel])[ + channel + ] + # No ancestors → no seed found, no writes accumulated. + assert "seed" not in result + assert result["writes"] == [] class TestBaseFallbackGetChannelWrites: - """Exercises the `BaseCheckpointSaver._get_channel_writes_history` default + """Exercises the `BaseCheckpointSaver.get_delta_channel_history` default implementation — the path third-party savers inherit when they don't - override `_get_channel_writes_history` themselves. + override `get_delta_channel_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 +439,11 @@ class TestBaseFallbackGetChannelWrites: """ class _ThirdPartyStyleSaver(InMemorySaver): - _get_channel_writes_history = ( - InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] + get_delta_channel_history = ( + InMemorySaver.__mro__[1].get_delta_channel_history # type: ignore[attr-defined] ) - _aget_channel_writes_history = ( - InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + aget_delta_channel_history = ( + InMemorySaver.__mro__[1].aget_delta_channel_history # type: ignore[attr-defined] ) saver = _ThirdPartyStyleSaver() @@ -477,10 +487,13 @@ class TestBaseFallbackGetChannelWrites: } } - result = saver._get_channel_writes_history(config, "messages") + result = saver.get_delta_channel_history(config=config, channels=["messages"])[ + "messages" + ] - assert result.seed is DELTA_SENTINEL - values = [v for _, _, v in result.writes] + # Walk reached root without a stored value → `seed` key absent. + assert "seed" not in result + values = [v for _, _, v in result["writes"]] assert values == [{"content": "first"}, {"content": "second"}] async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None: @@ -494,17 +507,19 @@ class TestBaseFallbackGetChannelWrites: } } - result = await saver._aget_channel_writes_history(config, "messages") + result = ( + await saver.aget_delta_channel_history(config=config, channels=["messages"]) + )["messages"] - assert result.seed is DELTA_SENTINEL - values = [v for _, _, v in result.writes] + assert "seed" not in result + values = [v for _, _, v in result["writes"]] assert values == [{"content": "first"}, {"content": "second"}] 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 + Two concurrent `aget_delta_channel_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=[]`. """ @@ -533,27 +548,32 @@ class TestBaseFallbackGetChannelWrites: } results = await asyncio.gather( - saver._aget_channel_writes_history(config, "messages"), - saver._aget_channel_writes_history(config, "messages"), + saver.aget_delta_channel_history(config=config, channels=["messages"]), + saver.aget_delta_channel_history(config=config, channels=["messages"]), ) expected_values = [{"content": "first"}, {"content": "second"}] - for result in results: - assert result.seed is DELTA_SENTINEL - values = [v for _, _, v in result.writes] + for result_map in results: + result = result_map["messages"] + assert "seed" not in result + values = [v for _, _, v in result["writes"]] assert values == expected_values class TestPreDeltaBlobTerminator: """Verify the pre-delta blob terminator: when the ancestor walk hits a - checkpoint whose blob for the channel is a real value (not - DELTA_SENTINEL), reconstruction seeds from it and stops. This guards + checkpoint whose blob for the channel is a real value, reconstruction + seeds from it and stops. This guards * back-compat: a thread written by pre-delta code, then extended under delta — reconstruction must return the correct value without walking past the last pre-delta ancestor; * perf: without the terminator, every reconstruct-after-migration would walk all the way to the thread root. + + Under the new public API a found seed populates `seed` in the + `DeltaChannelHistory` TypedDict; absence of the `seed` key means the walk + reached root without finding a stored value. """ def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]: @@ -625,14 +645,16 @@ class TestPreDeltaBlobTerminator: } } - result = saver._get_channel_writes_history(config, channel) + result = saver.get_delta_channel_history(config=config, channels=[channel])[ + channel + ] # Seed came from the pre-delta blob at cp1. - assert result.seed == ["A"] + assert result["seed"] == ["A"] # Delta-era writes from cp2 replay through the reducer on top of seed. # cp3 is the target — its own write is pending for the NEXT step and # must be excluded. - values = [v for _, _, v in result.writes] + values = [v for _, _, v in result["writes"]] assert values == ["B"] def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None: @@ -647,9 +669,11 @@ class TestPreDeltaBlobTerminator: } } - result = saver._get_channel_writes_history(config, channel) + result = saver.get_delta_channel_history(config=config, channels=[channel])[ + channel + ] - values = [v for _, _, v in result.writes] + values = [v for _, _, v in result["writes"]] # The pre-delta write under cp1 must not appear (the blob subsumes it). assert "PRE-DELTA-WRITE" not in values # And the pending write at the target is never folded in. diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index ed08f083f..ac1962e97 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -5,7 +5,7 @@ import copy as _copy from collections.abc import Callable, Sequence from typing import Any, Generic -from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite +from langgraph.checkpoint.base import PendingWrite from langgraph.checkpoint.serde.types import _DeltaSnapshot from typing_extensions import Self @@ -38,19 +38,17 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): This lets LangGraph replay checkpointed writes in larger batches than they were originally produced without changing reconstructed state. - `snapshot_frequency=None` (default): pure delta; stores only - `DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes. - - `snapshot_frequency=N`: `create_checkpoint` writes a full `_DeltaSnapshot` - blob every N steps, bounding replay depth to N. + Snapshot cadence is driven by per-channel update count. `create_checkpoint` + writes a full `_DeltaSnapshot` blob every `snapshot_frequency` updates to + this channel, bounding replay depth. Parameters: reducer: `(state, list[writes]) -> new_state`. Must be deterministic and batching-invariant as described above. typ: The value type (e.g. `list`, `dict`). Inferred automatically from the outer type when used inside `Annotated[T, DeltaChannel(...)]`. - snapshot_frequency: Every Nth pregel step writes a snapshot blob. - `None` (default) = pure delta, never snapshot. + snapshot_frequency: Every Nth update to this channel writes a snapshot + blob (default `1000`). Must be a positive int. """ __slots__ = ("value", "reducer", "snapshot_frequency") @@ -61,8 +59,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): reducer: Callable[[Any, Sequence[Any]], Any], typ: type[Value] | None = None, *, - snapshot_frequency: int | None = None, + snapshot_frequency: int = 1000, ) -> None: + if snapshot_frequency <= 0: + raise ValueError( + f"snapshot_frequency must be a positive int, got {snapshot_frequency}" + ) if typ is None: typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel super().__init__(typ) @@ -93,14 +95,6 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): def UpdateType(self) -> Any: return self.typ - def is_snapshot_step(self, step: int) -> bool: - """True if pregel should write a snapshot blob at this step.""" - return ( - self.snapshot_frequency is not None - and step > 0 - and step % self.snapshot_frequency == 0 - ) - def copy(self) -> Self: new = self.__class__( self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency @@ -110,18 +104,19 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return new def from_checkpoint(self, checkpoint: Any) -> Self: - """Initialize from a stored blob or sentinel. + """Initialize from a stored blob. - Blob types (dispatched via serde ext code, not dict key inspection): - * `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes. + Blob types: + * `MISSING`: start empty; caller replays writes. * `_DeltaSnapshot(value)`: restore value directly from snapshot. - * plain value (migration from old BinOp blobs): use directly. + * plain value (migration from old `BinaryOperatorAggregate` blobs): + use directly. """ new = self.__class__( self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency ) new.key = self.key - if checkpoint is MISSING or checkpoint is DELTA_SENTINEL: + if checkpoint is MISSING: new.value = self.typ() elif isinstance(checkpoint, _DeltaSnapshot): new.value = checkpoint.value @@ -186,12 +181,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return self.value is not MISSING def checkpoint(self) -> Any: - """Return stored representation: always `DELTA_SENTINEL`. + """Return stored representation: always `MISSING`. - Snapshot decisions are made by `create_checkpoint` in pregel (which - has the step number) via `is_snapshot_step`. `checkpoint()` is only - called for non-snapshot steps or when no checkpointer is available. + Snapshot decisions live in `create_checkpoint` (which has the channel + version) and write `_DeltaSnapshot(ch.get())` directly into + `channel_values`. For non-snapshot steps the channel does not appear + in `channel_values`; reconstruction walks ancestor writes via the + saver's `get_delta_channel_history`. """ - if self.value is MISSING: - return MISSING - return DELTA_SENTINEL + return MISSING diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index dc7675975..4e3bbffd6 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -5,7 +5,11 @@ from datetime import datetime, timezone from typing import Any, cast from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + Checkpoint, + CheckpointMetadata, +) from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.types import _DeltaSnapshot @@ -30,6 +34,30 @@ def empty_checkpoint() -> Checkpoint: ) +def _should_snapshot_delta( + name: str, + ch: DeltaChannel, + updates_since_snapshot: Mapping[str, int], + *, + force: bool, +) -> bool: + """Decide whether `ch` should write a `_DeltaSnapshot` this step. + + Triggers: + * `force` — always snapshot (used by `durability="exit"`). + * Update-count: this channel has accumulated at least + `snapshot_frequency` updates since its last snapshot. The count + is supplied by the caller via `updates_since_snapshot[name]` and + is reset to `0` whenever a snapshot fires. + + Version-format-independent: works for `int`, `float`, and `str` + versioning schemes alike. + """ + if force: + return True + return updates_since_snapshot.get(name, 0) >= ch.snapshot_frequency + + def create_checkpoint( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel] | None, @@ -39,21 +67,33 @@ def create_checkpoint( updated_channels: set[str] | None = None, get_next_version: GetNextVersion | None = None, force_delta_snapshot: bool = False, + updates_since_snapshot: Mapping[str, int] | None = None, + new_updates_since_snapshot: dict[str, int] | None = None, ) -> Checkpoint: """Create a checkpoint for the given channels. - For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a - `_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor - walk to at most N steps. Snapshots are eager: even if the channel had no - write this step, a version bump is forced (via `get_next_version`) so the - blob is stored by `put()`. Without `get_next_version` (e.g. static - contexts), snapshot steps gracefully fall back to sentinel. + For each `DeltaChannel`, a `_DeltaSnapshot(value)` blob is written into + `channel_values[k]` when this channel has accumulated at least + `snapshot_frequency` updates since its last snapshot (counter supplied + via `updates_since_snapshot`). Otherwise the channel is omitted from + `channel_values`; its `channel_versions` entry still bumps so that the + saver tracks the channel and the ancestor walk can replay writes. - `force_delta_snapshot` writes available `DeltaChannel` values as snapshots - regardless of `snapshot_frequency`. This is used by `durability="exit"`, - where intermediate writes are not stored as ancestor `checkpoint_writes`. + Snapshots are eager: even if the channel had no write this step, a + version bump is forced (via `get_next_version`) so `put()` includes + the channel in `new_versions` and stores the blob. + + `force_delta_snapshot` ignores the cadence and always snapshots — + used by `durability="exit"` where intermediate writes are not stored + as ancestor `checkpoint_writes`. + + If `new_updates_since_snapshot` is provided, the function resets the + counter to `0` for any channel that snapshotted this step. Counters + for channels that did not snapshot are left untouched (the caller is + responsible for incrementing them based on `updated_channels`). """ ts = datetime.now(timezone.utc).isoformat() + counts = updates_since_snapshot or {} if channels is None: values = checkpoint["channel_values"] channel_versions = checkpoint["channel_versions"] @@ -66,8 +106,13 @@ def create_checkpoint( ch = channels[k] if ( isinstance(ch, DeltaChannel) - and (force_delta_snapshot or ch.is_snapshot_step(step)) and ch.is_available() + and _should_snapshot_delta( + k, + ch, + counts, + force=force_delta_snapshot, + ) ): # Eager snapshot: bump version if not already written this step # so put() includes this channel in new_versions and stores blob. @@ -76,6 +121,8 @@ def create_checkpoint( ): channel_versions[k] = get_next_version(channel_versions[k], None) values[k] = _DeltaSnapshot(ch.get()) + if new_updates_since_snapshot is not None: + new_updates_since_snapshot[k] = 0 else: v = ch.checkpoint() if v is not MISSING: @@ -92,15 +139,15 @@ def create_checkpoint( def _needs_replay(spec: BaseChannel, stored: object) -> bool: - """True if `spec` is a `DeltaChannel` and the stored blob is a sentinel, - requiring an ancestor walk to reconstruct. + """True if `spec` is a `DeltaChannel` and no value is stored at this + checkpoint, requiring an ancestor walk to reconstruct. `_DeltaSnapshot` blobs and plain values (migration) resolve directly via - `from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay. + `from_checkpoint` — only absence (`MISSING`) triggers replay. """ if not isinstance(spec, DeltaChannel): return False - return stored is MISSING or stored is DELTA_SENTINEL + return stored is MISSING def channels_from_checkpoint( @@ -113,10 +160,12 @@ def channels_from_checkpoint( """Hydrate channels from a 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`. + is sufficient. `DeltaChannel` is the exception: when the channel is + absent from `channel_values`, an ancestor walk via + `saver.get_delta_channel_history` is required to find the nearest seed + (`_DeltaSnapshot` blob or pre-migration plain value) and accumulate + the writes between it and the target. All delta channels needing + replay are batched into a single saver call. """ channel_specs: dict[str, BaseChannel] = {} managed_specs: dict[str, ManagedValueSpec] = {} @@ -126,18 +175,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_delta_channel_history( + config=config, channels=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) - replay_ch = delta_spec.from_checkpoint(history.seed) - replay_ch.replay_writes(history.writes) + history = histories[k] + replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING)) + 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 +217,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_delta_channel_history( + config=config, channels=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) - replay_ch = delta_spec.from_checkpoint(history.seed) - replay_ch.replay_writes(history.writes) + history = histories[k] + replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING)) + 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 @@ -184,3 +253,17 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, updated_channels=checkpoint.get("updated_channels", None), ) + + +def read_delta_updates_since_snapshot( + metadata: CheckpointMetadata | None, +) -> dict[str, int]: + """Read the per-channel update counter from checkpoint metadata. + + Returns an empty dict for missing/None metadata; the dict is + `total=False` on `CheckpointMetadata`, so absence means "no prior + delta-channel activity tracked." + """ + if not metadata: + return {} + return dict(metadata.get("delta_updates_since_snapshot", {}) or {}) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index aeb8d52da..ad772016a 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -910,6 +910,22 @@ class PregelLoop: if exiting and self.checkpoint["id"] == self.checkpoint_id_saved: # checkpoint already saved return + # Carry per-delta-channel update bookkeeping forward across + # supersteps. Capture from the OLD metadata before potentially + # replacing it with a fresh dict that wouldn't contain it. Then + # increment for any delta channel updated this step (so the count + # reflects "supersteps that wrote to this channel since last + # snapshot"). create_checkpoint will reset entries to 0 for any + # channel that fires a snapshot this step. + prev_counts = dict( + self.checkpoint_metadata.get("delta_updates_since_snapshot", {}) or {} + ) + new_counts = dict(prev_counts) + if self.updated_channels: + for ch_name in self.updated_channels: + ch_obj = self.channels.get(ch_name) + if isinstance(ch_obj, DeltaChannel): + new_counts[ch_name] = new_counts.get(ch_name, 0) + 1 if not exiting: metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) @@ -929,7 +945,13 @@ class PregelLoop: if do_checkpoint else None, force_delta_snapshot=exiting and self.durability == "exit", + updates_since_snapshot=new_counts, + new_updates_since_snapshot=new_counts, ) + if new_counts: + self.checkpoint_metadata["delta_updates_since_snapshot"] = new_counts + elif "delta_updates_since_snapshot" in self.checkpoint_metadata: + del self.checkpoint_metadata["delta_updates_since_snapshot"] # sanitize TASK channel in the checkpoint before saving (durability=="exit") if TASKS in self.checkpoint["channel_values"] and any( isinstance(channel, UntrackedValue) for channel in self.channels.values() @@ -1450,7 +1472,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): new_versions: ChannelVersions, ) -> RunnableConfig: # Drain DeltaChannel write futures before committing the checkpoint so - # DELTA_SENTINEL blobs are never saved ahead of their backing writes. + # ancestor walks never see a checkpoint without its backing writes. if self._delta_write_futs: futs, self._delta_write_futs = self._delta_write_futs, [] await asyncio.gather(*futs) diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 6c15f3307..9dfa7158a 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -4,7 +4,6 @@ from typing import Annotated import pytest from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage -from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph.checkpoint.memory import InMemorySaver from langgraph.checkpoint.serde.types import _DeltaSnapshot from typing_extensions import NotRequired, TypedDict @@ -140,11 +139,11 @@ def test_delta_channel_basic_two_steps() -> None: ch.update([HumanMessage(content="hi", id="h1")]) d1 = ch.checkpoint() - assert d1 is DELTA_SENTINEL + assert d1 is MISSING ch.update([AIMessage(content="hello", id="a1")]) d2 = ch.checkpoint() - assert d2 is DELTA_SENTINEL + assert d2 is MISSING assert len(ch.get()) == 2 assert ch.get()[0].content == "hi" @@ -154,7 +153,7 @@ def test_delta_channel_basic_two_steps() -> None: def test_delta_channel_from_checkpoint_writes_list() -> None: """replay_writes on a fresh channel replays through the operator.""" spec = DeltaChannel(_messages_delta_reducer, list) - ch = spec.from_checkpoint(DELTA_SENTINEL) + ch = spec.from_checkpoint(MISSING) ch.replay_writes( [ ("t0", "messages", HumanMessage(content="hi", id="h1")), @@ -182,7 +181,7 @@ def test_delta_channel_overwrite() -> None: ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) d = ch.checkpoint() - assert d is DELTA_SENTINEL + assert d is MISSING assert len(ch.get()) == 1 assert ch.get()[0].content == "new" @@ -202,7 +201,7 @@ def test_delta_channel_remove_message_and_replay() -> None: ch.update([RemoveMessage(id="a1")]) assert ch.get() == [HumanMessage(content="hi", id="h1")] - ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2 = spec.from_checkpoint(MISSING) ch2.replay_writes( [ ("t0", "messages", HumanMessage(content="hi", id="h1")), @@ -222,7 +221,7 @@ def test_delta_channel_update_by_id_and_replay() -> None: ch.update([HumanMessage(content="updated", id="h1")]) assert ch.get() == [HumanMessage(content="updated", id="h1")] - ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2 = spec.from_checkpoint(MISSING) ch2.replay_writes( [ ("t0", "messages", HumanMessage(content="original", id="h1")), @@ -285,13 +284,18 @@ def test_messages_delta_reducer_tuple_write_is_one_message() -> None: assert result[0].content == "hi" -def test_delta_channel_checkpoint_returns_sentinel() -> None: - """checkpoint() always returns DELTA_SENTINEL regardless of state.""" +def test_delta_channel_checkpoint_returns_missing() -> None: + """checkpoint() always returns MISSING regardless of state. + + Pregel writes `_DeltaSnapshot(ch.get())` directly into `channel_values` + on snapshot steps; the channel itself never participates in snapshot + serialization, so its `checkpoint()` is always the absence sentinel. + """ ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) - assert ch.checkpoint() is DELTA_SENTINEL + assert ch.checkpoint() is MISSING ch.update([HumanMessage(content="hi", id="h1")]) - assert ch.checkpoint() is DELTA_SENTINEL + assert ch.checkpoint() is MISSING # --------------------------------------------------------------------------- @@ -299,12 +303,13 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None: # --------------------------------------------------------------------------- -def test_delta_channel_snapshot_step_based() -> None: - """Snapshots fire on every Nth step regardless of whether the channel was written. +def test_delta_channel_snapshot_version_based() -> None: + """Snapshots fire when a channel accumulates `snapshot_frequency` updates. - With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot - blob — even if the channel had no write that step (eager snapshot). This - bounds the ancestor walk to at most N steps on any read. + Under the version-delta cadence, every time the channel's + `current_version - last_snapshot_version >= snapshot_frequency` a + `_DeltaSnapshot` blob is written. Bounds the ancestor walk to at most + `snapshot_frequency` steps on any read for that channel. """ class State(TypedDict): @@ -347,51 +352,14 @@ def test_delta_channel_snapshot_step_based() -> None: assert len(state.values["messages"]) == 12 # 6 human + 6 AI -def test_delta_channel_snapshot_fires_even_when_not_written() -> None: - """Eager snapshot: _DeltaSnapshot stored at snapshot step even when the - channel had no write that step (node_b doesn't touch messages). - """ - - class State(TypedDict): - messages: Annotated[ - list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3) - ] - tick: int - - def writer(state: State) -> dict: - i = len(state["messages"]) // 2 - return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]} - - def ticker(state: State) -> dict: - return {"tick": state["tick"] + 1} - - g = StateGraph(State) - g.add_node("writer", writer) - g.add_node("ticker", ticker) - g.add_edge(START, "writer") - g.add_edge("writer", "ticker") - saver = InMemorySaver() - graph = g.compile(checkpointer=saver) - - config = {"configurable": {"thread_id": "t1"}} - for i in range(5): - graph.invoke( - {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0}, - config, - ) - - msg_blobs = { - k: saver.serde.loads_typed((t, b)) - for k, (t, b) in saver.blobs.items() - if k[2] == "messages" and t == "msgpack" and b - } - snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)} - assert snapshots, ( - "eager snapshots must fire even on steps where messages wasn't written" - ) - - state = graph.get_state(config) - assert len(state.values["messages"]) == 10 # 5 human + 5 AI +# TODO(delta-channel-cadence): the previous "snapshot fires even when channel +# was not written" test asserted eager step-based snapshotting; under the new +# version-delta cadence (`should_snapshot` triggers on per-channel update +# count, not superstep count), no snapshot fires for an unwritten channel. +# Replace with a test that exercises the version-delta trigger plus the +# durability="exit" force-snapshot branch — see +# `docs/superpowers/specs/2026-05-04-delta-channel-batched-reads-design.md` +# section "Snapshot cadence". # --------------------------------------------------------------------------- @@ -466,11 +434,11 @@ def test_delta_channel_dict_reducer_basic_updates() -> None: ch.update([{"a": 1}]) d1 = ch.checkpoint() - assert d1 is DELTA_SENTINEL + assert d1 is MISSING ch.update([{"b": 2}]) d2 = ch.checkpoint() - assert d2 is DELTA_SENTINEL + assert d2 is MISSING assert ch.get() == {"a": 1, "b": 2} @@ -485,7 +453,7 @@ def test_delta_channel_dict_reducer_writes_reconstruction() -> None: return result spec = _delta_channel_with_type(merge_dicts, dict) - ch = spec.from_checkpoint(DELTA_SENTINEL) + ch = spec.from_checkpoint(MISSING) ch.replay_writes( [ ("t0", "files", {"a": 1}), @@ -515,7 +483,7 @@ def test_delta_channel_dict_reducer_with_deletions() -> None: assert ch.get() == {"file2.py": "content2", "file3.py": "content3"} spec = _delta_channel_with_type(merge_files, dict) - ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2 = spec.from_checkpoint(MISSING) ch2.replay_writes( [ ("t0", "files", {"file1.py": "content1", "file2.py": "content2"}), @@ -550,7 +518,7 @@ def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None: return result spec = _delta_channel_with_type(merge_dicts, dict) - ch = spec.from_checkpoint(DELTA_SENTINEL) + ch = spec.from_checkpoint(MISSING) ch.replay_writes( [ ("t0", "files", {"a": 1}), @@ -691,7 +659,7 @@ def test_delta_channel_from_checkpoint_seed_without_writes() -> None: def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> None: """`seed=None` must start replay from None, not from an empty channel. - The DELTA_SENTINEL / MISSING sentinels mean 'no seed'; passing `None` + The `MISSING` absence sentinel means 'no seed'; passing `None` explicitly should feed None to the reducer as the left operand. """ diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index cf17d0f12..0e45ffdaf 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -1,35 +1,41 @@ -"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff. +"""Benchmark: DeltaChannel — multi-channel reads, mixed snapshot frequencies. Run directly: python tests/test_delta_channel_benchmark.py Run via pytest: pytest tests/test_delta_channel_benchmark.py -s -Part 1 — baseline (original): DeltaChannel(inf) vs add_messages (BinOp). -Part 2 — snapshot_frequency sweep: shows the storage/read-latency tradeoff - across frequencies [1, 5, 10, 50, inf] at scale. +Sweeps `(K delta channels, snapshot_frequency strategy, turn count)` and +reports per-scenario read latency, write latency, storage, and peak Python +heap usage during `get_state`. -Key insight: - snapshot_frequency=inf → O(N) storage, O(N) read depth (pure delta) - snapshot_frequency=N → O(N²/N) storage, O(N) read depth bounded by freq - snapshot_frequency=1 → O(N²) storage, O(1) read depth (full snapshot) +Scenarios cover the dimensions where this branch's optimizations matter: + + * K-channel batching — varying K (number of `DeltaChannel`s the graph + reads on hydrate) shows the effect of merging + per-channel reads into a single saver call. + * Mixed frequencies — channels with very different snapshot cadences + in one graph exercise the per-channel chain + bound in stage-2. + * Turn count — chain depth shows how paged stage-1 holds up. """ from __future__ import annotations import contextlib -import math +import gc import os import sys import time +import tracemalloc from typing import Annotated, Any import pytest -from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.messages import HumanMessage from langgraph.checkpoint.memory import MemorySaver from typing_extensions import TypedDict from langgraph.channels.delta import DeltaChannel from langgraph.graph import END, StateGraph -from langgraph.graph.message import _messages_delta_reducer, add_messages +from langgraph.graph.message import _messages_delta_reducer try: from langgraph.checkpoint.postgres import PostgresSaver @@ -42,24 +48,16 @@ try: except ImportError: _POSTGRES_AVAILABLE = False + # --------------------------------------------------------------------------- # Realistic message payload (~100 tokens / ~400 chars each) # --------------------------------------------------------------------------- _HUMAN_TEMPLATE = ( - "I need help understanding the implications of {topic} on our system architecture. " - "Specifically, I'm concerned about how this interacts with our existing {concern} " - "and whether we need to refactor the {component} layer before proceeding. " - "We've had prior incidents in this area and want to be deliberate. " - "What should we prioritize first, and are there known failure modes we should design around from the start?" -) - -_AI_TEMPLATE = ( - "Great question about {topic}. The key insight here is that {concern} introduces " - "a subtle ordering dependency that most teams overlook until they hit it in production. " - "For your {component} layer specifically, I'd recommend starting with a careful audit " - "of the interface boundaries before making any structural changes. This will give you " - "a clear picture of the blast radius and let you sequence the migration safely." + "I need help understanding the implications of {topic} on our system " + "architecture. Specifically, I'm concerned about how this interacts with " + "our existing {concern} and whether we need to refactor the {component} " + "layer before proceeding." ) _TOPICS = [ @@ -68,28 +66,9 @@ _TOPICS = [ "schema migration", "backpressure handling", "idempotency guarantees", - "cache invalidation", - "connection pooling", - "rate limiting", - "circuit breaking", - "observability pipelines", -] - -_CONCERNS = [ - "concurrency model", - "retry semantics", - "state management", - "error propagation", - "latency budget", -] - -_COMPONENTS = [ - "persistence", - "routing", - "ingestion", - "aggregation", - "serialization", ] +_CONCERNS = ["concurrency model", "retry semantics", "ordering guarantees"] +_COMPONENTS = ["persistence", "ingestion", "routing"] def _human_content(i: int) -> str: @@ -100,136 +79,132 @@ def _human_content(i: int) -> str: ) -def _ai_content(i: int) -> str: - return _AI_TEMPLATE.format( - topic=_TOPICS[i % len(_TOPICS)], - concern=_CONCERNS[i % len(_CONCERNS)], - component=_COMPONENTS[i % len(_COMPONENTS)], - ) - - # --------------------------------------------------------------------------- -# State definitions +# State / graph factory: K DeltaChannel fields with per-channel freqs # --------------------------------------------------------------------------- -class BinaryState(TypedDict): - messages: Annotated[list, add_messages] - - -class DeltaState(TypedDict): - messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] - - -def _make_delta_state(snapshot_frequency: int | float) -> type: - """Create a TypedDict with DeltaChannel at the given snapshot_frequency.""" - channel = DeltaChannel( - _messages_delta_reducer, snapshot_frequency=snapshot_frequency - ) - # Use the functional TypedDict form so the Annotated type is stored as an - # already-evaluated object rather than a forward-reference string (which - # would fail when get_type_hints tries to resolve 'snapshot_frequency'). +def _make_state_cls(freqs: list[int]) -> type: + """Build a TypedDict with one DeltaChannel per entry in `freqs`.""" + fields: dict[str, Any] = {} + for i, freq in enumerate(freqs): + ch = DeltaChannel(_messages_delta_reducer, snapshot_frequency=freq) + fields[f"ch{i}"] = Annotated[list, ch] return TypedDict( # type: ignore[return-value] - f"DeltaState_freq{snapshot_frequency}", - {"messages": Annotated[list, channel]}, + "_BenchState_" + "-".join(str(f) for f in freqs), + fields, ) -# --------------------------------------------------------------------------- -# Graph factory -# --------------------------------------------------------------------------- +def _make_graph(state_cls: type, K: int, checkpointer: Any = None) -> Any: + """Graph: one node, writes a fresh message into every channel each turn.""" - -def _make_graph(state_cls: type, checkpointer: Any = None) -> Any: - def human_node(state: Any) -> dict: - return {} - - def ai_node(state: Any) -> dict: - i = len(state["messages"]) // 2 - return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]} + def fanout(state: Any) -> dict[str, Any]: + i = max(len(state.get(f"ch{j}", [])) for j in range(K)) + # Each channel gets its own copy with a unique id so the reducer + # does meaningful per-channel state accumulation. + return { + f"ch{j}": [HumanMessage(content=_human_content(i), id=f"c{j}_{i}")] + for j in range(K) + } g = StateGraph(state_cls) - g.add_node("human", human_node) - g.add_node("ai", ai_node) - g.add_edge("human", "ai") - g.add_edge("ai", END) - g.set_entry_point("human") + g.add_node("fanout", fanout) + g.set_entry_point("fanout") + g.add_edge("fanout", END) return g.compile(checkpointer=checkpointer or MemorySaver()) # --------------------------------------------------------------------------- -# Measurement helpers +# Storage / memory measurement # --------------------------------------------------------------------------- -def _total_blob_bytes(saver: MemorySaver) -> int: - total = 0 - for (_, _, _, _), (type_tag, blob) in saver.blobs.items(): - if blob is not None: - total += len(blob) - return total +def _inmemory_blob_bytes(saver: MemorySaver) -> int: + return sum( + len(blob) for (_, _, _, _), (_, blob) in saver.blobs.items() if blob is not None + ) -def _run_turns( - n_turns: int, - state_cls: type, - checkpointer: Any = None, -) -> tuple[float, float, int]: - """Run n_turns conversation turns. - - Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes). - Read latency is the average of 5 get_state calls after the full history - is built — forces state rehydration including ancestor replay if needed. +def _postgres_storage_bytes(saver: Any, thread_id: str) -> int: + """Total bytes across checkpoints / checkpoint_blobs / checkpoint_writes + rows for this `thread_id`. Uses `pg_column_size` for an in-row payload + estimate; faster than full-table size and scoped to the thread.""" + sql = """ + SELECT COALESCE(SUM(pg_column_size(c.*)), 0) + + COALESCE((SELECT SUM(pg_column_size(b.*)) FROM checkpoint_blobs b + WHERE b.thread_id = %s), 0) + + COALESCE((SELECT SUM(pg_column_size(w.*)) FROM checkpoint_writes w + WHERE w.thread_id = %s), 0) + AS total + FROM checkpoints c + WHERE c.thread_id = %s """ - graph = _make_graph(state_cls, checkpointer) - config = {"configurable": {"thread_id": "bench"}} + with saver._cursor() as cur: + cur.execute(sql, (thread_id, thread_id, thread_id)) + row = cur.fetchone() + if row is None: + return 0 + if isinstance(row, dict): + return int(row.get("total") or 0) + return int(row[0] or 0) + +def _run_scenario( + freqs: list[int], + n_turns: int, + checkpointer: Any, + thread_id: str, +) -> dict[str, float | int]: + """Drive `n_turns` invocations of a K-channel graph, measure + write/read/storage/peak-memory.""" + K = len(freqs) + state_cls = _make_state_cls(freqs) + graph = _make_graph(state_cls, K, checkpointer) + config = {"configurable": {"thread_id": thread_id}} + + # Write phase t0 = time.perf_counter() for i in range(n_turns): - graph.invoke( - {"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]}, - config, - ) + graph.invoke({}, config) write_elapsed = time.perf_counter() - t0 + # Read phase + tracemalloc peak across get_state calls + gc.collect() + tracemalloc.start() t1 = time.perf_counter() for _ in range(5): graph.get_state(config) read_elapsed = (time.perf_counter() - t1) / 5 + _, peak_bytes = tracemalloc.get_traced_memory() + tracemalloc.stop() - blob_bytes = ( - _total_blob_bytes(graph.checkpointer) - if isinstance(graph.checkpointer, MemorySaver) - else -1 - ) - return write_elapsed, read_elapsed, blob_bytes + # Storage + if isinstance(graph.checkpointer, MemorySaver): + storage = _inmemory_blob_bytes(graph.checkpointer) + elif _POSTGRES_AVAILABLE and isinstance(graph.checkpointer, PostgresSaver): + storage = _postgres_storage_bytes(graph.checkpointer, thread_id) + else: + storage = -1 - -def _fmt_bytes(n: int) -> str: - if n >= 1_000_000: - return f"{n / 1_000_000:.1f} MB" - if n >= 1_000: - return f"{n / 1_000:.1f} KB" - return f"{n} B" - - -def _approx_tokens(n_turns: int) -> str: - tokens = n_turns * 200 - if tokens >= 1_000_000: - return f"~{tokens / 1_000_000:.1f}M tok" - if tokens >= 1_000: - return f"~{tokens / 1_000:.0f}K tok" - return f"~{tokens} tok" + return { + "K": K, + "turns": n_turns, + "write_total_s": write_elapsed, + "write_per_invoke_ms": (write_elapsed / n_turns) * 1000, + "read_avg_ms": read_elapsed * 1000, + "storage_bytes": storage, + "peak_mem_bytes": peak_bytes, + } # --------------------------------------------------------------------------- -# Checkpointer factories +# Postgres helpers # --------------------------------------------------------------------------- @contextlib.contextmanager -def _pg_saver(thread_id: str = "bench"): - """Context manager that yields a fresh PostgresSaver and cleans up after.""" +def _pg_saver(thread_id: str): with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver: saver.setup() with saver._cursor() as cur: @@ -242,7 +217,6 @@ def _pg_saver(thread_id: str = "bench"): def _checkpointers() -> list[tuple[str, Any]]: - """Return (label, saver_or_None) pairs for available checkpointers.""" result: list[tuple[str, Any]] = [("InMemory", None)] if _POSTGRES_AVAILABLE: try: @@ -256,223 +230,92 @@ def _checkpointers() -> list[tuple[str, Any]]: # --------------------------------------------------------------------------- -# Part 1: baseline DeltaChannel(inf) vs add_messages +# Scenarios # --------------------------------------------------------------------------- -BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500] -DELTA_ONLY_TURN_COUNTS = [1000] +SCENARIOS: list[tuple[str, list[int]]] = [ + ("K=1, freq=50", [50]), + ("K=3, freq=50 uniform", [50, 50, 50]), + ("K=3, freq=mixed", [50, 200, 1000]), + ("K=8, freq=50 uniform", [50] * 8), + ("K=8, freq=mixed", [25, 50, 100, 200, 500, 1000, 1000, 1000]), +] + +TURN_COUNTS = [100, 500] -def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None: - W = 72 +def _fmt_bytes(n: int) -> str: + if n < 0: + return "n/a" + if n >= 1_000_000: + return f"{n / 1_000_000:.1f} MB" + if n >= 1_000: + return f"{n / 1_000:.1f} KB" + return f"{n} B" - def _make_saver(): - if cp_hint is None: - return contextlib.nullcontext(None) - return _pg_saver() - rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = [] - for turns in BASELINE_TURN_COUNTS: - with _make_saver() as saver: - b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver) - with _make_saver() as saver: - d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) - rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt)) - for turns in DELTA_ONLY_TURN_COUNTS: - with _make_saver() as saver: - d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) - rows.append((turns, None, d_bytes, None, d_rt, None, d_wt)) - - def _bytes_or_na(v: Any) -> str: - if v is None or v < 0: - return "n/a" - return _fmt_bytes(v) - - def _ms_or_na(v: Any) -> str: - return "n/a" if v is None else f"{v * 1000:.1f}ms" - - print(f"\n [{cp_label}] Storage (blob bytes)") - print( - f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}" +def _print_scenario_table(cp_label: str, rows: list[dict]) -> None: + print(f"\n [{cp_label}]") + header = ( + f" {'scenario':<28}{'turns':>8}{'write_ms':>11}" + f"{'read_ms':>10}{'storage':>12}{'peak_mem':>12}" ) - print(" " + "-" * (W - 2)) - for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows: - if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0: - ratio_str = "n/a" - else: - ratio = b_bytes / d_bytes if d_bytes else float("inf") - ratio_str = f"{ratio:.0f}x" + print(header) + print(" " + "-" * (len(header) - 2)) + for row in rows: print( - f" {turns:>6} {_approx_tokens(turns):>10} " - f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} {ratio_str:>8}" - ) - - print(f"\n [{cp_label}] Read latency (avg of 5 get_state calls)") - print(f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}") - print(" " + "-" * (W - 2)) - for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows: - print( - f" {turns:>6} {_approx_tokens(turns):>10} " - f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}" + f" {row['scenario']:<28}" + f"{row['turns']:>8}" + f"{row['write_per_invoke_ms']:>10.1f}" + f"{row['read_avg_ms']:>10.2f}" + f"{_fmt_bytes(row['storage_bytes']):>12}" + f"{_fmt_bytes(row['peak_mem_bytes']):>12}" ) -def run_baseline_benchmark() -> None: - print() - print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency") - print("=" * 72) +def run_benchmark() -> list[dict]: + """Run the full sweep and return all measurement rows.""" + all_rows: list[dict] = [] for cp_label, cp_hint in _checkpointers(): - _run_baseline_for_checkpointer(cp_label, cp_hint) - print() + rows: list[dict] = [] + for scenario_label, freqs in SCENARIOS: + for turns in TURN_COUNTS: + thread_id = f"bench-{scenario_label.replace(' ', '_')}-{turns}".lower() + if cp_hint is None: + saver_ctx: Any = contextlib.nullcontext(None) + else: + saver_ctx = _pg_saver(thread_id) + with saver_ctx as saver: + measured = _run_scenario(freqs, turns, saver, thread_id) + measured["scenario"] = scenario_label + measured["saver"] = cp_label + rows.append(measured) + all_rows.append(measured) + _print_scenario_table(cp_label, rows) + return all_rows # --------------------------------------------------------------------------- -# Part 2: snapshot_frequency sweep -# --------------------------------------------------------------------------- - -# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta. -SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf] - -# Turn counts for the sweep — high enough to show storage divergence. -SWEEP_TURN_COUNTS = [50, 100, 500] - - -def _freq_label(freq: int | float) -> str: - if freq == math.inf: - return "inf" - return str(int(freq)) - - -def _run_sweep_for_checkpointer(cp_label: str, cp_hint: Any) -> None: - def _make_saver(): - if cp_hint is None: - return contextlib.nullcontext(None) - return _pg_saver() - - # Collect results: {turns: {freq_label: (write_s, read_s, bytes)}} - results: dict[int, dict[str, tuple[float, float, int]]] = {} - for turns in SWEEP_TURN_COUNTS: - results[turns] = {} - for freq in SNAPSHOT_FREQUENCIES: - state_cls = _make_delta_state(freq) - with _make_saver() as saver: - wt, rt, bb = _run_turns(turns, state_cls, saver) - results[turns][_freq_label(freq)] = (wt, rt, bb) - - freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES] - col_w = 12 - - header = f" {'turns':>6} {'ctx':>10}" + "".join( - f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels - ) - - print(f"\n [{cp_label}] Storage (blob bytes) — lower is better") - print(header) - print(" " + "-" * (len(header) - 2)) - for turns in SWEEP_TURN_COUNTS: - row = f" {turns:>6} {_approx_tokens(turns):>10}" - for label in freq_labels: - _, _, bb = results[turns][label] - row += f" {_fmt_bytes(bb) if bb >= 0 else 'n/a':>{col_w}}" - print(row) - - print(f"\n [{cp_label}] Read latency (avg of 5 get_state) — lower is better") - print(header) - print(" " + "-" * (len(header) - 2)) - for turns in SWEEP_TURN_COUNTS: - row = f" {turns:>6} {_approx_tokens(turns):>10}" - for label in freq_labels: - _, rt, _ = results[turns][label] - row += f" {f'{rt * 1000:.1f}ms':>{col_w}}" - print(row) - - print( - f"\n [{cp_label}] Per-invoke write latency (total / turns) — lower is better" - ) - print(header) - print(" " + "-" * (len(header) - 2)) - for turns in SWEEP_TURN_COUNTS: - row = f" {turns:>6} {_approx_tokens(turns):>10}" - for label in freq_labels: - wt, _, _ = results[turns][label] - row += f" {f'{(wt / turns) * 1000:.1f}ms':>{col_w}}" - print(row) - - -def run_snapshot_freq_benchmark() -> None: - print() - print("Part 2 — DeltaChannel snapshot_frequency sweep") - print("Lower freq → fewer snapshots → less storage but deeper read replay") - print("=" * 80) - for cp_label, cp_hint in _checkpointers(): - _run_sweep_for_checkpointer(cp_label, cp_hint) - print() - print("Legend:") - print( - " freq=1 snapshot every write (full blob always — same as add_messages / BinOp)" - ) - print(" freq=N snapshot every N writes; read walks at most N ancestor writes") - print(" freq=inf pure delta; read walks entire ancestor chain") - print() - - -# --------------------------------------------------------------------------- -# Pytest entry points +# Pytest entry point # --------------------------------------------------------------------------- @pytest.mark.skip( reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py" ) -def test_delta_channel_baseline_benchmark(capsys: Any) -> None: - """DeltaChannel(inf) uses less storage than add_messages at scale.""" +def test_delta_channel_benchmark(capsys: Any) -> None: + """Manual benchmark — see module docstring.""" with capsys.disabled(): - run_baseline_benchmark() - - for turns in [25, 50]: - _, _, b_bytes = _run_turns(turns, BinaryState) - _, _, d_bytes = _run_turns(turns, DeltaState) - assert d_bytes < b_bytes, ( - f"DeltaChannel should use less storage at {turns} turns, " - f"got delta={d_bytes} binary={b_bytes}" - ) - - -@pytest.mark.skip( - reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py" -) -def test_snapshot_freq_benchmark(capsys: Any) -> None: - """snapshot_frequency trades storage for bounded read depth.""" - with capsys.disabled(): - run_snapshot_freq_benchmark() - - # Correctness: results at all frequencies should agree on final state. - n_turns = 20 - states: dict[str, list] = {} - for freq in SNAPSHOT_FREQUENCIES: - state_cls = _make_delta_state(freq) - graph = _make_graph(state_cls) - config = {"configurable": {"thread_id": "correctness"}} - for i in range(n_turns): - graph.invoke( - {"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]}, - config, - ) - state = graph.get_state(config) - states[_freq_label(freq)] = [m.id for m in state.values["messages"]] - - ref = states["inf"] - for label, msg_ids in states.items(): - assert msg_ids == ref, ( - f"freq={label} produced different message IDs than freq=inf" - ) + run_benchmark() # --------------------------------------------------------------------------- # Script entry point # --------------------------------------------------------------------------- + if __name__ == "__main__": - run_baseline_benchmark() - run_snapshot_freq_benchmark() + print("DeltaChannel benchmark — multi-channel reads, mixed frequencies") + print("=" * 78) + run_benchmark() sys.exit(0) diff --git a/libs/langgraph/tests/test_delta_channel_migration.py b/libs/langgraph/tests/test_delta_channel_migration.py index 245ba2561..50e21ffcb 100644 --- a/libs/langgraph/tests/test_delta_channel_migration.py +++ b/libs/langgraph/tests/test_delta_channel_migration.py @@ -6,11 +6,14 @@ 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, -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 -it as the base value, and `replay_writes(writes)` folds on-path deltas. +Mechanism under test: the saver's public `get_delta_channel_history(config, +channels)` walks the parent chain; when it encounters an ancestor whose +`channel_values[channel]` is a real value, it populates that channel's +`seed` in the returned `DeltaChannelHistory`. If the walk reaches the root +without finding a stored value, the `seed` key is omitted (TypedDict +absence indicates "start empty"). `DeltaChannel.from_checkpoint(seed)` +uses it as the base value, and `replay_writes(writes)` folds on-path +deltas. Scenarios covered: @@ -29,8 +32,8 @@ 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 - same result as the optimized path. + `BaseCheckpointSaver.get_delta_channel_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, one migrated from pre-migration state — don't cross-contaminate. @@ -266,8 +269,8 @@ 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 - `BaseCheckpointSaver` rather than overriding it. + `get_delta_channel_history` implementation from `BaseCheckpointSaver` + rather than overriding it. We rebind the two methods to the base-class versions (via MRO) so the fallback path is exercised even though the storage layer is @@ -275,11 +278,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_delta_channel_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1].get_delta_channel_history # type: ignore[attr-defined] ) - _aget_channel_writes_history = ( # type: ignore[assignment] - InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + aget_delta_channel_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1].aget_delta_channel_history # type: ignore[attr-defined] ) @@ -326,7 +329,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_delta_channel_history` must be scoped to the target thread. """ checkpointer = InMemorySaver() @@ -429,9 +432,10 @@ async def test_tip_of_pre_migration_hydrates_directly_async() -> None: def test_update_state_after_migration_uses_written_value() -> None: """After migrating and running at least one post-migration super-step - (so the thread's tip has a `DELTA_SENTINEL`), `update_state` writes a - concrete value to a new checkpoint's `channel_values`. `get_state` - must reflect that concrete value.""" + (so the thread's tip has its delta channel absent from `channel_values`, + or stored as a `_DeltaSnapshot` on a snapshot step), `update_state` + writes a concrete value into a new checkpoint's `channel_values`. + `get_state` must reflect that concrete value.""" checkpointer = InMemorySaver() config = {"configurable": {"thread_id": "update-state"}} @@ -441,7 +445,8 @@ def test_update_state_after_migration_uses_written_value() -> None: _drive(binop, config, "u", 2) # Migrate and run one more super-step so the tip is a post-migration - # checkpoint with `DELTA_SENTINEL` in its own `channel_values`. + # checkpoint where the delta channel is absent from `channel_values` + # (no snapshot fired this step). delta = _delta_graph(checkpointer) delta.invoke({"items": ["post"]}, config) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d4cb624e9..c686a8071 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -25,7 +25,6 @@ from langchain_core.runnables import ( from langchain_core.runnables.graph import Edge from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( - DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint, CheckpointMetadata, @@ -9618,7 +9617,9 @@ async def test_delta_channel_durability_exit_stores_snapshot() -> None: async def test_delta_channel_async_write_ordering() -> None: """In async mode, DeltaChannel write futures are awaited before the checkpoint - is committed, so aput_writes always precedes aput for sentinel checkpoints.""" + is committed, so aput_writes always precedes aput for delta-channel + checkpoints (those where the delta channel had a versioned write but + is absent from `channel_values`, i.e. no snapshot fired this step).""" class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9637,10 +9638,17 @@ async def test_delta_channel_async_write_ordering() -> None: return result async def tracked_aput(self, config, checkpoint, metadata, new_versions): - has_sentinel = any( - v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values() + # A "delta" checkpoint here = `messages` versioned but absent from + # `channel_values` (no snapshot fired). When a snapshot does fire, + # `channel_values["messages"]` is a `_DeltaSnapshot` — also a delta + # checkpoint shape, since the writes still have to be persisted + # before the parent checkpoint commits. + channel_values = checkpoint.get("channel_values", {}) + is_delta_step = ( + "messages" in checkpoint.get("channel_versions", {}) + and "messages" not in channel_values ) - order.append("aput_sentinel" if has_sentinel else "aput_other") + order.append("aput_delta" if is_delta_step else "aput_other") return await original_aput(self, config, checkpoint, metadata, new_versions) InMemorySaver.aput_writes = tracked_aput_writes @@ -9657,18 +9665,18 @@ async def test_delta_channel_async_write_ordering() -> None: {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config ) - # Every aput_sentinel must be preceded by at least one aput_writes + # Every aput_delta must be preceded by at least one aput_writes for i, event in enumerate(order): - if event == "aput_sentinel": + if event == "aput_delta": preceding = order[:i] assert "aput_writes" in preceding, ( - f"aput_sentinel at {i} had no preceding aput_writes: {order}" + f"aput_delta at {i} had no preceding aput_writes: {order}" ) last_write_idx = max( j for j, e in enumerate(order[:i]) if e == "aput_writes" ) assert last_write_idx < i, ( - f"aput_writes at {last_write_idx} should precede aput_sentinel at {i}: {order}" + f"aput_writes at {last_write_idx} should precede aput_delta at {i}: {order}" ) finally: InMemorySaver.aput_writes = original_aput_writes