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

feat: public get_writes_history saver API + delta cadence rework (#7699)

Browse Source 
## Summary

- Promotes the private K-channel batched ancestor-walk to a stable
public `get_delta_channel_history` / `aget_delta_channel_history` API on
`BaseCheckpointSaver` (returns `Mapping[str, DeltaChannelHistory]`, a
TypedDict with `writes` always present and `seed` `NotRequired`)
- Removes `DELTA_SENTINEL` / `_DeltaSentinel` entirely — the saver layer
is now delta-agnostic on both write and read paths
- Reworks `DeltaChannel` snapshot cadence from "every Nth superstep" to
"every N updates to this channel," persisted in
`CheckpointMetadata.delta_updates_since_snapshot`
- Adds Postgres optimizations: paged stage-1 with cursor (1024-row
pages) and per-channel UNION ALL stage-2 (no over-fetch when channels
have different chain depths)
- Default `snapshot_frequency` becomes a positive int (default `1000`);
the previous `None` opt-out is removed

## Public API

```python
class DeltaChannelHistory(TypedDict):
    writes: list[PendingWrite]            # always present, possibly empty
    seed: NotRequired[Any]                # absent if walk reached root

def get_delta_channel_history(
    self, *, config: RunnableConfig, channels: Sequence[str]
) -> Mapping[str, DeltaChannelHistory]: ...

async def aget_delta_channel_history(
    self, *, config: RunnableConfig, channels: Sequence[str]
) -> Mapping[str, DeltaChannelHistory]: ...
```

`config` and `channels` are keyword-only so later additions (e.g.
`page_size`) don't shift the positional API.

The TypedDict-with-`NotRequired[seed]` shape matches the existing
checkpoint-package convention (`CheckpointMetadata` is
`TypedDict(total=False)`) — absence-via-key-omission rather than
introducing a new sentinel. Pregel translates `"seed" not in hist` to
`MISSING` on its side at consume time.

The default impl walks `get_tuple` + `parent_config` correctly but is
slow on long chains; savers that care override (`InMemorySaver`,
`PostgresSaver`).

## Sentinel removal

`DELTA_SENTINEL` and `_DeltaSentinel` are deleted entirely. The saver
layer becomes delta-agnostic:

- `DeltaChannel.checkpoint()` returns `MISSING` for non-snapshot steps;
pregel's `create_checkpoint` skips MISSING so delta channels without a
snapshot simply don't appear in `channel_values`
- `InMemorySaver.put` and Postgres `put` no longer filter sentinels
(they have nothing to filter)
- `_needs_replay` becomes `stored is MISSING`
- `DeltaChannel.from_checkpoint` accepts: `MISSING` → empty,
`_DeltaSnapshot(value)` → snapshot value, plain value → pre-migration
legacy

## Snapshot cadence

`DeltaChannel.snapshot_frequency: int` (default `1000`, positive). The
previous `None` opt-out is gone.

```python
def should_snapshot(ch_name, ch):
    if force_delta_snapshot:                                  # durability="exit"
        return True
    return updates_since_snapshot.get(ch_name, 0) >= ch.snapshot_frequency
```

Per-channel update counters are persisted in
`CheckpointMetadata.delta_updates_since_snapshot` (`NotRequired`,
`total=False`). The counter is incremented by `_put_checkpoint` for any
delta channel in `updated_channels` and reset to `0` by
`create_checkpoint` for channels that fire a snapshot this step.
Version-format-independent — works for `int`, `float`, and `str`
versioning schemes alike.

## Postgres optimization

Two improvements internal to the override:

**Stage-1 paged with cursor** (`LIMIT 1024` internal const, `AND
checkpoint_id < ?` for subsequent pages). The previous unpaged form
scanned every checkpoint in `(thread_id, ns)` and was pathological at
high thread depths.

**Stage-2 per-channel UNION ALL**: one `WHERE channel='X' AND
checkpoint_id = ANY(chain_X)` branch per channel plus one seed-blob
branch per channel with a seed. The previous form filtered by `channel =
ANY(channels) AND checkpoint_id = ANY(union_chain_cids)`, over-fetching
writes when channels had different chain depths (`K ×
max(chain_lengths)` vs the correct `sum(chain_lengths)`).

Both improvements stay internal to `PostgresSaver`/`AsyncPostgresSaver`;
the public contract returns a single `Mapping`.

## Benchmarks

`libs/langgraph/tests/test_delta_channel_benchmark.py`. Run via `python
libs/langgraph/tests/test_delta_channel_benchmark.py`. Postgres against
local pg:5441.

Results below trimmed to the high-signal cells. Sub-millisecond /
sub-100-turn rows omitted as warmup-bound; freq=1 omitted (chain depth =
1, nothing to optimize); peak read-time memory and Postgres storage are
flat between branches and omitted. Deep-thread reads and the
cadence-rework storage win are the load-bearing numbers.

### Postgres reads, 500 turns

| Scenario | main | branch | Δ |
|---|---:|---:|---:|
| Single-channel deep read | 17.7 ms | **6.1 ms** | **-66%** |
| Single-channel, 1000 turns | 35.0 ms | **14.3 ms** | **-59%** |
| K=3 channels, freq=50 uniform | 70.5 ms | **41.4 ms** | **-41%** |
| K=8 channels, freq=50 uniform | 214.2 ms | **139.4 ms** | **-35%** |
| K=8 channels, mixed freq (25/50/100/.../1000) | 295.6 ms | **214.4
ms** | **-27%** |

K-channel batching + paged stage-1 + per-channel UNION ALL stage-2 doing
exactly what they should at depth.

### InMemory reads, 500 turns

| Scenario | main | branch | Δ |
|---|---:|---:|---:|
| Single-channel deep read | 7.9 ms | **3.8 ms** | **-52%** |
| Single-channel, 1000 turns | 15.6 ms | **7.2 ms** | **-54%** |
| K=8 channels, freq=50 uniform | 112.3 ms | 94.6 ms | -16% |
| K=8 channels, mixed freq | 184.9 ms | **134.5 ms** | **-27%** |

### InMemory storage, 500 turns (cadence-rework win)

| Scenario | main | branch | Δ |
|---|---:|---:|---:|
| K=3, freq=50 uniform | 8.7 MB | **3.3 MB** | **-62%** |
| K=3 mixed freq | 3.8 MB | **1.3 MB** | **-66%** |
| K=8, freq=50 uniform | 23.1 MB | **8.7 MB** | **-62%** |
| K=8 mixed freq | 11.5 MB | **4.2 MB** | **-64%** |

Snapshot frequency now counts **channel updates** instead of
**supersteps**. On graphs where supersteps outpace per-channel updates
(e.g., input/end steps that don't write to channels), branch stores ~3×
fewer snapshot blobs.

### Tradeoff worth flagging

InMemory K=3 with mixed frequencies (50/200/1000) at 500 turns: **+64%
read latency** (46.6 → 76.5 ms). The mixed scenario has a channel with
`freq=1000` that goes the entire 500-turn run with no snapshot. On main,
the old superstep-counted cadence happened to fire at step=500 anyway.
New cadence gives users explicit control over walk depth via
`snapshot_frequency`. The K=8 mixed case still wins overall (-27%); this
regression is specific to the K=3 mixed shape.

Default `snapshot_frequency=1000` is the upper bound on walk depth —
it's a tunable knob.

## Tests

- New sqlite smoke test (`test_get_delta_channel_history.py`) exercises
the inherited default `BaseCheckpointSaver` impl via `SqliteSaver` /
`AsyncSqliteSaver` end-to-end with a real `DeltaChannel`-backed graph.
Sqlite uses the default unchanged — this validates the default path
actually works on a real second saver, not just on the optimized
override.
- Module-level `pytest.importorskip("langgraph.channels.delta")` guards
the test for sqlite's standalone CI environment (matches the postgres
pattern).

## Test plan

- [x] `libs/checkpoint`: 150 passed, 16 skipped
- [x] `libs/langgraph` (channels + delta migration): 41/41 (post-merge)
- [x] `libs/langgraph` (full pregel suite): 1784 passing — 6 "failures"
verified via `env -i` clean shell are local LangSmith env vars + `git
describe revision_id` polluting LangChain metadata fixtures; CI is
unaffected
- [x] `libs/checkpoint-postgres`: 40/40 saver tests + 3/3 delta channel
reconstruction tests against local Postgres
- [x] `libs/checkpoint-sqlite`: 105/105 (incl. retry-passed flake
`test_ttl_refresh`, unrelated to this PR)
- [x] Lint clean across all four libs (`ruff format`, `ruff check`,
`mypy`)
- [x] Branch-vs-main benchmarks — see results above

---------

Co-authored-by: Quanzheng Long <long@langchain.dev>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sydney Runkle
2026-05-04 15:18:43 -04:00
committed by GitHub
co-authored by Quanzheng Long Cursor
parent 35cea28707
commit 0a53c385b2
16 changed files with 1352 additions and 901 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+99 -44
View File
@@ -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),
)
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+82 -44
View File
@@ -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),
)
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+237 -97
View File
@@ -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,
libs/checkpoint-postgres/uv.lock
Generated
+16 -3
View File
@@ -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]]
libs/checkpoint-sqlite/tests/test_get_delta_channel_history.py
+271
View File
@@ -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"] == []
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+104 -100
View File
@@ -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.
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+74 -77
View File
@@ -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(
{
libs/checkpoint/langgraph/checkpoint/serde/types.py
+3 -18
View File
@@ -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.
libs/checkpoint/tests/test_memory.py
+60 -36
View File
@@ -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.
libs/langgraph/langgraph/channels/delta.py
+24 -29
View File
@@ -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
libs/langgraph/langgraph/pregel/_checkpoint.py
+114 -31
View File
@@ -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 {})
libs/langgraph/langgraph/pregel/_loop.py
+23 -1
View File
@@ -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)
libs/langgraph/tests/test_channels.py
+35 -67
View File
@@ -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.
"""
libs/langgraph/tests/test_delta_channel_benchmark.py
+170 -327
View File
@@ -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)
libs/langgraph/tests/test_delta_channel_migration.py
+23 -18
View File
@@ -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)
libs/langgraph/tests/test_pregel.py
+17 -9
View File
@@ -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