diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 3479c0166..d8548c2b5 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -25,14 +25,19 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): Works with all checkpointers. Savers with a dedicated blob store (InMemorySaver, PostgresSaver) use an O(1) fast-path per chain step; - all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal - bounded by `snapshot_every`. + all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal. + + Use `snapshot_every=N` to cap chain traversal depth at N steps. Every N + steps a full snapshot is written as the chain root; subsequent deltas + chain back to it, so `get_state` / reload never traverses more than N + checkpoints regardless of thread length. Recommended for savers without + a dedicated blob store. Usage:: class State(TypedDict): messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)] - # Cap reconstruction depth for non-Postgres savers: + # Cap reconstruction depth (recommended for SQLite / MongoDB savers): messages: Annotated[list[AnyMessage], DeltaChannel(add_messages, snapshot_every=50)] """ @@ -44,7 +49,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): "_base_version", "_last_checkpoint_id", "_overwritten", - "_steps_since_rehydrate", + "_steps_since_snapshot", ) def __init__( @@ -71,7 +76,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): self._base_version: str | None = None self._last_checkpoint_id: str | None = None self._overwritten: bool = False - self._steps_since_rehydrate: int = 0 + self._steps_since_snapshot: int = 0 def __eq__(self, other: object) -> bool: if not isinstance(other, DeltaChannel): @@ -101,7 +106,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): new._base_version = self._base_version new._last_checkpoint_id = self._last_checkpoint_id new._overwritten = self._overwritten - new._steps_since_rehydrate = self._steps_since_rehydrate + new._steps_since_snapshot = self._steps_since_snapshot return new def from_checkpoint(self, checkpoint: Any) -> Self: @@ -117,7 +122,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): new.value = accumulated # Seed the counter from actual chain depth so rehydration fires at # the right time regardless of how many prior invocations there were. - new._steps_since_rehydrate = len(checkpoint.deltas) + new._steps_since_snapshot = len(checkpoint.deltas) elif isinstance(checkpoint, DeltaValue): # Should never reach here — the pregel layer assembles DeltaValues # into DeltaChainValue before calling from_checkpoint. @@ -175,7 +180,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): def checkpoint(self) -> Any: if ( self.snapshot_every is not None - and self._steps_since_rehydrate >= self.snapshot_every + and self._steps_since_snapshot >= self.snapshot_every ): # Emit a full snapshot to cap chain depth at snapshot_every. # The saver stores this as a plain (non-diff) blob, so future @@ -191,10 +196,10 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): if self._base_version is None: pass # First call after from_checkpoint — anchor without counting a step. elif self.snapshot_every is not None: - if self._steps_since_rehydrate >= self.snapshot_every: - self._steps_since_rehydrate = 0 + if self._steps_since_snapshot >= self.snapshot_every: + self._steps_since_snapshot = 0 else: - self._steps_since_rehydrate += 1 + self._steps_since_snapshot += 1 self._base_version = version self._last_checkpoint_id = checkpoint_id self._pending = [] diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 3b28073c7..199b013a7 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -37,6 +37,7 @@ def _assemble_delta_channels( """ thread_id = str(config["configurable"]["thread_id"]) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + current_checkpoint_id = checkpoint.get("id") assembled: dict[str, Any] = {} for channel, value in checkpoint["channel_values"].items(): @@ -46,7 +47,8 @@ def _assemble_delta_channels( chain_deltas: list[list[Any]] = [] base: list[Any] | None = None cursor: DeltaValue = value - visited: set[str] = set() + # Pre-seed with current checkpoint ID to guard against self-referential chains. + visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set() while True: chain_deltas.append(cursor.delta) @@ -115,6 +117,7 @@ async def _aassemble_delta_channels( """Async version of _assemble_delta_channels.""" thread_id = str(config["configurable"]["thread_id"]) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + current_checkpoint_id = checkpoint.get("id") assembled: dict[str, Any] = {} for channel, value in checkpoint["channel_values"].items(): @@ -124,7 +127,7 @@ async def _aassemble_delta_channels( chain_deltas: list[list[Any]] = [] base: list[Any] | None = None cursor: DeltaValue = value - visited: set[str] = set() + visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set() while True: chain_deltas.append(cursor.delta) diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index a7a57ab29..231f491d8 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -358,3 +358,167 @@ def test_delta_channel_update_by_id_delta_and_replay() -> None: ch2 = spec.from_checkpoint(chain) assert len(ch2.get()) == 1 assert ch2.get()[0].content == "updated" + + +def test_delta_channel_snapshot_every_emits_plain_list() -> None: + """snapshot_every=N causes a plain-list snapshot after N steps; next deltas chain to it.""" + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.base import DeltaValue + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import add_messages + + SNAP = 3 + spec = DeltaChannel(add_messages, snapshot_every=SNAP) + ch = spec.from_checkpoint(MISSING) + # First after_checkpoint anchors _base_version without counting a step. + ch.after_checkpoint("v0", checkpoint_id="cid0") + + # Steps 1..SNAP: each should stay as DeltaValue; counter increments each step. + for i in range(1, SNAP + 1): + ch.update([HumanMessage(content=f"m{i}", id=f"h{i}")]) + ckpt = ch.checkpoint() + assert isinstance(ckpt, DeltaValue), f"expected DeltaValue at step {i}" + ch.after_checkpoint(f"v{i}", checkpoint_id=f"cid{i}") + + # Step SNAP+1: _steps_since_snapshot == SNAP → snapshot fires + ch.update([HumanMessage(content="snap", id="hsnap")]) + snap = ch.checkpoint() + assert isinstance(snap, list), "expected plain-list snapshot at snapshot_every step" + assert len(snap) == SNAP + 1 + + # After snapshot, counter resets — next step is DeltaValue again + ch.after_checkpoint("vsnap", checkpoint_id="cidsnap") + ch.update([HumanMessage(content="post", id="hpost")]) + post = ch.checkpoint() + assert isinstance(post, DeltaValue) + assert post.prev_checkpoint_id == "cidsnap" + + +def test_delta_channel_snapshot_every_end_to_end() -> None: + """Graph with snapshot_every: get_state returns correct accumulated value after snapshot.""" + from typing import Annotated + + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from typing_extensions import TypedDict + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph import START, StateGraph + from langgraph.graph.message import add_messages + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)] + + counter = {"n": 0} + + def respond(state: State) -> dict: + counter["n"] += 1 + return { + "messages": [ + AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}") + ] + } + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "snap-test"}} + + # Run 5 turns — snapshot fires after 2 steps, then again after 2 more + for i in range(5): + graph.invoke({"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + # 5 human + 5 AI = 10 total + assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}" + + +def test_delta_channel_assembly_fast_path_returns_delta_value() -> None: + """get_channel_blob returning a DeltaValue continues chain traversal (fast-path).""" + from unittest.mock import MagicMock + + from langgraph.checkpoint.base import ( + DeltaChainValue, + DeltaValue, + empty_checkpoint, + ) + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import add_messages + from langgraph.pregel._checkpoint import _assemble_delta_channels + + msg1 = {"type": "human", "content": "one"} + msg2 = {"type": "ai", "content": "two"} + msg3 = {"type": "human", "content": "three"} + + # cp3 → cp2 (DeltaValue) → cp1 (base list) + dv_cp2 = DeltaValue(delta=[msg2], prev_checkpoint_id="cp1") + cp3 = empty_checkpoint() + cp3["id"] = "cp3" + cp3["channel_values"]["messages"] = DeltaValue( + delta=[msg3], prev_checkpoint_id="cp2" + ) + + saver = MagicMock() + + def _get_blob(thread_id, ns, checkpoint_id, channel): + if checkpoint_id == "cp2": + return dv_cp2 # DeltaValue — chain continues + if checkpoint_id == "cp1": + return [msg1] # plain list — chain root + return NotImplemented + + saver.get_channel_blob.side_effect = _get_blob + + config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}} + assembled = _assemble_delta_channels(cp3, config, saver) + + chain = assembled["messages"] + assert isinstance(chain, DeltaChainValue) + assert chain.base == [msg1] + assert chain.deltas == [[msg2], [msg3]] + + spec = DeltaChannel(add_messages) + ch = spec.from_checkpoint(chain) + # add_messages converts dicts to message objects; check by type and content + + result = ch.get() + assert len(result) == 3 + assert result[0].content == "one" + assert result[1].content == "two" + assert result[2].content == "three" + + +def test_delta_channel_assembly_broken_chain_logs_warning() -> None: + """If a prev_checkpoint_id points to a missing checkpoint, log a warning and use partial chain.""" + from unittest.mock import MagicMock + + from langgraph.checkpoint.base import DeltaValue, empty_checkpoint + + from langgraph.pregel._checkpoint import _assemble_delta_channels + + cp = empty_checkpoint() + cp["id"] = "cp2" + cp["channel_values"]["messages"] = DeltaValue( + delta=["msg2"], prev_checkpoint_id="cp-missing" + ) + + saver = MagicMock() + saver.get_channel_blob.return_value = NotImplemented + saver.get_tuple.return_value = None # checkpoint not found + + config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}} + + assembled = _assemble_delta_channels(cp, config, saver) + + # Should still assemble — with partial chain (just the current delta, base=None) + assert "messages" in assembled + from langgraph.checkpoint.base import DeltaChainValue + + chain = assembled["messages"] + assert isinstance(chain, DeltaChainValue) + assert chain.base is None + assert chain.deltas == [["msg2"]]