feat(channels): step-based eager snapshot_frequency via _DeltaSnapshot ext type

Replaces the write-count (_write_count) snapshot mechanism with a clean
step-based approach: pregel's create_checkpoint fires snapshots every N
pregel steps regardless of whether the channel was written (eager).

Key changes:
- snapshot_frequency=None (default) for pure delta; int N for snapshot every N steps
- DeltaChannel.checkpoint() always returns DELTA_SENTINEL; snapshot logic
  lives in create_checkpoint which has the step number
- create_checkpoint bumps channel version via get_next_version when the
  channel wasn't written at a snapshot step (eager: always stores the blob)
- _DeltaSnapshot NamedTuple registered as EXT_DELTA_SNAPSHOT (code 7) in
  the msgpack serde — no dict key collision, type tag does the dispatch
- InMemorySaver and PostgresSaver _get_channel_writes_history updated:
  _DeltaSnapshot blobs collect pending_writes before terminating (they
  encode the NEXT step's transition, not subsumed by the snapshot unlike
  pre-delta migration blobs)

Tests confirm:
- Snapshots fire at every N steps even when channel has no write that step
- Correct accumulated state after reconstruction from snapshot + replay

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-27 20:43:59 -04:00
co-authored by Claude Sonnet 4.6
parent af9ef2a1f9
commit bf30da1c8f
8 changed files with 257 additions and 85 deletions
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
_ChannelWritesHistory,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.types import TASKS
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
from psycopg.types.json import Jsonb
MetadataInput = dict[str, Any] | None
@@ -281,13 +281,24 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
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:
if isinstance(blob_value, _DeltaSnapshot):
# Step-based snapshot: collect this ancestor's
# pending_writes first (they encode the NEXT step's
# transition, not subsumed by the snapshot blob).
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()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
# Pre-delta blob: subsumes this ancestor's writes.
collected.reverse()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
@@ -27,6 +27,7 @@ from langgraph.checkpoint.base import (
get_checkpoint_id,
get_checkpoint_metadata,
)
from langgraph.checkpoint.serde.types import _DeltaSnapshot
logger = logging.getLogger(__name__)
@@ -185,8 +186,31 @@ class InMemorySaver(
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.
if isinstance(blob_value, _DeltaSnapshot):
# Step-based snapshot: the blob is state AT this
# ancestor, but the ancestor's pending_writes
# encode the NEXT step's transition and are NOT
# subsumed by the snapshot — collect them first.
step_writes = self.writes.get(
(thread_id, checkpoint_ns, cp_id), {}
)
for (_task_id, _idx), (
tid,
ch,
serialized,
_,
) in sorted(step_writes.items(), reverse=True):
if ch != channel:
continue
collected.append(
(tid, ch, self.serde.loads_typed(serialized))
)
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
# Pre-delta blob: state AT this ancestor already
# subsumes its pending_writes — skip them.
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
@@ -33,7 +33,7 @@ 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 DELTA_SENTINEL, SendProtocol
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol, _DeltaSnapshot
from langgraph.store.base import Item
if TYPE_CHECKING:
@@ -296,10 +296,13 @@ EXT_METHOD_SINGLE_ARG = 3
EXT_PYDANTIC_V1 = 4
EXT_PYDANTIC_V2 = 5
EXT_NUMPY_ARRAY = 6
EXT_DELTA_SNAPSHOT = 7
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
if isinstance(obj, _DeltaSnapshot):
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
return ormsgpack.Ext(
EXT_PYDANTIC_V2,
_msgpack_enc(
@@ -613,7 +616,13 @@ def _create_msgpack_ext_hook(
return False
def ext_hook(code: int, data: bytes) -> Any:
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
if code == EXT_DELTA_SNAPSHOT:
return _DeltaSnapshot(
ormsgpack.unpackb(
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
)
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
try:
tup = ormsgpack.unpackb(
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
@@ -1,6 +1,7 @@
from collections.abc import Sequence
from typing import (
Any,
NamedTuple,
Protocol,
TypeVar,
runtime_checkable,
@@ -33,6 +34,20 @@ class _DeltaSentinel:
DELTA_SENTINEL = _DeltaSentinel()
class _DeltaSnapshot(NamedTuple):
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency.
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
The ancestor walk in `_get_channel_writes_history` terminates when it
encounters this type (any non-sentinel blob stops the walk).
`from_checkpoint` reconstructs the channel value directly from `.value`
without replaying writes — the snapshot IS the accumulated state.
"""
value: Any
Value = TypeVar("Value", covariant=True)
Update = TypeVar("Update", contravariant=True)
C = TypeVar("C")
+31 -51
View File
@@ -1,11 +1,11 @@
from __future__ import annotations
import copy as _copy
import math
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import Self
from langgraph._internal._constants import OVERWRITE
@@ -40,38 +40,34 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]:
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
"""Fold-reducer channel with configurable snapshot cadence.
`snapshot_frequency=math.inf` (default): pure delta mode — stores only
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes
through `operator`. O(N) storage, O(N) read replay depth.
`snapshot_frequency=None` (default): pure delta — stores only
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
`snapshot_frequency=N` (integer > 1): writes a full snapshot blob every
N steps; on non-snapshot steps stores `DELTA_SENTINEL`. Reads walk at
most N ancestors before finding the snapshot, bounding replay depth to N.
Storage is O(N²/N) = O(N·avg_msg_size) — still linear per step,
but snapshots grow with accumulated messages so total is O(N²/N).
`snapshot_frequency=N`: pregel's `create_checkpoint` writes a full
`_DeltaSnapshot` blob every N steps (eagerly, even if the channel had
no write that step). Reads walk at most N ancestor checkpoints before
hitting the snapshot, bounding replay depth to N regardless of thread
length.
Parameters:
operator: Binary reducer `(Value, Value) -> Value` applied pairwise.
Must be associative when `snapshot_frequency > 1` since replayed
writes fold through the same operator as live writes.
snapshot_frequency: Every Nth step writes a full snapshot blob.
Default `math.inf` (pure delta, never snapshot).
operator: Binary reducer `(Value, Value) -> Value`.
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
`None` (default) = pure delta, never snapshot.
"""
__slots__ = ("value", "operator", "snapshot_frequency", "_write_count")
__slots__ = ("value", "operator", "snapshot_frequency")
value: Value | Any
def __init__(
self,
operator: Callable[[Any, Any], Any],
*,
snapshot_frequency: int | float = math.inf,
snapshot_frequency: int | None = None,
) -> None:
super().__init__(list)
self.operator = operator
self.snapshot_frequency = snapshot_frequency
self.value: Any = []
self._write_count: int = 0
def __eq__(self, other: object) -> bool:
if not isinstance(other, DeltaChannel):
@@ -93,19 +89,23 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
def UpdateType(self) -> Any:
return self.typ
def is_snapshot_step(self, step: int) -> bool:
"""True if pregel should write a snapshot blob at this step."""
return (
self.snapshot_frequency is not None and step % self.snapshot_frequency == 0
)
def _clone_empty(self) -> Self:
new = self.__class__.__new__(self.__class__)
new.typ = self.typ
new.key = self.key
new.operator = self.operator
new.snapshot_frequency = self.snapshot_frequency
new._write_count = 0
new.value = MISSING
return new
def copy(self) -> Self:
new = self._clone_empty()
new._write_count = self._write_count
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
return new
@@ -123,36 +123,24 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
def from_checkpoint(self, checkpoint: Any) -> Self:
"""Initialize from a stored blob or sentinel.
Blob formats:
* `DELTA_SENTINEL` / `MISSING`: start empty; caller will replay writes.
* `{"__delta_v__": value, "__delta_wc__": n}`: snapshot blob — restore
value and write count so future writes trigger the next snapshot at
the right cadence.
* plain value (migration from old DeltaChannel blobs): restore value,
reset write count to 0.
Blob types (dispatched via serde ext code, not dict key inspection):
* `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes.
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
* plain value (migration from old BinOp blobs): use directly.
"""
new = self._clone_empty()
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
new.value = _empty(self.typ)
new._write_count = 0
elif isinstance(checkpoint, dict) and "__delta_v__" in checkpoint:
new.value = checkpoint["__delta_v__"]
new._write_count = checkpoint["__delta_wc__"]
new.value = _empty(new.typ)
elif isinstance(checkpoint, _DeltaSnapshot):
new.value = checkpoint.value
else:
new.value = checkpoint
new._write_count = 0
return new
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
"""Fold ancestor writes oldest→newest into current value.
Also increments `_write_count` per replayed write so the snapshot
cadence stays correct across invocations. The count resumes from
wherever the seed's `__delta_wc__` left off (set by `from_checkpoint`).
"""
"""Fold ancestor writes oldest→newest into current value."""
for _, _, value in writes:
self.value = self._apply_write(self.value, value)
self._write_count += 1
def update(self, values: Sequence[Any]) -> bool:
if not values:
@@ -171,7 +159,6 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
elif seen_overwrite:
continue
self.value = self._apply_write(self.value, value)
self._write_count += 1
return True
def get(self) -> Any:
@@ -183,19 +170,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
return self.value is not MISSING
def checkpoint(self) -> Any:
"""Return stored representation.
"""Return stored representation: always `DELTA_SENTINEL`.
Pure delta mode (`snapshot_frequency=math.inf`) always returns
`DELTA_SENTINEL`. For finite `snapshot_frequency`, returns a snapshot
blob `{"__delta_v__": value, "__delta_wc__": n}` every
`snapshot_frequency` writes, with the current write count embedded so
`from_checkpoint` can restore the counter and maintain correct cadence
across invocations. All other writes return `DELTA_SENTINEL`.
Snapshot decisions are made by `create_checkpoint` in pregel (which
has the step number) via `is_snapshot_step`. `checkpoint()` is only
called for non-snapshot steps or when no checkpointer is available.
"""
if self.value is MISSING:
return MISSING
if self.snapshot_frequency == math.inf:
return DELTA_SENTINEL
if self._write_count > 0 and self._write_count % self.snapshot_frequency == 0:
return {"__delta_v__": self.value, "__delta_wc__": self._write_count}
return DELTA_SENTINEL
+39 -27
View File
@@ -1,19 +1,23 @@
from __future__ import annotations
from collections.abc import Mapping
from collections.abc import Callable, Mapping
from datetime import datetime, timezone
from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel
from langgraph.channels.delta import DeltaChannel # used in _needs_replay
from langgraph.channels.delta import DeltaChannel
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 4
GetNextVersion = Callable[[Any, None], Any]
def empty_checkpoint() -> Checkpoint:
return Checkpoint(
@@ -33,31 +37,50 @@ def create_checkpoint(
*,
id: str | None = None,
updated_channels: set[str] | None = None,
get_next_version: GetNextVersion | None = None,
) -> Checkpoint:
"""Create a checkpoint for the given channels.
`DeltaChannel.checkpoint()` returns `DELTA_SENTINEL` or the full value
depending on its internal write count, so no special handling is needed
here. The snapshot is stored at the exact step of the triggering write.
For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a
`_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor
walk to at most N steps. Snapshots are eager: even if the channel had no
write this step, a version bump is forced (via `get_next_version`) so the
blob is stored by `put()`. Without `get_next_version` (e.g. static
contexts), snapshot steps gracefully fall back to sentinel.
"""
ts = datetime.now(timezone.utc).isoformat()
if channels is None:
values = checkpoint["channel_values"]
channel_versions = checkpoint["channel_versions"]
else:
values = {}
channel_versions = dict(checkpoint["channel_versions"])
for k in channels:
if k not in checkpoint["channel_versions"]:
if k not in channel_versions:
continue
ch = channels[k]
v = ch.checkpoint()
if v is not MISSING:
values[k] = v
if (
isinstance(ch, DeltaChannel)
and ch.is_snapshot_step(step)
and ch.is_available()
):
# Eager snapshot: bump version if not already written this step
# so put() includes this channel in new_versions and stores blob.
if get_next_version is not None and (
updated_channels is None or k not in updated_channels
):
channel_versions[k] = get_next_version(channel_versions[k], None)
values[k] = _DeltaSnapshot(ch.get())
else:
v = ch.checkpoint()
if v is not MISSING:
values[k] = v
return Checkpoint(
v=LATEST_VERSION,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
channel_versions=channel_versions,
versions_seen=checkpoint["versions_seen"],
updated_channels=None if updated_channels is None else sorted(updated_channels),
)
@@ -67,9 +90,8 @@ def _needs_replay(spec: BaseChannel, stored: object) -> bool:
"""True if `spec` is a `DeltaChannel` and the stored blob is a sentinel,
requiring an ancestor walk to reconstruct.
Snapshot blobs are dicts with `__delta_v__`; sentinels are `DELTA_SENTINEL`.
Plain non-sentinel values are migration blobs (old DeltaChannel format).
All non-sentinel values resolve directly via `from_checkpoint`.
`_DeltaSnapshot` blobs and plain values (migration) resolve directly via
`from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay.
"""
if not isinstance(spec, DeltaChannel):
return False
@@ -86,16 +108,10 @@ def channels_from_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 on non-snapshot steps is
`DELTA_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 operator. The walk stops at the nearest full-snapshot
blob, so read depth is bounded by `snapshot_frequency`. Without saver/config
(static contexts — graph drawing, unit tests), delta-mode channels fall back
to empty.
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
ancestor walk via `saver._get_channel_writes_history`. The walk terminates
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain
value, so read depth is bounded by `snapshot_frequency`.
"""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
@@ -110,10 +126,6 @@ def channels_from_checkpoint(
ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING)
if _needs_replay(spec, stored) and saver is not None and config is not None:
# Walk ancestors for seed + writes. The saver's walk stops at
# the nearest non-sentinel blob (natural terminator under
# snapshot_frequency > 1; pre-migration blobs also act as
# terminators if the spec was changed mid-thread).
assert isinstance(spec, DeltaChannel)
history = saver._get_channel_writes_history(config, k)
replay_ch = spec.from_checkpoint(history.seed)
+3
View File
@@ -891,6 +891,9 @@ class PregelLoop:
self.step,
id=self.checkpoint["id"] if exiting else None,
updated_channels=self.updated_channels,
get_next_version=self.checkpointer_get_next_version
if do_checkpoint
else None,
)
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
if TASKS in self.checkpoint["channel_values"] and any(
+118
View File
@@ -275,6 +275,124 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None:
assert ch.checkpoint() is DELTA_SENTINEL
def test_delta_channel_snapshot_step_based() -> None:
"""Snapshots fire on every Nth step regardless of whether the channel was written.
With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot
blob — even if the channel had no write that step (eager snapshot). This
bounds the ancestor walk to at most N steps on any read.
"""
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
# snapshot_frequency=5: snapshot every 5 pregel steps
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=5)]
other: str
def node_a(state: State) -> dict:
# writes to messages
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
def node_b(state: State) -> dict:
# writes ONLY to other, not messages — snapshot must still fire at step N
return {"other": "y"}
g = StateGraph(State)
g.add_node("a", node_a)
g.add_node("b", node_b)
g.add_edge(START, "a")
g.add_edge("a", "b")
saver = InMemorySaver()
graph = g.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "t1"}}
for i in range(6):
graph.invoke(
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
config,
)
# Confirm at least one snapshot blob exists for messages
msg_blob_values = [
saver.serde.loads_typed((type_tag, blob))
for k, (type_tag, blob) in saver.blobs.items()
if k[2] == "messages" and type_tag == "msgpack" and blob
]
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
# Final state must be correct regardless of snapshot cadence
state = graph.get_state(config)
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
def test_delta_channel_snapshot_fires_even_when_not_written() -> None:
"""Eager snapshot: _DeltaSnapshot stored at snapshot step even when the
channel had no write that step (node_b doesn't touch messages).
"""
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages, snapshot_frequency=3)]
tick: int
def writer(state: State) -> dict:
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
def ticker(state: State) -> dict:
# never writes messages
return {"tick": state["tick"] + 1}
g = StateGraph(State)
g.add_node("writer", writer)
g.add_node("ticker", ticker)
g.add_edge(START, "writer")
g.add_edge("writer", "ticker")
saver = InMemorySaver()
graph = g.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "t1"}}
for i in range(5):
graph.invoke(
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0},
config,
)
# Count distinct message channel blob versions
msg_blobs = {
k: saver.serde.loads_typed((t, b))
for k, (t, b) in saver.blobs.items()
if k[2] == "messages" and t == "msgpack" and b
}
snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)}
# There must be snapshots (ticker steps are snapshot steps too)
assert snapshots, (
"eager snapshots must fire even on steps where messages wasn't written"
)
# All get_state calls must return the correct accumulated value
state = graph.get_state(config)
assert len(state.values["messages"]) == 10 # 5 human + 5 AI
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
from typing import Annotated