From bf30da1c8f7c1731e743205ea13a65fda2b03b7c Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Mon, 27 Apr 2026 20:43:59 -0400 Subject: [PATCH] feat(channels): step-based eager snapshot_frequency via _DeltaSnapshot ext type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../langgraph/checkpoint/postgres/base.py | 15 ++- .../langgraph/checkpoint/memory/__init__.py | 28 ++++- .../langgraph/checkpoint/serde/jsonplus.py | 15 ++- .../langgraph/checkpoint/serde/types.py | 15 +++ libs/langgraph/langgraph/channels/delta.py | 82 +++++------- .../langgraph/langgraph/pregel/_checkpoint.py | 66 ++++++---- libs/langgraph/langgraph/pregel/_loop.py | 3 + libs/langgraph/tests/test_channels.py | 118 ++++++++++++++++++ 8 files changed, 257 insertions(+), 85 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index b56c503fd..5f1c010ef 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -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, []): diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 39bb133dc..d9d704336 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -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 diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 48ef785ca..d23633526 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -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 diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index ea4cd8180..d68ec413a 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -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") diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 992f49a92..bb7945bb6 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -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 diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 3f56a09d8..aea3ed583 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -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) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index b2ce1d5f9..0ebb14042 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -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( diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 0b2d4ce47..ab86a4174 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -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