Compare commits

...
Author SHA1 Message Date
Quanzheng Long 9800e80296 add_test 2026-04-28 11:19:23 -07:00
Sydney Runkle 0ae81f3cff format, lint, restructure 2026-04-24 07:46:57 -04:00
Sydney Runkle afec98f369 internal for now 2026-04-24 07:29:52 -04:00
Sydney RunkleandClaude Opus 4.7 f25d1935ef fix(postgres): handle missing checkpoint_id in _get_channel_writes_history; update test signatures
Two fixes exposed by running the postgres test suite against a local
postgres instance:

1. `PostgresSaver._get_channel_writes_history` /
   `AsyncPostgresSaver._aget_channel_writes_history` required
   `checkpoint_id` in the passed config, raising `KeyError` when called
   with just `thread_id` (e.g. `graph.aget_state({"thread_id": "..."})`).
   Now resolves to the latest checkpoint via `get_tuple`/`aget_tuple`
   when the id is missing.

2. `test_get_checkpoint_no_channel_values` (sync + async) monkeypatched
   `_load_checkpoint_tuple` with the old `(value, cur)` signature. Method
   now takes `(value)` only since delta reconstruction moved out of the
   tuple-load path — updated both tests.

Local postgres (`brew install pgvector postgresql@16`, running on port
5441) now exercises all 40 non-vector postgres tests green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:47:22 -04:00
Sydney RunkleandClaude Opus 4.7 3a7ed5b454 refactor(delta-channel): honest data model, private experimental API
Restructure DeltaChannel reconstruction so the hydration path matches
pregel's storage axes (blobs + writes) without leaking internal DTOs
into the public checkpoint contract.

Key changes:

* Deleted `DeltaChannelWrites` dataclass and `SEED_UNSET` sentinel.
  Reconstruction data no longer flows through `Checkpoint.channel_values`
  as a wrapped DTO — that field now carries a value or `DELTA_SENTINEL`,
  never a reconstruction shape.
* Added private `_ChannelWritesHistory(seed: Any, writes: list[PendingWrite])`
  NamedTuple as the return type for the new storage-level query.
* Added private, experimental `_get_channel_writes_history` /
  `_aget_channel_writes_history` on `BaseCheckpointSaver` — reference
  impl via `get_tuple` + `parent_config` walk, overridden on
  `InMemorySaver` / `PostgresSaver` / `AsyncPostgresSaver` for perf.
  Fixes a latent migration bug in the base fallback (now inspects
  ancestor `channel_values` for pre-delta seed).
* `DeltaChannel.from_checkpoint(seed)` simplified to two cases
  (sentinel/MISSING → empty, else → seed). New `replay_writes` method
  folds `list[PendingWrite]` through the reducer.
* Delta hydration consolidated inside `channels_from_checkpoint` via
  optional `saver` + `config` kwargs (+ async mirror
  `achannels_from_checkpoint`). All six pregel call sites updated.
  `get_tuple` no longer patches `channel_values` — removed
  `_resolve_delta_channels` (memory) and per-tuple reconstruction from
  `_load_checkpoint_tuple` (postgres sync + async).
* Hydration short-circuits on the target's own blob: if
  `channel_values[k]` is a real value (pre-migration tip, `update_state`
  result), use it directly. Only walks ancestors when the target holds
  sentinel or is missing. Fixes a correctness bug where migration-tip
  and `update_state` values would be lost.
* New test_delta_channel_migration.py: 10 scenarios covering
  BinaryOperatorAggregate → DeltaChannel migration (basic + async,
  time-travel, fork, `update_state`, tip-of-pre-migration, base-saver
  fallback parity, cross-thread isolation).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:15:02 -04:00
Sydney Runkle 31ef0e942a refactor(delta-channel): drop snapshot_every and saver Overwrite terminator
snapshot_every was a knob for bounding reconstruction cost on deep threads.
Benchmarks (notes/add_messages_replay_problem.md + scratch work on
sr/add-messages-replay-bench) showed the add_messages fast-path
(optimize/add-messages-fast-path) closes the quadratic replay cost for
threads under ~1000 turns, where the crossover to snapshots makes sense.
For deeper threads we'll ship a first-class compaction primitive instead.

Removals:

* DeltaChannel: snapshot_every ctor param, _writes_since_snapshot counter,
  should_snapshot() / snapshot_write() methods, counter threading through
  _apply_write / update / from_checkpoint / copy.
* Pregel loop: post-checkpoint snapshot-injection block and
  SNAPSHOT_TASK_ID import + constant.
* Checkpoint base: _overwrite_types() helper and the ancestor-walk
  short-circuit on user-emitted Overwrite in sync + async
  get_channel_writes.
* InMemory + Postgres savers: same walk-terminator shortcut. The
  pre-delta blob terminator (seed-from-ancestor-blob) stays — it's
  required for migration correctness, not a snapshot optimization.
* Tests for all of the above.

Preserved:

* Channel-level Overwrite semantics in DeltaChannel / BinOpAggregate:
  Overwrite still resets the value at reducer level; same-super-step
  dedup and InvalidUpdateError on multiple Overwrites still enforced.
* Pre-delta migration seeding.
2026-04-23 09:54:12 -04:00
Sydney Runkle d120f127ca refactor(delta-channel): plain SELECT WHERE replaces recursive CTE
The recursive CTE was bottlenecked by a JSON-expression join
(`bl.version = checkpoint->'channel_versions'->>bl.channel`) that the
planner could not index, producing an O(ancestors x blobs) nested-loop.
At depth 1000 it ran ~275 ms and removed ~2M filter rows; the recursion
itself was 2.4 ms.

Switch to three plain indexed SELECTs per delta channel
(checkpoints, checkpoint_writes, checkpoint_blobs); a pure helper on
BasePostgresSaver walks the parent chain and assembles
DeltaChannelWrites. Sync (__init__.py) and async (aio.py) each own
their three-roundtrip I/O wrappers.

Bench numbers (notes/delta_channel_query_bench.md): 3x at depth 50,
15x at depth 200, ~100x at depth 1000. Plain over-fetches sibling rows
when the thread branches but still wins at every realistic depth on
both local and remote postgres.

Multi-channel coalescing dropped — reconstruction is per-channel now.
Same shape as InMemorySaver. Can come back as a SQL-level
optimization later if needed.
2026-04-23 08:21:00 -04:00
Sydney Runkle 9e330c96dc contextvar 2026-04-23 07:42:51 -04:00
Sydney RunkleandClaude Opus 4.7 cd8fad5905 fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk
Four fixes from an independent review of the reconstruction pipeline, plus
a structural cleanup:

1. Ancestor walk excludes the target checkpoint itself (matches pregel:
   writes stored under checkpoint_id=T are pending for the NEXT step and
   applied separately via apply_writes). Memory saver previously included
   them, diverging from Postgres and causing pending writes to be folded
   into the reconstructed snapshot — visible via get_state during
   interrupts and time-travel into a non-leaf checkpoint.

2. Pre-delta blob terminator. When the walk hits an ancestor whose blob
   for the channel is a real value (not DELTA_SENTINEL), bind that blob
   as DeltaChannelWrites.seed and stop. Without this, threads migrated
   from pre-delta storage would replay ancestor writes to the root
   forever AND lose any value that lived only in the old blob
   (e.g. from update_state). Per-ancestor, the blob is checked BEFORE
   its writes — a pre-delta blob subsumes writes at the same checkpoint,
   so including them would double-count.

3. Base-fallback get_channel_writes follows parent_checkpoint_id instead
   of list(before=...). The previous form returned every tuple with
   id<target, including sibling branches on forked threads.

4. seed replaces the Overwrite-wrapping hack for pre-delta values.
   DeltaChannelWrites(writes, seed=SEED_UNSET) makes the saver's
   reconstruction terminator semantically explicit; drops the lazy
   _make_overwrite import dance. User-emitted Overwrite still reset the
   chain via _apply_write as before.

Postgres: recursive CTE enumerates on-path ancestors and joins once
against checkpoint_writes and once against checkpoint_blobs for every
delta channel in the get_tuple — one roundtrip instead of the previous
3 queries × N channels.

Tests added:
- Pre-delta blob seeding (seed binding, no double-counting of ancestor
  writes at the terminator, pending-at-target excluded).
- Root checkpoint returns empty writes.
- Seed-based from_checkpoint replay (three scenarios: with writes,
  seed-only, seed=None distinct from SEED_UNSET).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-22 22:09:08 -04:00
Sydney Runkle acc7eda8c5 optimizations i sure hope 2026-04-22 20:47:51 -04:00
Sydney Runkle 5b7fdf5655 eh 2026-04-22 18:41:06 -04:00
Sydney RunkleandClaude Opus 4.7 4cad68f767 fix(delta-channel): unwrap NotRequired[X] for dict/set reducers
Annotated[NotRequired[dict[...]], DeltaChannel(reducer)] (the shape used
by deepagents' filesystem middleware) fell through type inference to
`list`, so the first operator call blew up with
"'list' object is not a mapping". `_is_field_channel` now unwraps a
parameterized Required[X]/NotRequired[X] before stripping extras, which
lets dict/set/mapping outer types reach the abc normalization block.

Also type-annotates the `new` locals in DeltaChannel.copy() and
from_checkpoint() so mypy can infer them through the abstract return
type.

Adds tests covering: dict Overwrite in update and in writes replay,
snapshot_write with a dict reducer, dict backwards-compat checkpoints,
NotRequired type inference, and a filesystem-shaped end-to-end graph.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-22 17:58:36 -04:00
Sydney Runkle 9d8c0be068 arbitrary 2026-04-22 17:54:12 -04:00
Sydney Runkle 51154be4ab lint 2026-04-22 17:47:53 -04:00
Sydney Runkle 2e7edb2b60 lint 2026-04-22 17:30:22 -04:00
Sydney Runkle 325cb42f19 lint and snapshot every 2026-04-22 17:06:53 -04:00
Sydney Runkle b9fad696ec cleanup 2026-04-22 16:45:32 -04:00
26 changed files with 2201 additions and 677 deletions
@@ -4,15 +4,17 @@ import threading
from collections import defaultdict
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, cast
from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
_ChannelWritesHistory,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -23,7 +25,12 @@ from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
Conn = _internal.Conn # For backward compatibility
@@ -32,7 +39,7 @@ Conn = _internal.Conn # For backward compatibility
class PostgresSaver(BasePostgresSaver):
"""Checkpointer that stores checkpoints in a Postgres database."""
lock: threading.RLock
lock: threading.Lock
def __init__(
self,
@@ -48,7 +55,7 @@ class PostgresSaver(BasePostgresSaver):
self.conn = conn
self.pipe = pipe
self.lock = threading.RLock()
self.lock = threading.Lock()
self.supports_pipeline = Capabilities().has_pipeline()
@classmethod
@@ -430,6 +437,42 @@ 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`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`) each filtered by `(thread_id, checkpoint_ns)` and
the per-table key. Plain SELECTs let the planner pick straight index
scans; rationale + benchmark in `notes/delta_channel_query_bench.md`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
# Caller didn't specify a target — resolve to the latest
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
# returns the newest; its config carries the resolved id.
target = self.get_tuple(config)
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
with self._cursor() as cur:
cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
parents_rows = cur.fetchall()
cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
writes_rows = cur.fetchall()
cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
blobs_rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
)
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
Convert a database row into a CheckpointTuple object.
@@ -442,22 +485,7 @@ class PostgresSaver(BasePostgresSaver):
including its configuration, metadata, parent checkpoint (if any),
and pending writes.
"""
from langgraph.checkpoint.base import DeltaChannelSentinel
channel_values = self._load_blobs(value["channel_values"])
if any(isinstance(v, DeltaChannelSentinel) for v in channel_values.values()):
cp_config = cast(
RunnableConfig,
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
}
},
)
with self._cursor() as cur:
self._resolve_delta_channels(cp_config, channel_values, cur)
return CheckpointTuple(
{
"configurable": {
@@ -8,12 +8,13 @@ from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
DeltaChannelSentinel,
_ChannelWritesHistory,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -24,7 +25,12 @@ from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
Conn = _ainternal.Conn # For backward compatibility
@@ -392,57 +398,44 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
async def _aget_channel_writes_cur(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
cur: Any,
) -> list[Any]:
"""Async version of _get_channel_writes_cur — see sync version for rationale."""
await cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s",
(thread_id, checkpoint_ns),
)
parent_map: dict[str, str | None] = {
row["checkpoint_id"]: row["parent_checkpoint_id"]
for row in await cur.fetchall()
}
ancestor_ids: list[str] = []
cid: str | None = parent_map.get(checkpoint_id)
while cid is not None:
ancestor_ids.append(cid)
cid = parent_map.get(cid)
if not ancestor_ids:
return []
await cur.execute(
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
" AND checkpoint_id = ANY(%s) "
"ORDER BY task_id, idx",
(thread_id, checkpoint_ns, channel, ancestor_ids),
)
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
for row in await cur.fetchall():
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
result = []
for cid in reversed(ancestor_ids):
for type_tag, blob in writes_by_cp.get(cid, []):
result.append(self.serde.loads_typed((type_tag, blob)))
return result
async def aget_channel_writes(
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> list[Any]:
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`); rows assembled by the shared pure helper on
`BasePostgresSaver`. Rationale + benchmark in
`notes/delta_channel_query_bench.md`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"]["checkpoint_id"]
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=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
async with self._cursor() as cur:
return await self._aget_channel_writes_cur(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
await cur.execute(
SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns)
)
parents_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel)
)
writes_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel)
)
blobs_rows = await cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
)
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
@@ -458,23 +451,10 @@ class AsyncPostgresSaver(BasePostgresSaver):
"""
thread_id = value["thread_id"]
checkpoint_ns = value["checkpoint_ns"]
checkpoint_id = value["checkpoint_id"]
blob_values = value["channel_values"]
channel_values: dict[str, Any] = {}
if blob_values:
channel_values = self._load_blobs(blob_values)
delta_channels = [
ch
for ch, v in channel_values.items()
if isinstance(v, DeltaChannelSentinel)
]
if delta_channels:
async with self._cursor() as cur:
for channel in delta_channels:
channel_values[channel] = await self._aget_channel_writes_cur(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
)
return CheckpointTuple(
{
@@ -2,17 +2,18 @@ from __future__ import annotations
import random
import warnings
from collections import defaultdict
from collections.abc import Sequence
from importlib.metadata import version as get_version
from typing import Any, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
DeltaChannelSentinel,
PendingWrite,
_ChannelWritesHistory,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.types import TASKS
@@ -154,6 +155,30 @@ INSERT_CHECKPOINT_WRITES_SQL = """
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
# DeltaChannel reconstruction: three plain indexed SELECTs per channel.
# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE
# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are
# 3x-100x faster in the realistic depth range and the Python walk is O(n).
SELECT_DELTA_PARENTS_SQL = """
SELECT checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
"""
SELECT_DELTA_WRITES_SQL = """
SELECT checkpoint_id, type, blob, task_id, idx
FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
SELECT_DELTA_BLOBS_SQL = """
SELECT version, type, blob
FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
@@ -199,66 +224,78 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
result[k.decode()] = self.serde.loads_typed((type_tag, v))
return result
def _resolve_delta_channels(
def _build_delta_channel_writes_history(
self,
config: RunnableConfig,
channel_values: dict[str, Any],
cur: Any,
) -> None:
for channel, value in list(channel_values.items()):
if isinstance(value, DeltaChannelSentinel):
channel_values[channel] = self._get_channel_writes_cur(
config["configurable"]["thread_id"],
config["configurable"].get("checkpoint_ns", ""),
config["configurable"]["checkpoint_id"],
channel,
cur,
)
def _get_channel_writes_cur(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
*,
channel: str,
cur: Any,
) -> list[Any]:
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.
target_id: str,
parents_rows: Sequence[Any],
writes_rows: Sequence[Any],
blobs_rows: Sequence[Any],
) -> _ChannelWritesHistory:
"""Reconstruct one delta channel's history from rows of the three SELECTs.
Two queries:
1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, IDs only.
2. Walk the ancestor chain in Python, then fetch writes with a plain ANY() filter.
Pure data transform shared by sync (`PostgresSaver`) and async
(`AsyncPostgresSaver`); both paths run the queries themselves and
feed the rows here.
Walk is newest → oldest from the target's parent. A non-sentinel
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
walk and is returned as the seed so replay starts from it.
Writes stored at `target_id` itself are pending writes for the next
step and are excluded — the walk begins at the target's parent.
"""
cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
"WHERE thread_id = %s AND checkpoint_ns = %s",
(thread_id, checkpoint_ns),
)
parent_map: dict[str, str | None] = {
row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall()
}
ancestor_ids: list[str] = []
cid: str | None = parent_map.get(checkpoint_id)
parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {}
for r in parents_rows:
cid = r["checkpoint_id"]
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
ancestors: list[str] = []
cid = parent_of.get(target_id)
while cid is not None:
ancestor_ids.append(cid)
cid = parent_map.get(cid)
if not ancestor_ids:
return []
cur.execute(
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
" AND checkpoint_id = ANY(%s) "
"ORDER BY task_id, idx",
(thread_id, checkpoint_ns, channel, ancestor_ids),
)
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
for row in cur.fetchall():
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
result = []
for cid in reversed(ancestor_ids):
for type_tag, blob in writes_by_cp.get(cid, []):
result.append(self.serde.loads_typed((type_tag, blob)))
return result
ancestors.append(cid)
cid = parent_of.get(cid)
if not ancestors:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
ancestor_set = set(ancestors)
# Group writes by ancestor cid; sort within (task_id DESC, idx DESC)
# to match the prior CTE ordering — newest write first per ancestor.
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
for r in writes_rows:
cid = r["checkpoint_id"]
if cid not in ancestor_set:
continue
writes_by_cid.setdefault(cid, []).append(
(r["type"], r["blob"], r["task_id"], r["idx"])
)
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
blob_by_ver: dict[str, tuple[str, bytes]] = {
r["version"]: (r["type"], r["blob"]) for r in blobs_rows
}
collected: list[PendingWrite] = [] # newest first; reversed at the end
for cid in ancestors:
# Pre-delta blob terminator: subsumes any writes at this ancestor.
ver = ver_of.get(cid)
if ver is not None:
seed_blob = blob_by_ver.get(ver)
if seed_blob is not None and seed_blob[0] != "empty":
blob_value = self.serde.loads_typed(seed_blob)
if blob_value is not DELTA_SENTINEL:
collected.reverse()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
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.reverse() # oldest → newest
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
def _dump_blobs(
self,
+4 -4
View File
@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
load_checkpoint_tuple = saver._load_checkpoint_tuple
def patched_load_checkpoint_tuple(value):
async def patched_load_checkpoint_tuple(value):
value["checkpoint"].pop("channel_values", None)
return load_checkpoint_tuple(value)
return await load_checkpoint_tuple(value)
monkeypatch.setattr(
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
@@ -377,13 +377,13 @@ async def test_get_checkpoint_no_channel_values(
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
pytest.importorskip(
"langgraph.channels.delta", reason="langgraph core not installed"
"langgraph.channels._delta", reason="langgraph core not installed"
)
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
@@ -1,14 +1,12 @@
from __future__ import annotations
import contextvars
import copy
import dataclasses
import logging
import threading
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import ( # noqa: UP035
from typing import (
Any,
Generic,
List,
Literal,
NamedTuple,
TypedDict,
@@ -21,6 +19,9 @@ 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,
@@ -32,19 +33,14 @@ from langgraph.checkpoint.serde.types import (
V = TypeVar("V", int, float, str)
PendingWrite = tuple[str, str, Any]
# Task-local guard: ContextVar is copied per asyncio Task, so concurrent
# requests on the same event-loop thread do not share this flag. A plain
# `threading.local()` would leak across tasks and let one in-flight
# reconstruction silently short-circuit another.
_DELTA_RECONSTRUCTION: contextvars.ContextVar[bool] = contextvars.ContextVar(
"_DELTA_RECONSTRUCTION", default=False
)
@dataclasses.dataclass
class DeltaChannelSentinel:
"""Marker stored in checkpoint_blobs for a DeltaChannel field.
No data is stored here — the actual per-step writes live in checkpoint_writes
and are replayed through the reducer at load time.
"""
pass
_DELTA_RECONSTRUCTION: threading.local = threading.local()
logger = logging.getLogger(__name__)
@@ -137,6 +133,30 @@ class CheckpointTuple(NamedTuple):
pending_writes: list[PendingWrite] | None = None
class _ChannelWritesHistory(NamedTuple):
"""Result of `BaseCheckpointSaver._get_channel_writes_history`.
Storage-level view of what one channel wrote 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.
"""
seed: Any
writes: list[PendingWrite]
class BaseCheckpointSaver(Generic[V]):
"""Base class for creating a graph checkpointer.
@@ -475,55 +495,103 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006
"""Collect all writes for `channel` across this checkpoint's ancestry, oldest→newest.
Default implementation walks the full thread history via `list()`. Savers can
override with a more efficient query (InMemorySaver and PostgresSaver do this).
"""
# Guard against re-entrant calls: when list() triggers reconstruction which
# calls list() again, the inner call returns tuples with DeltaChannelSentinel
# in channel_values (which get_channel_writes ignores — it only reads
# pending_writes). This breaks the recursion safely.
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
_DELTA_RECONSTRUCTION.active = True
try:
result: list[Any] = []
target_id = config["configurable"].get("checkpoint_id")
for tup in self.list(config):
if tup.config["configurable"].get("checkpoint_id") == target_id:
continue # skip the checkpoint itself; we want its ancestors' writes
if tup.pending_writes:
for _, ch, value in tup.pending_writes:
if ch == channel:
result.append(value)
result.reverse() # list() yields newest→oldest; we want oldest→newest
return result
finally:
_DELTA_RECONSTRUCTION.active = False
async def aget_channel_writes(
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> List[Any]: # noqa: UP006
"""Async version of get_channel_writes."""
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
_DELTA_RECONSTRUCTION.active = True
) -> _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.
Walks the **parent chain** (not `list(before=...)`): for forked
threads, only on-path ancestors contribute.
Reference implementation walks `get_tuple` + `parent_config`,
inspecting each ancestor's `channel_values[channel]` for the seed
terminator. Savers with direct storage access (`InMemorySaver`,
`PostgresSaver`) override for performance; the return contract is
fixed here.
Underscore-prefixed because the method surface is experimental.
"""
# Guard against re-entrant calls: when get_tuple() triggers
# reconstruction which calls get_tuple() again, the inner call
# short-circuits here.
if _DELTA_RECONSTRUCTION.get():
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
token = _DELTA_RECONSTRUCTION.set(True)
try:
result: list[Any] = []
target_id = config["configurable"].get("checkpoint_id")
async for tup in self.alist(config):
if tup.config["configurable"].get("checkpoint_id") == target_id:
continue
collected: list[PendingWrite] = [] # newest first; reversed at the end
target_tuple = self.get_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = self.get_tuple(cursor_config)
if tup is None:
break
# Pre-delta seed terminator: if the ancestor has a stored
# (non-sentinel) value for this channel, that snapshot
# subsumes any earlier writes on the chain. 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)
if tup.pending_writes:
for _, ch, value in tup.pending_writes:
if ch == channel:
result.append(value)
result.reverse()
return result
# 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)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
finally:
_DELTA_RECONSTRUCTION.active = False
_DELTA_RECONSTRUCTION.reset(token)
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""Async version of `_get_channel_writes_history`. See docstring there."""
if _DELTA_RECONSTRUCTION.get():
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
token = _DELTA_RECONSTRUCTION.set(True)
try:
collected: list[PendingWrite] = []
target_tuple = await self.aget_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = await self.aget_tuple(cursor_config)
if tup is None:
break
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)
if tup.pending_writes:
for write in reversed(tup.pending_writes):
if write[1] != channel:
continue
collected.append(write)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
finally:
_DELTA_RECONSTRUCTION.reset(token)
def get_next_version(self, current: V | None, channel: None) -> V:
"""Generate the next version ID for a channel.
@@ -14,14 +14,16 @@ from typing import Any, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
DeltaChannelSentinel,
PendingWrite,
SerializerProtocol,
_ChannelWritesHistory,
get_checkpoint_id,
get_checkpoint_metadata,
)
@@ -138,24 +140,20 @@ class InMemorySaver(
result[k] = self.serde.loads_typed(vv)
return result
def _resolve_delta_channels(
self,
config: RunnableConfig,
channel_values: dict[str, Any],
) -> None:
"""Replace DeltaChannelSentinel entries with reconstructed write lists."""
for channel, value in list(channel_values.items()):
if isinstance(value, DeltaChannelSentinel):
channel_values[channel] = self.get_channel_writes(config, channel)
def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]:
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
# Walk the parent chain newest→oldest collecting checkpoint IDs.
# Walk the parent chain newest→oldest. Skip the target itself —
# writes stored AT `checkpoint_id` are pending for the next step
# (pregel applies them via `apply_writes`; they aren't part of the
# snapshot value AT `checkpoint_id`).
chain: list[str] = []
current: str | None = checkpoint_id
target_entry = ns_storage.get(checkpoint_id)
current: str | None = target_entry[2] if target_entry is not None else None
while current is not None:
entry = ns_storage.get(current)
if entry is None:
@@ -163,19 +161,54 @@ class InMemorySaver(
chain.append(current)
_, _, parent = entry
current = parent
# Collect writes oldest→newest.
result: list[Any] = []
for cp_id in reversed(chain):
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
for (_task_id, _idx), (_, ch, serialized, _) in sorted(step_writes.items()):
if ch == channel:
result.append(self.serde.loads_typed(serialized))
return result
# Scan newest→oldest. A pre-delta blob on an ancestor terminates the
# walk and is bound as `seed`; without this, a thread migrated from
# pre-delta storage would replay ancestor writes all the way to the
# root AND miss any value that lived only in the old blob (e.g. from
# `update_state`).
#
# At each ancestor, check the blob BEFORE processing its pending
# writes: a pre-delta blob represents the state AT that ancestor,
# which already subsumes any writes stored under it. Processing
# those writes first would fold them into the reconstructed value
# twice (once via the blob, once via replay).
collected: list[PendingWrite] = [] # newest first
for cp_id in chain: # newest → oldest
entry = ns_storage.get(cp_id)
if entry is not None:
ckpt = self.serde.loads_typed(entry[0])
ver = ckpt.get("channel_versions", {}).get(channel)
if ver is not None:
blob_entry = self.blobs.get(
(thread_id, checkpoint_ns, channel, ver)
)
if blob_entry is not None and blob_entry[0] != "empty":
blob_value = self.serde.loads_typed(blob_entry)
if blob_value is not DELTA_SENTINEL:
# Pre-delta snapshot terminator. Skip this
# ancestor's writes — the blob subsumes them.
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
async def aget_channel_writes(
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:
continue
val = self.serde.loads_typed(serialized)
collected.append((tid, ch, val))
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> list[Any]:
return self.get_channel_writes(config, channel)
) -> _ChannelWritesHistory:
return self._get_channel_writes_history(config, channel)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Get a checkpoint tuple from the in-memory storage.
@@ -203,7 +236,6 @@ class InMemorySaver(
checkpoint_ns,
checkpoint_["channel_versions"],
)
self._resolve_delta_channels(config, channel_values)
return CheckpointTuple(
config=config,
checkpoint={
@@ -247,7 +279,6 @@ class InMemorySaver(
checkpoint_ns,
checkpoint_["channel_versions"],
)
self._resolve_delta_channels(resolved_config, channel_values)
return CheckpointTuple(
config=resolved_config,
checkpoint={
@@ -362,7 +393,6 @@ class InMemorySaver(
checkpoint_ns,
checkpoint_["channel_versions"],
)
self._resolve_delta_channels(list_config, channel_values)
yield CheckpointTuple(
config=list_config,
@@ -33,14 +33,13 @@ from langchain_core.load.load import Reviver
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
from langgraph.checkpoint.serde.types import SendProtocol
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol
from langgraph.store.base import Item
if TYPE_CHECKING:
from langgraph.checkpoint.serde._msgpack import (
AllowedMsgpackModules,
)
from langgraph.checkpoint.serde.types import SendProtocol
LC_REVIVER = Reviver()
EMPTY_BYTES = b""
@@ -64,14 +63,6 @@ def _warn_once(
logger.warning(msg, *args)
def _get_delta_sentinel_cls() -> type:
from langgraph.checkpoint.base import (
DeltaChannelSentinel,
) # lazy import avoids circular dep
return DeltaChannelSentinel
class JsonPlusSerializer(SerializerProtocol):
"""Serializer that uses ormsgpack, with optional fallbacks.
@@ -260,12 +251,12 @@ class JsonPlusSerializer(SerializerProtocol):
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
if obj is None:
return "null", EMPTY_BYTES
elif obj is DELTA_SENTINEL:
return "delta", EMPTY_BYTES
elif isinstance(obj, bytes):
return "bytes", obj
elif isinstance(obj, bytearray):
return "bytearray", obj
elif isinstance(obj, _get_delta_sentinel_cls()):
return "delta", b""
else:
try:
return "msgpack", _msgpack_enc(obj)
@@ -289,9 +280,7 @@ class JsonPlusSerializer(SerializerProtocol):
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
elif type_ == "delta":
from langgraph.checkpoint.base import DeltaChannelSentinel
return DeltaChannelSentinel()
return DELTA_SENTINEL
elif self.pickle_fallback and type_ == "pickle":
return pickle.loads(data_)
else:
@@ -14,6 +14,25 @@ INTERRUPT = "__interrupt__"
RESUME = "__resume__"
TASKS = "__pregel_tasks"
class _DeltaSentinel:
"""Singleton marker stored (as zero bytes) in checkpoint_blobs for a
DeltaChannel field. The actual per-step writes live in checkpoint_writes
and are replayed through the reducer at load time.
Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same
module-level instance.
"""
__slots__ = ()
def __repr__(self) -> str:
return "DELTA_SENTINEL"
DELTA_SENTINEL = _DeltaSentinel()
Value = TypeVar("Value", covariant=True)
Update = TypeVar("Update", contravariant=True)
C = TypeVar("C")
+6 -5
View File
@@ -999,13 +999,14 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
assert result == obj
def test_delta_channel_sentinel_serde_round_trip() -> None:
from langgraph.checkpoint.base import DeltaChannelSentinel
def test_delta_sentinel_serde_round_trip() -> None:
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
original = DeltaChannelSentinel()
type_tag, blob = serde.dumps_typed(original)
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
# Zero-byte "delta" tag — no allowlist change needed.
assert type_tag == "delta"
assert blob == b""
loaded = serde.loads_typed((type_tag, blob))
assert isinstance(loaded, DeltaChannelSentinel)
assert loaded is DELTA_SENTINEL
+292 -20
View File
@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
@@ -208,8 +209,6 @@ class TestMemorySaver:
async def test_memory_saver() -> None:
from langgraph.checkpoint.memory import InMemorySaver
memory_saver = InMemorySaver()
assert isinstance(memory_saver, InMemorySaver)
@@ -324,20 +323,14 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
class TestInMemorySaverDeltaChannel:
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DeltaChannelSentinel for delta channels (reconstruction deferred)."""
from langgraph.checkpoint.base import (
DeltaChannelSentinel,
empty_checkpoint,
)
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0000000000000000"
sentinel = DeltaChannelSentinel()
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(sentinel)
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
@@ -348,12 +341,12 @@ class TestInMemorySaverDeltaChannel:
result = saver._load_blobs(thread_id, ns, {channel: v1})
assert channel in result
assert isinstance(result[channel], DeltaChannelSentinel)
def test_get_channel_writes_collects_writes(self) -> None:
"""get_channel_writes collects per-step writes oldest→newest."""
from langgraph.checkpoint.base import empty_checkpoint
assert result[channel] is DELTA_SENTINEL
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)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
@@ -367,18 +360,20 @@ class TestInMemorySaverDeltaChannel:
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
}
# cp1 has a write for channel
# Writes stored at cp1 produced the cp1 snapshot; part of history.
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "hi"}),
"",
)
# cp2 has a write for channel
# Writes stored at cp2 are pending — they will produce cp3 when the
# step that loaded cp2 completes. They MUST NOT appear in the
# reconstructed snapshot value at cp2.
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed({"content": "bye"}),
serde.dumps_typed({"content": "pending"}),
"",
)
@@ -389,5 +384,282 @@ class TestInMemorySaverDeltaChannel:
"checkpoint_id": "cp2",
}
}
result = saver.get_channel_writes(config, channel)
assert result == [{"content": "hi"}, {"content": "bye"}]
result = saver._get_channel_writes_history(config, channel)
assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes]
assert values == [{"content": "hi"}]
def test_get_channel_writes_at_root_returns_empty(self) -> None:
"""Reconstructing the root checkpoint's state: no ancestors → []."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
}
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "pending"}),
"",
)
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": "cp1",
}
}
result = saver._get_channel_writes_history(config, channel)
assert result.seed is DELTA_SENTINEL
assert result.writes == []
class TestBaseFallbackGetChannelWrites:
"""Exercises the `BaseCheckpointSaver._get_channel_writes_history` default
implementation — the path third-party savers inherit when they don't
override `_get_channel_writes_history` themselves.
Regression guard for a bug where the fallback passed the caller's config
(with `checkpoint_id`) straight to `self.list()`, which most savers
collapse to a single row — causing the fallback to return `[]`.
"""
def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]:
"""Build an InMemorySaver with a 3-checkpoint chain and per-step writes
for a `messages` channel.
Returns `(saver, thread_id, namespace)`. The saver subclass deletes the
InMemorySaver override so the base class fallback is exercised.
"""
class _ThirdPartyStyleSaver(InMemorySaver):
_get_channel_writes_history = (
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
)
_aget_channel_writes_history = (
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
)
saver = _ThirdPartyStyleSaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
cp0 = empty_checkpoint()
cp0["id"] = "00000000000000000000000000000001.0000000000000000"
cp1 = empty_checkpoint()
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
cp2 = empty_checkpoint()
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
saver.storage[thread_id][ns] = {
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]),
}
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "first"}),
"",
)
saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed({"content": "second"}),
"",
)
return saver, thread_id, ns
def test_fallback_returns_ancestor_writes_oldest_first(self) -> None:
saver, thread_id, ns = self._build_saver_with_chain()
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
result = saver._get_channel_writes_history(config, "messages")
assert result.seed is DELTA_SENTINEL
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:
saver, thread_id, ns = self._build_saver_with_chain()
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
result = await saver._aget_channel_writes_history(config, "messages")
assert result.seed is DELTA_SENTINEL
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
`threading.local()` guard would let whichever task set it first
short-circuit the other to `writes=[]`.
"""
import asyncio
saver, thread_id, ns = self._build_saver_with_chain()
# Force the two tasks to interleave across the `set(True)` boundary:
# each `aget_tuple` yields control, so if the guard were thread-local
# the second task would observe `active=True` set by the first.
orig_aget_tuple = saver.aget_tuple
async def slow_aget_tuple(config: RunnableConfig) -> Any:
await asyncio.sleep(0)
return await orig_aget_tuple(config)
saver.aget_tuple = slow_aget_tuple # type: ignore[method-assign]
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
results = await asyncio.gather(
saver._aget_channel_writes_history(config, "messages"),
saver._aget_channel_writes_history(config, "messages"),
)
expected_values = [{"content": "first"}, {"content": "second"}]
for result in results:
assert result.seed is DELTA_SENTINEL
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
* back-compat: a thread written by pre-delta code, then extended under
delta — reconstruction must return the correct value without walking
past the last pre-delta ancestor;
* perf: without the terminator, every reconstruct-after-migration would
walk all the way to the thread root.
"""
def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]:
"""Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta,
write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield
seed=[A] + writes=[B, C].
Returns `(saver, thread_id, ns, channel, cp3_id)`.
"""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0"
v2 = "00000000000000000000000000000002.0"
v3 = "00000000000000000000000000000003.0"
# Pre-delta: cp1 stored a real blob for the channel.
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
cp2 = empty_checkpoint()
cp2["id"] = "cp2"
cp2["channel_versions"][channel] = v2
cp3 = empty_checkpoint()
cp3["id"] = "cp3"
cp3["channel_versions"][channel] = v3
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
}
# Write under cp1 would be from the pre-delta era and MUST be ignored
# (the blob already captures it). We add one and assert it is not
# folded into the reconstructed result.
saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = (
"task0",
channel,
serde.dumps_typed("PRE-DELTA-WRITE"),
"",
)
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed("B"),
"",
)
saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = (
"task3",
channel,
serde.dumps_typed("PENDING-AT-TARGET"),
"",
)
return saver, thread_id, ns, channel, "cp3"
def test_seed_from_pre_delta_ancestor_blob(self) -> None:
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver._get_channel_writes_history(config, channel)
# Seed came from the pre-delta blob at cp1.
assert result.seed == ["A"]
# Delta-era writes from cp2 replay through the reducer on top of seed.
# cp3 is the target — its own write is pending for the NEXT step and
# must be excluded.
values = [v for _, _, v in result.writes]
assert values == ["B"]
def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None:
"""Writes stored at the pre-delta ancestor itself must not be replayed
(the blob subsumes them)."""
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver._get_channel_writes_history(config, channel)
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.
assert "PENDING-AT-TARGET" not in values
@@ -1,7 +1,6 @@
from langgraph.channels.any_value import AnyValue
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
from langgraph.channels.named_barrier_value import (
@@ -21,7 +20,6 @@ __all__ = (
"UntrackedValue",
"EphemeralValue",
"BinaryOperatorAggregate",
"DeltaChannel",
"NamedBarrierValue",
"NamedBarrierValueAfterFinish",
# topics
+177
View File
@@ -0,0 +1,177 @@
from __future__ import annotations
import copy as _copy
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
from typing_extensions import Self
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel, Value
from langgraph.channels.binop import _get_overwrite
from langgraph.errors import EmptyChannelError
__all__ = ("DeltaChannel",)
def _empty(typ: Any) -> Any:
try:
return typ()
except Exception:
return []
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
"""Experimental — private API, subject to change or removal without notice.
Imported from the underscored module `langgraph.channels._delta` on purpose;
not re-exported from `langgraph.channels`. Intended for internal use only
while we validate the design on real workloads.
A channel that stores only a sentinel in checkpoints; per-step writes are
stored in checkpoint_writes and replayed through the operator at load time.
Use with append-style reducers (e.g. `add_messages`) on long-running threads
to eliminate O(N²) blob growth — storage is O(N) using the writes table that
every checkpointer already maintains.
Reconstruction replays every ancestor write through the operator, so
per-get cost scales with thread depth. Compaction for deep threads is
a follow-up — today, use this on threads of a few hundred turns.
Usage::
from langgraph.channels._delta import DeltaChannel
class State(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
"""
__slots__ = (
"value",
"operator",
)
def __init__(
self,
operator: Callable[[Any, Any], Any],
) -> None:
super().__init__(list)
self.operator = operator
self.value: Any = []
def __eq__(self, other: object) -> bool:
if not isinstance(other, DeltaChannel):
return False
if (
self.operator.__name__ != "<lambda>"
and other.operator.__name__ != "<lambda>"
):
return self.operator is other.operator
return True
@property
def ValueType(self) -> Any:
return self.typ
@property
def UpdateType(self) -> Any:
return self.typ
def copy(self) -> Self:
new: DeltaChannel[Value] = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
return new
def _apply_write(self, value: Any, write: Any) -> Any:
"""Apply one write to `value` and return the new value.
An `Overwrite` replaces the value; any other write is folded through
the operator. Centralizes the Overwrite/reducer branching used by both
`update` (live super-step) and `from_checkpoint` (ancestor replay).
"""
is_overwrite, overwrite_value = _get_overwrite(write)
if is_overwrite:
return (
_copy.copy(overwrite_value)
if overwrite_value is not None
else _empty(self.typ)
)
base = _empty(self.typ) if value is MISSING else value
return self.operator(base, write)
def from_checkpoint(self, checkpoint: Any) -> Self:
"""Initialize from a seed value.
Pregel's hydration path calls this with the `seed` returned by
`saver.get_channel_history`:
* `MISSING` / `DELTA_SENTINEL` → channel starts empty. The walk
either reached the root (fresh delta thread) or found nothing
to seed from.
* any other value → use as the base value. Typically a pre-delta
blob preserved across a channel-type migration; `replay_writes`
folds subsequent deltas on top.
"""
new: DeltaChannel[Value] = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
new.value = _empty(new.typ)
else:
new.value = checkpoint
return new
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
"""Fold a sequence of `PendingWrite` tuples into the current value.
Called after `from_checkpoint` during pregel hydration to replay
per-step deltas from on-path ancestors through the reducer. Writes
are oldest→newest. `Overwrite` values inside the stream reset the
reducer state at that point, same as during a live super-step.
The `task_id` and `channel` fields of each `PendingWrite` are
ignored — `_get_channel_writes_history` has already filtered to
this channel.
"""
for _, _, value in writes:
self.value = self._apply_write(self.value, value)
def update(self, values: Sequence[Any]) -> bool:
if not values:
return False
seen_overwrite = False
for value in values:
is_overwrite, _ = _get_overwrite(value)
if is_overwrite:
if seen_overwrite:
from langgraph.errors import (
ErrorCode,
InvalidUpdateError,
create_error_message,
)
msg = create_error_message(
message="Can receive only one Overwrite value per super-step.",
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
)
raise InvalidUpdateError(msg)
seen_overwrite = True
elif seen_overwrite:
# Post-Overwrite writes within the same super-step are dropped.
continue
self.value = self._apply_write(self.value, value)
return True
def get(self) -> Any:
if self.value is MISSING:
raise EmptyChannelError()
return self.value
def is_available(self) -> bool:
return self.value is not MISSING
def checkpoint(self) -> Any:
return DELTA_SENTINEL
@@ -119,12 +119,3 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
Returns `True` if the channel was updated, `False` otherwise.
"""
return False
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
"""Called after checkpoint() with the assigned version, and after
from_checkpoint() with the current channel version.
No-op by default. Override in channels that track their own version
for incremental checkpointing (e.g. DeltaChannel).
"""
pass
-137
View File
@@ -1,137 +0,0 @@
from __future__ import annotations
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DeltaChannelSentinel
from typing_extensions import Self
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel, Value
from langgraph.channels.binop import _get_overwrite
from langgraph.errors import EmptyChannelError
__all__ = ("DeltaChannel",)
class DeltaChannel(
Generic[Value], BaseChannel[list[Value], Value, DeltaChannelSentinel]
):
"""A channel that stores only a sentinel in checkpoints; per-step writes are
stored in checkpoint_writes and replayed through the operator at load time.
Use with append-style reducers (e.g. `add_messages`) on long-running threads
to eliminate O(N²) blob growth — storage is O(N) using the writes table that
every checkpointer already maintains.
Works with all checkpointers. Savers with dedicated implementations
(InMemorySaver, PostgresSaver) reconstruct in one pass; others fall back to
walking the checkpoint list.
Usage::
class State(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
# Dict-type reducer (type inferred from the Annotated outer type):
files: Annotated[dict, DeltaChannel(merge_files)]
"""
__slots__ = ("value", "operator")
def __init__(
self,
operator: Callable[[list[Value], Any], list[Value]],
) -> None:
super().__init__(list)
self.operator = operator
self.value: list[Value] = []
def __eq__(self, other: object) -> bool:
if not isinstance(other, DeltaChannel):
return False
if (
self.operator.__name__ != "<lambda>"
and other.operator.__name__ != "<lambda>"
):
return self.operator is other.operator
return True
@property
def ValueType(self) -> Any:
return list[self.typ] # type: ignore[name-defined]
@property
def UpdateType(self) -> Any:
return self.typ | list[self.typ] # type: ignore[name-defined]
def copy(self) -> Self:
new = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
new.value = self.value if self.value is MISSING else self.value.copy()
return new
def from_checkpoint(self, checkpoint: Any) -> Self:
new = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
if checkpoint is MISSING:
try:
new.value = new.typ()
except Exception:
new.value = []
elif isinstance(checkpoint, list):
# Flat list of write values (oldest→newest) from get_channel_writes.
try:
value: Any = new.typ()
except Exception:
value = []
for write in checkpoint:
value = new.operator(value, write)
new.value = value
else:
# Backward compat: plain accumulated value (e.g. from a migrated thread).
try:
new.value = list(checkpoint)
except Exception:
new.value = []
return new
def update(self, values: Sequence[Any]) -> bool:
if not values:
return False
seen_overwrite = False
for value in values:
is_overwrite, overwrite_value = _get_overwrite(value)
if is_overwrite:
if seen_overwrite:
from langgraph.errors import (
ErrorCode,
InvalidUpdateError,
create_error_message,
)
msg = create_error_message(
message="Can receive only one Overwrite value per super-step.",
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
)
raise InvalidUpdateError(msg)
self.value = (
list(overwrite_value) if overwrite_value is not None else self.typ()
)
seen_overwrite = True
elif not seen_overwrite:
base = self.typ() if self.value is MISSING else self.value
self.value = self.operator(base, value)
return True
def get(self) -> list[Value]:
if self.value is MISSING:
raise EmptyChannelError()
return self.value
def is_available(self) -> bool:
return self.value is not MISSING
def checkpoint(self) -> DeltaChannelSentinel:
return DeltaChannelSentinel()
+20 -2
View File
@@ -47,9 +47,9 @@ from langgraph._internal._fields import (
from langgraph._internal._pydantic import create_model
from langgraph._internal._runnable import coerce_to_runnable
from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
from langgraph.channels.named_barrier_value import (
@@ -1671,12 +1671,30 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
for item in meta:
if isinstance(item, BaseChannel):
if isinstance(item, DeltaChannel) and hasattr(typ, "__origin__"):
outer = _strip_extras(typ.__origin__)
origin = typ.__origin__
# Unwrap parameterized Required[X]/NotRequired[X] to X
# (e.g. Annotated[NotRequired[dict[...]], ...]).
if hasattr(origin, "__origin__") and origin.__origin__ in (
Required,
NotRequired,
):
origin = origin.__args__[0]
outer = _strip_extras(origin)
if outer in (
collections.abc.Sequence,
collections.abc.MutableSequence,
):
outer = list
elif outer in (
collections.abc.Mapping,
collections.abc.MutableMapping,
):
outer = dict
elif outer in (
collections.abc.Set,
collections.abc.MutableSet,
):
outer = set
item.typ = outer
try:
item.value = outer()
+76 -8
View File
@@ -1,20 +1,19 @@
from __future__ import annotations
import logging
from collections.abc import Mapping
from datetime import datetime, timezone
from langgraph.checkpoint.base import Checkpoint
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph._internal._typing import MISSING
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.base import BaseChannel
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 4
logger = logging.getLogger(__name__)
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
@@ -61,8 +60,22 @@ def create_checkpoint(
def channels_from_checkpoint(
specs: Mapping[str, BaseChannel | ManagedValueSpec],
checkpoint: Checkpoint,
*,
saver: BaseCheckpointSaver | None = None,
config: RunnableConfig | None = None,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Get channels from a checkpoint."""
"""Hydrate channels from a checkpoint.
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
is sufficient — the stored value IS the reconstructed state.
`DeltaChannel` is the exception: its stored value is a sentinel; the
full state is spread across `checkpoint_writes` along the ancestor
chain. When `saver` and `config` are provided, this function fetches
that history via `saver._get_channel_writes_history` and folds it
through the channel's reducer. Without them (static contexts — graph
drawing, unit tests), delta channels fall back to empty.
"""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -70,10 +83,65 @@ def channels_from_checkpoint(
channel_specs[k] = v
else:
managed_specs[k] = v
channels: dict[str, BaseChannel] = {}
for k, v in channel_specs.items():
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
ch.after_checkpoint(checkpoint["channel_versions"].get(k), checkpoint.get("id"))
for k, spec in channel_specs.items():
ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING)
if (
isinstance(spec, DeltaChannel)
and saver is not None
and config is not None
and (stored is MISSING or stored is DELTA_SENTINEL)
):
# Target's own blob is empty/sentinel — walk ancestors for
# seed + writes. Skipping this when `stored` is a real value
# preserves state written via `update_state` or sitting at the
# tip of a pre-migration thread: the saver's ancestor walk
# intentionally excludes the target's own blob, so without
# this short-circuit we'd lose it.
history = saver._get_channel_writes_history(config, k)
delta_ch = spec.from_checkpoint(history.seed)
delta_ch.replay_writes(history.writes)
ch = delta_ch
else:
ch = spec.from_checkpoint(stored)
channels[k] = ch
return channels, managed_specs
async def achannels_from_checkpoint(
specs: Mapping[str, BaseChannel | ManagedValueSpec],
checkpoint: Checkpoint,
*,
saver: BaseCheckpointSaver | None = None,
config: RunnableConfig | None = None,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Async version of `channels_from_checkpoint`. See docstring there."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
if isinstance(v, BaseChannel):
channel_specs[k] = v
else:
managed_specs[k] = v
channels: dict[str, BaseChannel] = {}
for k, spec in channel_specs.items():
ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING)
if (
isinstance(spec, DeltaChannel)
and saver is not None
and config is not None
and (stored is MISSING or stored is DELTA_SENTINEL)
):
history = await saver._aget_channel_writes_history(config, k)
delta_ch = spec.from_checkpoint(history.seed)
delta_ch.replay_writes(history.writes)
ch = delta_ch
else:
ch = spec.from_checkpoint(stored)
channels[k] = ch
return channels, managed_specs
+10 -9
View File
@@ -92,6 +92,7 @@ from langgraph.pregel._algo import (
task_path_str,
)
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -891,12 +892,6 @@ class PregelLoop:
id=self.checkpoint["id"] if exiting else None,
updated_channels=self.updated_channels,
)
if do_checkpoint and self.channels:
for k, ch in self.channels.items():
ch.after_checkpoint(
self.checkpoint["channel_versions"].get(k),
self.checkpoint.get("id"),
)
# 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()
@@ -1279,7 +1274,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
)
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
@@ -1482,8 +1480,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
)
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
self.channels, self.managed = await achannels_from_checkpoint(
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
+31 -12
View File
@@ -122,6 +122,7 @@ from langgraph.pregel._algo import (
)
from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -1049,14 +1050,17 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
checkpoint = saved.checkpoint
channels, managed = channels_from_checkpoint(
self.channels,
checkpoint,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1169,14 +1173,17 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
checkpoint = saved.checkpoint
channels, managed = channels_from_checkpoint(
channels, managed = await achannels_from_checkpoint(
self.channels,
checkpoint,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1522,8 +1529,9 @@ class Pregel(
saved = checkpointer.get_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
@@ -1542,6 +1550,11 @@ class Pregel(
channels, managed = channels_from_checkpoint(
self.channels,
checkpoint,
saver=self.checkpointer
if saved is not None
and isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config if saved is not None else None,
)
values, as_node = updates[0][:2]
@@ -1967,8 +1980,9 @@ class Pregel(
saved = await checkpointer.aget_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
@@ -1984,9 +1998,14 @@ class Pregel(
)
if saved:
checkpoint_config = patch_configurable(config, saved.config[CONF])
channels, managed = channels_from_checkpoint(
channels, managed = await achannels_from_checkpoint(
self.channels,
checkpoint,
saver=self.checkpointer
if saved is not None
and isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config if saved is not None else None,
)
values, as_node = updates[0][:2]
# no values, just clear all tasks
+272 -56
View File
@@ -2,13 +2,17 @@ import operator
from collections.abc import Sequence
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph._internal._typing import MISSING
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.graph.message import add_messages
pytestmark = pytest.mark.anyio
@@ -121,9 +125,9 @@ def test_untracked_value() -> None:
def test_delta_channel_basic_two_steps() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DeltaChannelSentinel
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
@@ -131,12 +135,12 @@ def test_delta_channel_basic_two_steps() -> None:
# Step 1: one message added
ch.update([HumanMessage(content="hi", id="h1")])
d1 = ch.checkpoint()
assert isinstance(d1, DeltaChannelSentinel)
assert d1 is DELTA_SENTINEL
# Step 2: another message
ch.update([AIMessage(content="hello", id="a1")])
d2 = ch.checkpoint()
assert isinstance(d2, DeltaChannelSentinel)
assert d2 is DELTA_SENTINEL
# Full accumulated value is preserved in memory
assert len(ch.get()) == 2
@@ -145,20 +149,21 @@ def test_delta_channel_basic_two_steps() -> None:
def test_delta_channel_from_checkpoint_writes_list() -> None:
"""from_checkpoint with a flat list of individual writes replays them through the operator."""
"""replay_writes on a fresh channel replays through the operator."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
# Each element is one write value (as stored in checkpoint_writes)
writes = [
HumanMessage(content="hi", id="h1"),
AIMessage(content="hello", id="a1"),
HumanMessage(content="bye", id="h2"),
]
ch = spec.from_checkpoint(writes)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "messages", HumanMessage(content="hi", id="h1")),
("t1", "messages", AIMessage(content="hello", id="a1")),
("t2", "messages", HumanMessage(content="bye", id="h2")),
]
)
msgs = ch.get()
assert len(msgs) == 3
assert msgs[0].content == "hi"
@@ -169,7 +174,7 @@ def test_delta_channel_from_checkpoint_writes_list() -> None:
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
from langchain_core.messages import HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
# Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat
@@ -181,9 +186,9 @@ def test_delta_channel_from_checkpoint_backwards_compat() -> None:
def test_delta_channel_overwrite() -> None:
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import DeltaChannelSentinel
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
from langgraph.types import Overwrite
@@ -192,7 +197,7 @@ def test_delta_channel_overwrite() -> None:
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
d = ch.checkpoint()
assert isinstance(d, DeltaChannelSentinel)
assert d is DELTA_SENTINEL
# After overwrite, value is reset to only the new message
assert len(ch.get()) == 1
assert ch.get()[0].content == "new"
@@ -202,7 +207,7 @@ def test_delta_channel_remove_message_and_replay() -> None:
"""RemoveMessage must round-trip correctly when writes are replayed."""
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
@@ -221,12 +226,14 @@ def test_delta_channel_remove_message_and_replay() -> None:
assert ch.get() == [HumanMessage(content="hi", id="h1")]
# Replay the writes list from scratch — must reproduce the post-remove state
writes = [
HumanMessage(content="hi", id="h1"),
AIMessage(content="hello", id="a1"),
RemoveMessage(id="a1"),
]
ch2 = spec.from_checkpoint(writes)
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "messages", HumanMessage(content="hi", id="h1")),
("t1", "messages", AIMessage(content="hello", id="a1")),
("t2", "messages", RemoveMessage(id="a1")),
]
)
assert ch2.get() == [HumanMessage(content="hi", id="h1")]
@@ -234,7 +241,7 @@ def test_delta_channel_update_by_id_and_replay() -> None:
"""Updating a message by ID must round-trip correctly through writes replay."""
from langchain_core.messages import HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
@@ -248,29 +255,31 @@ def test_delta_channel_update_by_id_and_replay() -> None:
assert ch.get() == [HumanMessage(content="updated", id="h1")]
# Replay writes — must produce the updated message, not the original
writes = [
HumanMessage(content="original", id="h1"),
HumanMessage(content="updated", id="h1"),
]
ch2 = spec.from_checkpoint(writes)
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "messages", HumanMessage(content="original", id="h1")),
("t1", "messages", HumanMessage(content="updated", id="h1")),
]
)
assert len(ch2.get()) == 1
assert ch2.get()[0].content == "updated"
def test_delta_channel_checkpoint_returns_sentinel() -> None:
"""checkpoint() always returns DeltaChannelSentinel regardless of state."""
from langgraph.checkpoint.base import DeltaChannelSentinel
"""checkpoint() always returns DELTA_SENTINEL regardless of state."""
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
assert isinstance(ch.checkpoint(), DeltaChannelSentinel)
assert ch.checkpoint() is DELTA_SENTINEL
from langchain_core.messages import HumanMessage
ch.update([HumanMessage(content="hi", id="h1")])
assert isinstance(ch.checkpoint(), DeltaChannelSentinel)
assert ch.checkpoint() is DELTA_SENTINEL
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
@@ -281,7 +290,7 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
@@ -304,16 +313,13 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
# get_tuple must return a resolved list (not DeltaChannelSentinel)
from langgraph.checkpoint.base import DeltaChannelSentinel
# get_tuple returns raw storage shape — channel_values stores DELTA_SENTINEL
# for delta channels; the reconstructed writes flow separately via
# saver._get_channel_writes_history.
saved = saver.get_tuple(config)
assert saved is not None
assert "messages" in saved.checkpoint["channel_values"]
assert not isinstance(
saved.checkpoint["channel_values"]["messages"], DeltaChannelSentinel
)
assert isinstance(saved.checkpoint["channel_values"]["messages"], list)
assert saved.checkpoint["channel_values"]["messages"] is DELTA_SENTINEL
state = graph.get_state(config)
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
@@ -323,7 +329,7 @@ def _delta_channel_with_type(operator, typ):
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
from typing import Annotated
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.state import _get_channel
return _get_channel("_test", Annotated[typ, DeltaChannel(operator)])
@@ -343,7 +349,7 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None:
def test_delta_channel_dict_reducer_basic_updates() -> None:
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
from langgraph.checkpoint.base import DeltaChannelSentinel
from langgraph.checkpoint.base import DELTA_SENTINEL
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
@@ -352,25 +358,30 @@ def test_delta_channel_dict_reducer_basic_updates() -> None:
ch.update([{"a": 1}])
d1 = ch.checkpoint()
assert isinstance(d1, DeltaChannelSentinel)
assert d1 is DELTA_SENTINEL
ch.update([{"b": 2}])
d2 = ch.checkpoint()
assert isinstance(d2, DeltaChannelSentinel)
assert d2 is DELTA_SENTINEL
assert ch.get() == {"a": 1, "b": 2}
def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
"""from_checkpoint with a writes list replays correctly through a dict merge reducer."""
"""replay_writes on a fresh channel replays through a dict merge reducer."""
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
# Each element is one write value (oldest→newest)
writes = [{"a": 1}, {"b": 2}, {"c": 3}]
ch = spec.from_checkpoint(writes)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "files", {"a": 1}),
("t1", "files", {"b": 2}),
("t2", "files", {"c": 3}),
]
)
assert ch.get() == {"a": 1, "b": 2, "c": 3}
@@ -398,10 +409,215 @@ def test_delta_channel_dict_reducer_with_deletions() -> None:
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
# Confirm writes reconstruction produces the same result
writes = [
{"file1.py": "content1", "file2.py": "content2"},
{"file1.py": None, "file3.py": "content3"},
]
spec = _delta_channel_with_type(merge_files, dict)
ch2 = spec.from_checkpoint(writes)
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "files", {"file1.py": "content1", "file2.py": "content2"}),
("t1", "files", {"file1.py": None, "file3.py": "content3"}),
]
)
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
"""Overwrite(dict) in update() must preserve dict shape, not coerce to list."""
from langgraph.types import Overwrite
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
ch.update([{"a": 1}])
ch.update([Overwrite({"b": 2, "c": 3})])
assert ch.get() == {"b": 2, "c": 3}
def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
"""Overwrite(dict) embedded in replayed writes must reconstruct as dict."""
from langgraph.types import Overwrite
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "files", {"a": 1}),
("t1", "files", Overwrite({"x": 10, "y": 20})),
("t2", "files", {"z": 30}),
]
)
assert ch.get() == {"x": 10, "y": 20, "z": 30}
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`.
This is the shape the deepagents filesystem middleware uses for its
`files` field; without unwrapping NotRequired we'd fall through to `list`
and blow up on the first dict operator call.
"""
from typing import Annotated
from typing_extensions import NotRequired
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.state import _get_channel
def merge_dicts(left: dict | None, right: dict) -> dict:
if left is None:
return dict(right)
return {**left, **right}
annotation = Annotated[
NotRequired[dict[str, int]],
DeltaChannel(merge_dicts),
]
ch = _get_channel("files", annotation).from_checkpoint(MISSING)
assert ch.get() == {}
ch.update([{"a": 1}])
ch.update([{"b": 2}])
assert ch.get() == {"a": 1, "b": 2}
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel.
Mirrors the deepagents filesystem pattern: `files: Annotated[dict, reducer]`
where the reducer merges dicts and treats None values as deletions.
"""
from typing import Annotated
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
def merge_files(left: dict | None, right: dict) -> dict:
if left is None:
return {k: v for k, v in right.items() if v is not None}
result = {**left}
for k, v in right.items():
if v is None:
result.pop(k, None)
else:
result[k] = v
return result
class State(TypedDict):
files: Annotated[dict[str, str], DeltaChannel(merge_files)]
turn = {"v": 0}
def write_file(state: State) -> dict:
turn["v"] += 1
n = turn["v"]
return {"files": {f"/doc_{n}.txt": f"content for turn {n}"}}
builder = StateGraph(State)
builder.add_node("write_file", write_file)
builder.add_edge(START, "write_file")
saver = InMemorySaver()
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "fs"}}
for _ in range(3):
graph.invoke({"files": {}}, config)
# Checkpoint stores only the sentinel — per-step writes live in checkpoint_writes.
saved = saver.get_tuple(config)
assert saved is not None
cv = saved.checkpoint["channel_values"]["files"]
assert cv is DELTA_SENTINEL
state = graph.get_state(config)
assert state.values["files"] == {
"/doc_1.txt": "content for turn 1",
"/doc_2.txt": "content for turn 2",
"/doc_3.txt": "content for turn 3",
}
# Deletion path must round-trip through writes replay.
def delete_file(state: State) -> dict:
return {"files": {"/doc_1.txt": None}}
builder2 = StateGraph(State)
builder2.add_node("write_file", write_file)
builder2.add_node("delete_file", delete_file)
builder2.add_edge(START, "write_file")
builder2.add_edge("write_file", "delete_file")
turn["v"] = 0
saver2 = InMemorySaver()
graph2 = builder2.compile(checkpointer=saver2)
config2 = {"configurable": {"thread_id": "fs2"}}
graph2.invoke({"files": {}}, config2)
state2 = graph2.get_state(config2)
assert state2.values["files"] == {}
def test_delta_channel_dict_reducer_backwards_compat() -> None:
"""A pre-DeltaChannel dict checkpoint must load as a dict, not be listified."""
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
old_value = {"a": 1, "b": 2}
ch = spec.from_checkpoint(old_value)
assert ch.get() == {"a": 1, "b": 2}
# ---------------------------------------------------------------------------
# seed / pre-delta migration
# ---------------------------------------------------------------------------
def test_delta_channel_from_checkpoint_honors_seed() -> None:
"""A non-sentinel value to from_checkpoint is used as the pre-delta seed.
Guards the pre-delta migration path: when the saver's ancestor walk hits
a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs
the post-migration state correctly rather than replaying from empty.
"""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="pre-delta", id="p1")]
ch = spec.from_checkpoint(seed)
ch.replay_writes(
[
("t0", "messages", AIMessage(content="delta-1", id="d1")),
("t1", "messages", HumanMessage(content="delta-2", id="d2")),
]
)
msgs = ch.get()
assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"]
def test_delta_channel_from_checkpoint_seed_without_writes() -> None:
"""Reconstruction at a pre-delta ancestor with no newer deltas returns
just the seed — the saver's terminator fired immediately."""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="only-snap", id="s1")]
ch = spec.from_checkpoint(seed)
ch.replay_writes([])
assert ch.get() == seed
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`
explicitly should feed None to the reducer as the left operand.
"""
def replace(left, right):
return right
spec = DeltaChannel(replace)
ch = spec.from_checkpoint(None)
ch.replay_writes([("t0", "x", "after")])
# Reducer replaces; seed=None → first write produces "after".
assert ch.get() == "after"
@@ -26,7 +26,7 @@ from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
@@ -54,7 +54,9 @@ except ImportError:
_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."
"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 = (
@@ -221,6 +223,10 @@ def _approx_tokens(n_turns: int) -> str:
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
TURN_COUNTS = [10, 25, 50, 100, 500]
# Deep-thread counts where add_messages blob storage would exceed 1 GB;
# only DeltaChannel runs here.
DELTA_ONLY_TURN_COUNTS = [1000]
def _checkpointer_factories() -> list[tuple[str, Any]]:
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
@@ -242,7 +248,7 @@ def run_benchmark() -> None:
import psycopg
psycopg.connect(_POSTGRES_URI).close()
checkpointers.append(("Postgres (recursive CTE)", "postgres"))
checkpointers.append(("Postgres (plain SELECT)", "postgres"))
except Exception:
pass
@@ -276,36 +282,49 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
with SqliteSaver.from_conn_string(f.name) as saver:
yield saver
rows = []
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
for turns in 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))
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))
# ── Table 1: Storage ─────────────────────────────────────────────────────
W = 70
W = 64
print("Storage (checkpoint blob bytes)")
print("=" * W)
print(
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}"
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} "
f"{'savings':>8}"
)
print("-" * W)
storage_results = []
for turns, b_bytes, d_bytes, b_rt, d_rt in rows:
if b_bytes < 0:
print(
f"{turns:>6} {_approx_tokens(turns):>10} {'n/a':>12} {'n/a':>12} {'n/a':>8}"
)
def _bytes_or_na(v: Any) -> str:
if v is None:
return "n/a"
if 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"
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")
storage_results.append((turns, b_bytes, d_bytes, ratio))
print(
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} "
f"{ratio:>7.0f}x"
)
ratio_str = f"{ratio:.0f}x"
print(
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} "
f"{ratio_str:>8}"
)
print("=" * W)
print()
@@ -314,29 +333,36 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
print("=" * W)
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
print("-" * W)
for turns, b_bytes, d_bytes, b_rt, d_rt in rows:
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"{b_rt * 1000:>10.1f}ms {d_rt * 1000:>10.1f}ms"
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
)
print("=" * W)
print()
if storage_results:
turns, b_bytes, d_bytes, ratio = storage_results[-1]
b_rt = rows[-1][-2]
d_rt = rows[-1][-1]
# ── Table 3: Per-invoke latency (total write_elapsed / turns) ─────────────
print("Per-invoke latency (total graph.invoke time / turns)")
print("=" * W)
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
print("-" * W)
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
def _per(wt: Any) -> str:
if wt is None:
return "n/a"
return f"{(wt / turns) * 1000:.1f}ms"
print(
f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(d_bytes)} ({ratio:.0f}x less storage); "
f"read {b_rt * 1000:.1f}ms → {d_rt * 1000:.1f}ms"
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_per(b_wt):>12} {_per(d_wt):>12}"
)
print()
print("=" * W)
print()
print("Legend:")
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
print(
" delta = DeltaChannel(add_messages) — O(N) storage, full chain replay"
)
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
print(" delta = Annotated[list, DeltaChannel(add_messages)] — O(N) storage")
print()
@@ -0,0 +1,504 @@
"""Tests for the BinaryOperatorAggregate -> DeltaChannel migration path.
A thread written under `BinaryOperatorAggregate(...)` must keep working
after its annotation is swapped to `DeltaChannel(...)` on the same
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.
Scenarios covered:
1. **Basic migration (sync + async)**: build pre-migration state with
`BinaryOperatorAggregate`, swap the annotation to `DeltaChannel` on
the same checkpointer, and verify that every settled pre-migration
super-step boundary (`next=('__start__',)`) round-trips exactly
under the delta-channel view.
2. **Time travel into a pre-migration checkpoint** after migration —
`graph.get_state(pre_migration_config)` at a settled ancestor
returns the same state as under the binop channel.
3. **Continuing a migrated thread**: driving one more super-step after
migration produces a state that includes the pre-migration settled
prefix plus the new delta write — proving `from_checkpoint(seed)` +
`replay_writes` correctly fold post-migration deltas onto the
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.
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.
The parent-chain walk is scoped to the thread.
TODO: add postgres variants in the existing `libs/checkpoint-postgres`
test files (different fixture setup; not this file).
"""
from __future__ import annotations
import operator
from typing import Annotated, Any
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.graph import END, START, StateGraph
pytestmark = pytest.mark.anyio
# ---------------------------------------------------------------------------
# Graph factories
#
# A minimal reducer (`operator.add` on lists of str) with a noop node keeps
# state change localized to the HumanMessage-like payload passed through
# `invoke`. That isolates the pre/post-migration parity assertions to
# channel-hydration semantics.
# ---------------------------------------------------------------------------
def _noop(_state: Any) -> dict:
return {}
def _binop_graph(checkpointer: Any) -> Any:
class BinopState(TypedDict):
items: Annotated[list, BinaryOperatorAggregate(list, operator.add)]
return (
StateGraph(BinopState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def _delta_graph(checkpointer: Any) -> Any:
class DeltaState(TypedDict):
items: Annotated[list, DeltaChannel(operator.add)]
return (
StateGraph(DeltaState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def _drive(graph: Any, config: dict, tag: str, n: int) -> None:
for i in range(n):
graph.invoke({"items": [f"{tag}{i}"]}, config)
async def _adrive(graph: Any, config: dict, tag: str, n: int) -> None:
for i in range(n):
await graph.ainvoke({"items": [f"{tag}{i}"]}, config)
def _settled_boundaries(history: list) -> list[tuple[dict, list]]:
"""Return `[(config, items), ...]` for every checkpoint in `history`
whose `next == ('__start__',)` — the stable boundaries between invokes.
"""
return [
(s.config, list(s.values.get("items", [])))
for s in history
if s.next == ("__start__",)
]
# ---------------------------------------------------------------------------
# 1. Basic migration (sync + async)
# ---------------------------------------------------------------------------
def test_basic_migration_preserves_pre_migration_state() -> None:
"""Build state under `BinaryOperatorAggregate`, migrate to
`DeltaChannel` on the same checkpointer, and verify that every
settled pre-migration super-step boundary round-trips exactly.
Settled boundaries (`next=('__start__',)`) are the stable hydration
targets for the migration path: writes that produced the NEXT
super-step are kept as `pending_writes` on the ancestor, so walking
from a descendant finds the ancestor's blob as the seed and
reconstructs the correct state.
"""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "basic-sync"}}
# Pre-migration: accumulate items across 3 invokes.
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
assert len(pre_boundaries) >= 2, "expected multiple settled boundaries"
# Migrate: swap the annotation on the same checkpointer.
delta = _delta_graph(checkpointer)
for cfg, items in pre_boundaries:
snap = delta.get_state(cfg)
assert list(snap.values.get("items", [])) == items, (
f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: "
f"expected {items}, got {snap.values.get('items', [])}"
)
async def test_basic_migration_preserves_pre_migration_state_async() -> None:
"""Async variant of the basic migration scenario."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "basic-async"}}
binop = _binop_graph(checkpointer)
await _adrive(binop, config, "u", 3)
pre_history = [s async for s in binop.aget_state_history(config)]
pre_boundaries = _settled_boundaries(pre_history)
assert len(pre_boundaries) >= 2
delta = _delta_graph(checkpointer)
for cfg, items in pre_boundaries:
snap = await delta.aget_state(cfg)
assert list(snap.values.get("items", [])) == items, (
f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}"
)
# ---------------------------------------------------------------------------
# 2. Time travel into a pre-migration checkpoint after migration
# ---------------------------------------------------------------------------
def test_time_travel_into_pre_migration_checkpoint() -> None:
"""After migration, `graph.get_state(pre_migration_config)` at a
settled ancestor returns the state as stored at that point."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "time-travel"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
assert pre_boundaries, "no settled ancestors to time-travel to"
delta = _delta_graph(checkpointer)
# Pick the oldest non-empty boundary — a long distance to walk back.
non_empty = [(cfg, items) for cfg, items in pre_boundaries if items]
assert non_empty, "expected at least one non-empty boundary"
target_cfg, expected_items = non_empty[-1]
snap = delta.get_state(target_cfg)
assert list(snap.values.get("items", [])) == expected_items
# ---------------------------------------------------------------------------
# 3. Continuing a migrated thread: deltas fold onto pre-migration seed
# ---------------------------------------------------------------------------
def test_continuing_migrated_thread_folds_deltas_on_seed() -> None:
"""Resume a pre-migration settled ancestor via `invoke(None, cfg)`
under the delta-channel graph. Since the pre-migration checkpoint
has an existing `pending_writes` entry (the input for the NEXT
super-step), re-running from that ancestor reproduces the same
post-ancestor state as the original binop run.
This proves the seed-terminator + write-replay pipeline works
end-to-end across the migration boundary.
"""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "continue"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 2)
# Pick the oldest settled boundary with non-empty state.
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
target_cfg, seed_items = next(
(cfg, items) for cfg, items in reversed(pre_boundaries) if items
)
assert seed_items, "need a non-empty seed boundary"
# Migrate and resume from the pre-migration ancestor. `invoke(None,
# cfg)` replays the pending writes staged at `cfg` under the new
# channel; the reducer folds those deltas onto the seed.
delta = _delta_graph(checkpointer)
result = delta.invoke(None, target_cfg)
# The resumed state must include the pre-migration seed items in order.
result_items = list(result.get("items", []))
for idx, prefix_item in enumerate(seed_items):
assert result_items[idx] == prefix_item, (
f"pre-migration seed item at {idx} not preserved: "
f"got {result_items[: idx + 1]}, expected {seed_items}"
)
# ---------------------------------------------------------------------------
# 4. Base-saver fallback path
# ---------------------------------------------------------------------------
class _ThirdPartyStyleSaver(InMemorySaver):
"""Simulates a third-party saver that inherits the reference
`_get_channel_writes_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
still the in-memory one.
"""
# MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...]
_get_channel_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
)
_aget_channel_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
)
def test_base_saver_fallback_matches_optimized_override() -> None:
"""The reference `BaseCheckpointSaver` implementation must produce
the same migration behavior as the optimized `InMemorySaver`
override. We drive the same migration scenario through both savers
and assert per-snapshot parity in the delta-channel view."""
# Fast path: optimized InMemorySaver override.
fast_saver = InMemorySaver()
fast_config = {"configurable": {"thread_id": "fast"}}
fast_binop = _binop_graph(fast_saver)
_drive(fast_binop, fast_config, "u", 3)
fast_delta = _delta_graph(fast_saver)
fast_history = [
(s.next, list(s.values.get("items", [])))
for s in fast_delta.get_state_history(fast_config)
]
# Slow path: base-class fallback.
slow_saver = _ThirdPartyStyleSaver()
slow_config = {"configurable": {"thread_id": "slow"}}
slow_binop = _binop_graph(slow_saver)
_drive(slow_binop, slow_config, "u", 3)
slow_delta = _delta_graph(slow_saver)
slow_history = [
(s.next, list(s.values.get("items", [])))
for s in slow_delta.get_state_history(slow_config)
]
assert slow_history == fast_history, (
"base-saver fallback should match optimized-override behavior; "
f"fast={fast_history}, slow={slow_history}"
)
# ---------------------------------------------------------------------------
# 5. Thread isolation under mixed-generation storage
# ---------------------------------------------------------------------------
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.
"""
checkpointer = InMemorySaver()
migrated_cfg = {"configurable": {"thread_id": "migrated"}}
fresh_cfg = {"configurable": {"thread_id": "fresh"}}
# Thread A: pre-migration build-up.
binop = _binop_graph(checkpointer)
_drive(binop, migrated_cfg, "m", 2)
# Thread B: fresh delta-channel run.
delta = _delta_graph(checkpointer)
_drive(delta, fresh_cfg, "f", 2)
# Thread A: migrate and confirm its state is anchored in its own
# thread's pre-migration history (tag 'm'), never mixing in tag 'f'.
migrated_boundaries = _settled_boundaries(
list(delta.get_state_history(migrated_cfg))
)
assert migrated_boundaries, "migrated thread has no settled boundaries"
for _, items in migrated_boundaries:
for it in items:
assert it.startswith("m"), (
f"migrated thread leaked item from other thread: {it}"
)
# Thread B: settled boundaries must only contain 'f' tags.
fresh_boundaries = _settled_boundaries(list(delta.get_state_history(fresh_cfg)))
assert fresh_boundaries, "fresh thread has no settled boundaries"
for _, items in fresh_boundaries:
for it in items:
assert it.startswith("f"), (
f"fresh thread leaked item from migrated thread: {it}"
)
# ---------------------------------------------------------------------------
# 6. Tip-of-pre-migration hydration: the latest checkpoint from a binop-run
# thread has a real accumulated value in its own `channel_values["items"]`.
# When hydrated under the delta-channel graph via `get_state(config)` with no
# `checkpoint_id`, the short-circuit must use that value directly instead of
# walking ancestors (which would skip the tip's own blob).
# ---------------------------------------------------------------------------
def test_tip_of_pre_migration_hydrates_directly() -> None:
"""`graph.get_state(config)` at the latest (pre-migration) checkpoint
returns the full accumulated list stored in that checkpoint's own
`channel_values`. The hydration must not walk ancestors past it."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "tip-sync"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
binop_tip = binop.get_state(config)
expected_items = list(binop_tip.values.get("items", []))
assert expected_items == ["u0", "u1", "u2"], (
f"sanity: pre-migration tip should accumulate all 3 items, got {expected_items}"
)
delta = _delta_graph(checkpointer)
snap = delta.get_state(config)
assert list(snap.values.get("items", [])) == expected_items, (
f"tip hydration mismatch: expected {expected_items}, "
f"got {snap.values.get('items', [])}"
)
async def test_tip_of_pre_migration_hydrates_directly_async() -> None:
"""Async variant of the tip-of-pre-migration hydration scenario."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "tip-async"}}
binop = _binop_graph(checkpointer)
await _adrive(binop, config, "u", 3)
binop_tip = await binop.aget_state(config)
expected_items = list(binop_tip.values.get("items", []))
assert expected_items == ["u0", "u1", "u2"]
delta = _delta_graph(checkpointer)
snap = await delta.aget_state(config)
assert list(snap.values.get("items", [])) == expected_items, (
f"async tip hydration mismatch: expected {expected_items}, "
f"got {snap.values.get('items', [])}"
)
# ---------------------------------------------------------------------------
# 7. `update_state` after migration writes a real value to the new
# checkpoint's `channel_values` (not a sentinel). Hydration must use it
# directly — the ancestor walk would skip this blob and return stale state.
# ---------------------------------------------------------------------------
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."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "update-state"}}
# Pre-migration: accumulate a little state.
binop = _binop_graph(checkpointer)
_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`.
delta = _delta_graph(checkpointer)
delta.invoke({"items": ["post"]}, config)
# `update_state` writes a concrete value into a new checkpoint's blob
# via the reducer against the hydrated prior state.
delta.update_state(config, {"items": ["x", "y"]})
snap = delta.get_state(config)
updated_items = list(snap.values.get("items", []))
# Must include the "x","y" update; without the hydration fix, the
# update_state-written blob would be skipped in favor of an ancestor
# walk, and the update values would disappear.
assert "x" in updated_items and "y" in updated_items, (
f"update_state values missing from snapshot: {updated_items}"
)
# The "x","y" items should be folded onto the prior accumulated state,
# not stand alone. This verifies the update-written blob is used
# directly by `get_state` (no ancestor walk past it).
assert len(updated_items) >= 4, (
f"update_state snapshot should preserve pre-update state, got {updated_items}"
)
assert updated_items[-2:] == ["x", "y"], (
f"update_state deltas should be at the tail, got {updated_items}"
)
# ---------------------------------------------------------------------------
# 8. Fork from an `update_state` checkpoint: a new run branched off the
# update_state-produced checkpoint must see that checkpoint's concrete
# `channel_values` as its base, with new deltas folded on top.
# ---------------------------------------------------------------------------
def test_fork_from_update_state_checkpoint() -> None:
"""Branching a new run from the checkpoint produced by `update_state`
must use that checkpoint's concrete blob as the base. Additional
deltas from the forked run fold onto it through the reducer."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "fork"}}
# Pre-migration build-up, then migrate and add one post-migration step.
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 2)
delta = _delta_graph(checkpointer)
delta.invoke({"items": ["post"]}, config)
# Apply `update_state` and capture the returned config (references
# the new checkpoint produced by the update).
update_cfg = delta.update_state(config, {"items": ["x", "y"]})
update_snap = delta.get_state(update_cfg)
base_items = list(update_snap.values.get("items", []))
assert "x" in base_items and "y" in base_items, (
f"update_state values missing from snapshot: {base_items}"
)
assert base_items[-2:] == ["x", "y"], (
f"sanity: update_state deltas should be at the tail, got {base_items}"
)
# Fork: invoke from the update_state checkpoint with a new delta.
forked = delta.invoke({"items": ["fork0"]}, update_cfg)
forked_items = list(forked.get("items", []))
# The fork must see the update_state-written blob as its base (not
# walk past it), and the new delta must fold on top of it.
assert forked_items[: len(base_items)] == base_items, (
f"fork lost update_state base: base={base_items}, forked={forked_items}"
)
assert forked_items[-1] == "fork0", f"fork delta not appended: {forked_items}"
+4 -4
View File
@@ -9407,7 +9407,7 @@ async def test_delta_channel_end_to_end_inmemory() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
@@ -9449,7 +9449,7 @@ async def test_delta_channel_time_travel() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
@@ -9507,7 +9507,7 @@ async def test_delta_channel_remove_message_end_to_end() -> None:
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
@@ -9554,7 +9554,7 @@ async def test_delta_channel_update_by_id_end_to_end() -> None:
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
+1 -3
View File
@@ -1161,9 +1161,7 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
assert called == ["step_a", "ask_human"]
# Resume with explicit head checkpoint_id in config
head_checkpoint_id = graph.get_state(config).config["configurable"][
"checkpoint_id"
]
head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"]
called.clear()
resume_config = {
"configurable": {
+137
View File
@@ -0,0 +1,137 @@
# Delta-channel reconstruction: query strategy benchmark
**Branch:** `delta-channel-writes-based`
**Question (Nuno):** Is the recursive CTE the right query shape for reconstructing a delta channel inside `get_tuple`, or would a plain `SELECT WHERE` be cheaper even though it returns more rows?
**Answer:** Plain `SELECT WHERE` wins at every realistic depth. The recursion isn't the problem — the JSON-expression join inside the CTE is.
## Setup
- Postgres 16 on `localhost:5441` (the `compose-postgres.yml` instance, run directly without docker for this round)
- Single delta channel `messages`, one write per checkpoint, `DELTA_SENTINEL` blob per checkpoint
- Linear chain (`branch=1`) and 5-way branching at every step (`branch=5`) — branching is the case where plain over-fetches sibling rows
- Median of 20 timed runs after 3 warmups, fresh psycopg cursor per strategy
- Bench script: `bench_get_tuple_strategies.py` at repo root
Three strategies compared:
| name | roundtrips | shape |
|------|-----------|-------|
| `cte` | 1 | Current prod: recursive CTE walks ancestors, LEFT JOINs writes + blobs |
| `plain` | 3 | Nuno's suggestion: thread-wide `SELECT WHERE` per table, Python walks parent chain and filters |
| `cte+narrow` | 2 | CTE returns ancestor IDs only, then one `UNION ALL` of writes + blobs filtered by `ANY(ids)` |
## Results (ms per get_tuple, median of 20)
```
depth branch cte plain cte+narrow rows_cte rows_plain plain/cte
10 1 0.14ms 0.26ms 0.24ms 9 30 1.89x
10 5 0.21ms 0.23ms 0.17ms 9 110 1.11x
50 1 0.89ms 0.27ms 0.33ms 49 150 0.31x
50 5 2.35ms 0.66ms 0.51ms 49 550 0.28x
200 1 11.61ms 0.78ms 1.30ms 199 600 0.07x
200 5 34.79ms 2.33ms 3.07ms 199 2200 0.07x
1000 1 274.60ms 2.59ms 13.29ms 999 3000 0.01x
1000 5 856.01ms 10.14ms 15.31ms 999 11000 0.01x
```
Lower is better. `plain/cte < 1` means plain is faster.
### Headline numbers
- depth 50: plain is **3x** faster
- depth 200: plain is **15x** faster
- depth 1000: plain is **~100x** faster
- Branching makes plain over-fetch (3000 rows → 11000 rows at d=1000), but it remains ~85x faster than the CTE
## Why the CTE collapses
`EXPLAIN (ANALYZE, BUFFERS)` of the CTE at depth 1000 (linear). Excerpt with the load-bearing nodes:
```
Sort ... actual time=137.798..137.827 rows=999
CTE ancestors
-> Recursive Union ... actual time=0.005..2.443 rows=999
^^^^^^
recursion is 2.4 ms — fine
-> Nested Loop Left Join ... actual time=2.676..137.529 rows=999
Join Filter: (cw.checkpoint_id = a.cid)
Rows Removed by Join Filter: 998001
^^^^^^^
999 ancestors x ~1000 writes
-> Nested Loop Left Join ... actual time=2.669..85.061 rows=999
Join Filter: (bl.version = ((c.checkpoint -> 'channel_versions'::text) ->> bl.channel))
Rows Removed by Join Filter: 998001
^^^^^^^
same quadratic blow-up on the blob join
```
Two pathological things are happening:
1. **The blob join filter is on a JSON expression**: `bl.version = (c.checkpoint -> 'channel_versions' ->> bl.channel)`. The planner cannot push this into an index lookup, so it materializes `checkpoint_blobs` for the thread and does a nested-loop comparison against every ancestor — a Cartesian product that grows as `O(ancestors × blobs_in_thread)`.
2. **The writes join is similar**: writes for the thread are materialized once, then nested-loop joined against ancestors with a `Join Filter` rather than a hash/merge join over the indexed `checkpoint_id`.
At depth 1000 that's **~2 million rows evaluated, 99.9% of them discarded**. The recursion itself is a rounding error.
For comparison, the plain Q1 (`SELECT … FROM checkpoints WHERE thread_id=? AND checkpoint_ns=?`) at depth 1000:
```
Seq Scan on checkpoints ... actual time=0.012..0.121 rows=1000
Execution Time: 0.140 ms
```
A simple seq scan over 57 buffers. Q2 and Q3 follow the same shape and complete in well under 1 ms each.
## Crossover and remote-DB reasoning
- Pure local Postgres: plain wins from depth ~30 onward; CTE wins by fractions of a ms below that
- Remote Postgres at ~5 ms RTT adds ~10 ms to plain (3 roundtrips vs 1). Crossover shifts to ~depth 30. Above that, the CTE's quadratic SQL cost still dominates the RTT savings.
There is no realistic conversation depth where the CTE wins on a remote DB. At depth 200+ (anything resembling a real multi-turn agent run) plain is faster regardless of network.
## Recommendation
**Switch to plain SELECT WHERE, one delta channel at a time.**
Three indexed queries per delta channel:
```sql
-- Q1: parent chain + per-checkpoint version of this channel
SELECT checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> 'channel_name' AS ver
FROM checkpoints
WHERE thread_id = ? AND checkpoint_ns = ?;
-- Q2: writes for this channel, anywhere in the thread
SELECT checkpoint_id, type, blob, task_id, idx
FROM checkpoint_writes
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
-- Q3: blobs for this channel, anywhere in the thread
SELECT version, type, blob
FROM checkpoint_blobs
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
```
Python then:
- Builds `parent_of: dict[cid, parent_cid]` from Q1
- Walks from target's parent newest → oldest
- Filters Q2 rows by `ancestor_set`, processes oldest → newest, applies overwrite-terminator
- Picks seed blob via the per-ancestor `ver` map, terminates at first non-sentinel blob
All O(n) on n = thread checkpoints, with tight constants (dict lookups). No recursion, no JSON-expression joins, no quadratic plans.
If the 3-roundtrip cost ever shows up on remote-DB benchmarks, fold Q2 + Q3 into one `UNION ALL` to get back to 2 roundtrips. Bench says it isn't worth the SQL complexity right now.
## Bonus: code simplification from single-channel scope
Multi-channel reconstruction in the current `_reconstruct_delta_channels_cur` carries:
- `rows_by_cid` nested dicts, keyed by cid then channel
- `seen_blob: set[(cid, channel)]` and `seen_write: set[(cid, channel, task_id, idx)]` dedup
- `collected: dict[channel, list]`, `done: set[channel]`, `seeds: dict[channel, value]`
- Inner `for ch in channels_list` loops and an early-exit `if len(done) == len(channels_list)`
Single-channel collapses these to a single list, a single bool, and one `Optional[Any]`. Roughly half the Python in that function, plus an obvious shape for splitting pure post-processing into `base.py` so sync and async stop duplicating it.
If multi-channel coalescing turns out to matter later, it can come back as a SQL-level optimization without re-introducing the bookkeeping in Python.
-133
View File
@@ -1,133 +0,0 @@
# feat(channels): DeltaChannel — O(N) incremental checkpoint storage
## The problem
LangGraph checkpoints store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, that means each checkpoint blob contains the entire conversation history up to that point.
Storage cost grows **O(N²)** in the number of turns:
| Step | Checkpoint blob |
|------|----------------|
| 1 | [msg_1] |
| 2 | [msg_1, msg_2] |
| N | [msg_1, ..., msg_N] |
At 100K tokens of conversation data, a single thread accumulates ~250 MB; with large messages or file attachments costs scale even faster.
## The fix: `DeltaChannel`
`DeltaChannel` is an opt-in wrapper around any binary reducer that stores only a **sentinel marker** in `checkpoint_blobs` rather than the full accumulated value. The actual per-step writes stay in `checkpoint_writes` (which every checkpointer already writes unconditionally). At read time the saver walks the ancestor chain, collects all writes for the channel, and replays them through the reducer.
Storage scales **O(N)** — the sentinel blob is effectively zero bytes, and the writes table already exists.
```python
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
class State(TypedDict):
# Before: O(N²) storage
messages: Annotated[list[AnyMessage], add_messages]
# After: O(N) storage
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
```
## Benchmarks
Simulated with realistic paragraph-length messages (~100 tokens each, ~400 chars). Each turn = one human + one AI message (~200 tokens total).
### Storage (InMemorySaver)
| turns | ctx | add_msgs | delta | savings |
|------:|----:|---------:|------:|--------:|
| 10 | ~2K tok | 108.6 KB | 4.0 KB | 27x |
| 25 | ~5K tok | 649.0 KB | 10.1 KB | 64x |
| 50 | ~10K tok | 2.6 MB | 20.2 KB | 126x |
| 100 | ~20K tok | 10.2 MB | 40.5 KB | 251x |
| 500 | ~100K tok | 252.6 MB | 202.8 KB | 1245x |
Savings grow with N because `add_messages` is O(N²) while `DeltaChannel` is O(N). The sentinel blob itself is essentially zero bytes.
### Read latency (avg of 5 `get_state` calls = cost per `invoke`)
| turns | ctx | add_msgs | delta |
|------:|----:|---------:|------:|
| 10 | ~2K tok | 0.1ms | 0.2ms |
| 25 | ~5K tok | 0.2ms | 0.5ms |
| 50 | ~10K tok | 0.4ms | 1.5ms |
| 100 | ~20K tok | 0.7ms | 4.8ms |
| 500 | ~100K tok | 5.8ms | 114.9ms |
**This cost is paid once per `invoke`/`stream` call, not per node.** Within a single invocation, all channels are loaded into memory once at the start and shared across every node — there is no per-node reconstruction. The 114.9ms at 500 turns is what you pay each time a user sends a new message, not on each step of the graph.
## How it works
**Write:** `DeltaChannel.checkpoint()` always emits `DeltaChannelSentinel()` — a tiny marker (zero payload bytes) stored in `checkpoint_blobs`. Per-step writes flow into `checkpoint_writes` as they normally do for every channel.
**Read:** The saver detects `DeltaChannelSentinel` values in `channel_values` and replaces them by calling `get_channel_writes` / `aget_channel_writes`, which walks the ancestor checkpoint chain and collects all writes for that channel (oldest→newest). `DeltaChannel.from_checkpoint()` replays those writes through the operator to reconstruct the full value.
**Saver implementations:**
- `InMemorySaver` — direct dict traversal of `self.storage` and `self.writes`, no I/O
- `PostgresSaver` (sync + async) — two queries: one cheap ID walk across the thread, one `ANY()` fetch of writes; no recursive CTE
- All other savers — `BaseCheckpointSaver.get_channel_writes` fallback via `list()`, with a re-entrancy guard to prevent infinite recursion
## Changes
**`libs/checkpoint`**
- `base/__init__.py` — add `DeltaChannelSentinel` marker dataclass; add `get_channel_writes` / `aget_channel_writes` to `BaseCheckpointSaver` with a `list()`-based fallback and re-entrancy guard
**`libs/checkpoint/memory`**
- `memory/__init__.py` — `get_channel_writes` via direct dict traversal; `_resolve_delta_channels` helper called in `get_tuple` / `aget_tuple` to replace sentinels with reconstructed write lists
**`libs/langgraph`**
- `channels/delta.py` — `DeltaChannel` implementation: `checkpoint()` always emits sentinel, `from_checkpoint()` replays writes list
- `channels/__init__.py` — export `DeltaChannel`
- `graph/state.py` — recognize `DeltaChannel` as a valid channel annotation
- `pregel/_checkpoint.py` / `pregel/_loop.py` — wire `after_checkpoint` hook; call it after each checkpointing step so `DeltaChannel` can advance internal state
**`libs/checkpoint-postgres`**
- `postgres/base.py` — `_get_channel_writes_cur` two-query ancestor walk (sync); `_resolve_delta_channels` called after `_load_blobs`
- `postgres/aio.py` — `_aget_channel_writes_cur` (async counterpart)
## Open questions
**Should we add a compile-time capability check?**
Currently misconfiguring `DeltaChannel` with an unsupported saver only errors at runtime on first reload. A protocol-based check at `compile()` time would give an early warning without requiring a manual boolean flag.
**`snapshot_every` for bounded reconstruction cost?**
Both per-invoke read latency and total write wall time grow O(N) per invoke / O(N²) total as the conversation lengthens. A `snapshot_every` parameter — periodically store a full snapshot in `checkpoint_blobs` to cap chain depth — would bound reconstruction cost and is a natural follow-up once the core design is stable.
## Backwards compatibility
| Scenario | Behaviour |
|----------|-----------|
| Existing graph using `add_messages` | Unaffected — no code or schema changes |
| `DeltaChannel` loading an old full-list checkpoint blob | Handled via backwards-compat path in `from_checkpoint` |
| `DeltaChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
| Time-travel to a past checkpoint | Ancestor walk uses the version at that checkpoint — correct by construction |
| `Overwrite` value | Resets the effective chain; reconstruction starts from that step |
## Test plan
- [x] `DeltaChannel` unit tests: `update` → `checkpoint` lifecycle, `from_checkpoint` chain replay, backwards-compat with plain list, `Overwrite` resets chain
- [x] `InMemorySaver` `get_channel_writes`: assembles write list from dict storage
- [x] Serde round-trip for `DeltaChannelSentinel`
- [x] End-to-end graph tests: multi-turn conversations accumulate correctly, time-travel reconstructs correct partial history
- [x] `PostgresSaver` two-query chain reconstruction (sync + async)
- [x] `BaseCheckpointSaver` fallback path via `list()` with re-entrancy guard
- [x] Storage benchmark: `DeltaChannel` uses strictly less storage than `add_messages` at all measured turn counts
---
## Changes from previous base branch
The previous version stored `DeltaValue` objects (containing the per-step writes) directly in `checkpoint_blobs` and used a `DeltaChainValue` to represent the assembled chain. Reconstruction required a dedicated `get_delta_chain` / `aget_delta_chain` protocol and a recursive CTE in Postgres.
This version pivots to a simpler design:
- **Sentinel in blobs, writes in `checkpoint_writes`** — `checkpoint_blobs` stores only a zero-byte `DeltaChannelSentinel` marker. The actual per-step data already lives in `checkpoint_writes` (written unconditionally by every checkpointer), so blob storage is essentially free. This is why storage savings jump to 1245x at 500 turns.
- **No custom serde type for the delta payload** — `DeltaValue` / `DeltaChainValue` and the `"delta"` serde type tag are gone. Writes are deserialized with the same serde path they were originally written with.
- **Postgres: two queries instead of a recursive CTE** — fetch all `(checkpoint_id, parent_checkpoint_id)` pairs for the thread, walk the ancestor chain in Python, then fetch writes with a plain `ANY()` filter.
- **Universal fallback on `BaseCheckpointSaver`** — the base class now provides `get_channel_writes` via `list()`, so any third-party saver works without modification.
- **`snapshot_every` removed** — deferred as a follow-up; the simpler design is easier to reason about and delivers larger storage savings.
+217
View File
@@ -0,0 +1,217 @@
"""Verify that SEND values land in checkpoint_blobs (TASKS channel),
not just in checkpoint_writes.
Strategy:
1. Build a tiny graph: node "1" emits a Send to node "2" via
conditional edges. interrupt_before=["2"] freezes the loop after
super-step 1 completes (cp_1 persisted) but before "2" runs.
2. Run graph.invoke against PostgresSaver.
3. Read checkpoint_blobs raw via psycopg.
4. Find the row with channel = '__pregel_tasks' and decode the blob
using the saver's serde. Assert it contains a non-empty list of
Send objects.
5. As a control, also dump checkpoint_writes to confirm the *same*
value is independently present in the writes table.
"""
from __future__ import annotations
import operator
from typing import Annotated
from typing_extensions import TypedDict
from uuid import uuid4
from psycopg import Connection
from psycopg.rows import dict_row
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.serde.types import TASKS
from langgraph.graph import START, StateGraph
from langgraph.types import Send
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/"
class State(TypedDict):
history: Annotated[list[str], operator.add]
def make_graph(checkpointer: PostgresSaver, *, with_interrupt: bool = True):
def node_one(state: State) -> State:
return {"history": ["1"]}
def node_two(state: dict) -> State:
return {"history": [f"2:{state['payload']}"]}
def fanout(state: State):
return [Send("two", {"payload": x}) for x in ("a", "b", "c")]
builder = StateGraph(State)
builder.add_node("one", node_one)
builder.add_node("two", node_two)
builder.add_edge(START, "one")
builder.add_conditional_edges("one", fanout, ["two"])
return builder.compile(
checkpointer=checkpointer,
interrupt_before=["two"] if with_interrupt else [],
)
def main() -> None:
database = f"verify_{uuid4().hex[:12]}"
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
uri = DEFAULT_POSTGRES_URI + database
with Connection.connect(
uri, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
saver = PostgresSaver(conn)
saver.setup()
graph = make_graph(saver)
thread_id = "t1"
cfg = {"configurable": {"thread_id": thread_id}}
result = graph.invoke({"history": []}, cfg)
print("=== invoke result (interrupted before 'two') ===")
print(result)
print()
# Inspect raw checkpoint_blobs.
rows = conn.execute(
"""
SELECT thread_id, checkpoint_ns, channel, version, type,
octet_length(blob) AS nbytes, blob
FROM checkpoint_blobs
WHERE thread_id = %s
ORDER BY channel, version
""",
(thread_id,),
).fetchall()
print(f"=== checkpoint_blobs ({len(rows)} rows) ===")
for r in rows:
preview = (
"(empty)"
if r["blob"] is None
else f"{r['blob'][:60]!r}{'...' if len(r['blob']) > 60 else ''}"
)
print(
f" channel={r['channel']!r:30s} v={r['version']!s:18s} "
f"type={r['type']:8s} nbytes={r['nbytes']!s:>5} blob={preview}"
)
print()
tasks_rows = [r for r in rows if r["channel"] == TASKS]
assert tasks_rows, (
f"expected at least one row in checkpoint_blobs for {TASKS!r}, "
f"got channels={[r['channel'] for r in rows]}"
)
# Decode the latest TASKS blob using the saver's serde.
tasks_row = tasks_rows[-1]
decoded = saver.serde.loads_typed((tasks_row["type"], tasks_row["blob"]))
print(f"=== decoded {TASKS!r} blob ({tasks_row['type']}) ===")
print(f" type(decoded) = {type(decoded).__name__}")
print(f" len(decoded) = {len(decoded)}")
for i, v in enumerate(decoded):
print(f" [{i}] type={type(v).__name__} value={v!r}")
print()
assert isinstance(decoded, list), (
f"expected list of Sends, got {type(decoded).__name__}"
)
assert len(decoded) == 3, f"expected 3 Sends (a, b, c), got {len(decoded)}"
for v in decoded:
assert isinstance(v, Send), (
f"expected Send instance in blob, got {type(v).__name__}"
)
assert v.node == "two"
print("✅ ASSERTION PASSED: checkpoint_blobs contains the Sends.")
print()
# Control: confirm the same Sends are also in checkpoint_writes.
wrows = conn.execute(
"""
SELECT task_id, idx, channel, type, octet_length(blob) AS nbytes, blob
FROM checkpoint_writes
WHERE thread_id = %s AND channel = %s
ORDER BY task_id, idx
""",
(thread_id, TASKS),
).fetchall()
print(
f"=== checkpoint_writes (channel={TASKS!r}, {len(wrows)} rows) ==="
)
for r in wrows:
v = saver.serde.loads_typed((r["type"], r["blob"]))
print(
f" task={r['task_id']} idx={r['idx']} type={r['type']} "
f"nbytes={r['nbytes']} -> {v!r}"
)
assert len(wrows) == 3, (
f"expected 3 writes-table rows for {TASKS!r}, got {len(wrows)}"
)
print(
"✅ ASSERTION PASSED: checkpoint_writes also contains the Sends."
)
print()
print(
"Conclusion (cp_1): SEND is double-stored. blobs row carries "
"the authoritative Send list; writes rows carry the same "
"payload as a redundant per-task record."
)
# ----------------------------------------------------------------
# Now resume: let 'two' actually run, so the Sends get consumed.
# We expect the next checkpoint's TASKS blob to be cleared,
# proving the accumulate=False lifecycle.
# ----------------------------------------------------------------
graph2 = make_graph(saver, with_interrupt=False)
result2 = graph2.invoke(None, cfg)
print()
print("=== invoke result (resumed; 'two' fanout completed) ===")
print(result2)
print()
rows2 = conn.execute(
"""
SELECT channel, version, type, octet_length(blob) AS nbytes, blob
FROM checkpoint_blobs
WHERE thread_id = %s AND channel = %s
ORDER BY version
""",
(thread_id, TASKS),
).fetchall()
print(f"=== checkpoint_blobs rows for {TASKS!r} across all cps ===")
for r in rows2:
decoded_v = saver.serde.loads_typed((r["type"], r["blob"]))
print(
f" v={r['version']} type={r['type']} nbytes={r['nbytes']} "
f"-> {decoded_v!r}"
)
assert len(rows2) >= 2, (
"expected at least 2 versions of TASKS blob across the run"
)
latest = saver.serde.loads_typed(
(rows2[-1]["type"], rows2[-1]["blob"])
)
assert latest == [], (
f"expected latest TASKS blob to be cleared (==[]) "
f"after Sends were consumed, got {latest!r}"
)
print()
print(
"✅ ASSERTION PASSED: after Sends were consumed, the latest "
"TASKS blob is [] (Topic with accumulate=False clears itself "
"via update(EMPTY_SEQ) in apply_writes line 304-311)."
)
finally:
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
if __name__ == "__main__":
main()