From fe2bc286fc5a2afe21f45e75ab47e90e15730c45 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Fri, 17 Apr 2026 14:35:46 -0400 Subject: [PATCH] feat(channels): implement DiffChannel for incremental checkpoint storage Adds DiffChannel, a new channel type that stores only per-step write deltas in checkpoints and reconstructs the full list by replaying the chain through the operator at load time. Co-Authored-By: Claude Sonnet 4.6 --- libs/langgraph/langgraph/channels/__init__.py | 2 + libs/langgraph/langgraph/channels/diff.py | 157 ++++++++++++++++++ libs/langgraph/tests/test_channels.py | 120 +++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 libs/langgraph/langgraph/channels/diff.py diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index a69c230b5..213fd6778 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,6 +1,7 @@ from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.diff import DiffChannel from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue, LastValueAfterFinish from langgraph.channels.named_barrier_value import ( @@ -20,6 +21,7 @@ __all__ = ( "UntrackedValue", "EphemeralValue", "BinaryOperatorAggregate", + "DiffChannel", "NamedBarrierValue", "NamedBarrierValueAfterFinish", # topics diff --git a/libs/langgraph/langgraph/channels/diff.py b/libs/langgraph/langgraph/channels/diff.py new file mode 100644 index 000000000..56c925c3d --- /dev/null +++ b/libs/langgraph/langgraph/channels/diff.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import collections.abc +from collections.abc import Callable, Sequence +from typing import Any, Generic + +from langgraph.checkpoint.base import DiffChainValue, DiffDelta +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, _strip_extras +from langgraph.errors import EmptyChannelError + +__all__ = ("DiffChannel",) + + +class DiffChannel(Generic[Value], BaseChannel[list[Value], Value, DiffDelta]): + """A channel that stores only per-step write deltas in checkpoints. + + Reconstructs the full accumulated list at load time by replaying the + chain of deltas through the operator. Use with append-style reducers + (e.g. `add_messages`) on long-running threads to reduce checkpoint + storage from O(N²) to O(N). + + Requires InMemorySaver or PostgresSaver; SqliteSaver is not supported. + + Usage:: + + class State(TypedDict): + messages: Annotated[list[AnyMessage], DiffChannel(add_messages)] + """ + + __slots__ = ("value", "operator", "_pending", "_base_version", "_overwritten") + + def __init__( + self, + operator: Callable[[list[Value], Any], list[Value]], + typ: type = list, + ) -> None: + typ = _strip_extras(typ) + if typ in ( + collections.abc.Sequence, + collections.abc.MutableSequence, + ): + typ = list + super().__init__(typ) + self.operator = operator + try: + self.value: list[Value] = typ() + except Exception: + self.value = [] + self._pending: list[Any] = [] + self._base_version: str | None = None + self._overwritten: bool = False + + def __eq__(self, other: object) -> bool: + if not isinstance(other, DiffChannel): + return False + if ( + self.operator.__name__ != "" + and other.operator.__name__ != "" + ): + 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 = DiffChannel(self.operator, self.typ) + new.key = self.key + new.value = self.value[:] + new._pending = self._pending[:] + new._base_version = self._base_version + new._overwritten = self._overwritten + return new + + def from_checkpoint(self, checkpoint: Any) -> Self: + new = DiffChannel(self.operator, self.typ) + new.key = self.key + if checkpoint is MISSING: + new.value = [] + elif isinstance(checkpoint, DiffChainValue): + accumulated: list[Value] = list(checkpoint.base) if checkpoint.base else [] + for step_writes in checkpoint.deltas: + for write in step_writes: + accumulated = new.operator(accumulated, write) + new.value = accumulated + elif isinstance(checkpoint, DiffDelta): + raise ValueError( + "DiffChannel received a raw DiffDelta from the checkpoint saver. " + "Your saver does not support incremental channel storage. " + "Use InMemorySaver or PostgresSaver." + ) + else: + # Backwards compat: plain list from old BinaryOperatorAggregate checkpoint. + new.value = list(checkpoint) + new._pending = [] + new._base_version = None # set by the subsequent after_checkpoint() call + new._overwritten = False + 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._pending = list(self.value) + self._overwritten = True + seen_overwrite = True + elif not seen_overwrite: + self.value = self.operator(self.value, value) + self._pending.append(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 and self.value is not None + + def checkpoint(self) -> DiffDelta: + return DiffDelta( + delta=self._pending[:], + prev_version=None if self._overwritten else self._base_version, + ) + + def after_checkpoint(self, version: Any) -> None: + if version != self._base_version: + self._base_version = version + self._pending = [] + self._overwritten = False diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 37a5eb432..aee6a1cea 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -117,3 +117,123 @@ def test_untracked_value() -> None: new_channel = UntrackedValue(dict).from_checkpoint(checkpoint) with pytest.raises(EmptyChannelError): new_channel.get() + + +def test_diff_channel_basic_two_steps() -> None: + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.base import DiffDelta + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + ch = DiffChannel(add_messages).from_checkpoint(MISSING) + ch.after_checkpoint(None) + + # Step 1: one message added + ch.update([HumanMessage(content="hi", id="h1")]) + d1 = ch.checkpoint() + assert isinstance(d1, DiffDelta) + assert len(d1.delta) == 1 + assert d1.prev_version is None # first ever step + ch.after_checkpoint("v1") + + # Step 2: another message + ch.update([AIMessage(content="hello", id="a1")]) + d2 = ch.checkpoint() + assert d2.prev_version == "v1" + assert len(d2.delta) == 1 + ch.after_checkpoint("v2") + + # Full accumulated value is preserved in memory + assert len(ch.get()) == 2 + assert ch.get()[0].content == "hi" + assert ch.get()[1].content == "hello" + + +def test_diff_channel_after_checkpoint_no_op_when_unchanged() -> None: + from langchain_core.messages import HumanMessage + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + ch = DiffChannel(add_messages).from_checkpoint(MISSING) + ch.after_checkpoint(None) + ch.update([HumanMessage(content="hi", id="h1")]) + ch.after_checkpoint("v1") + + # Same version: no-op + ch.after_checkpoint("v1") + assert ch._base_version == "v1" + assert ch._pending == [] + + +def test_diff_channel_from_checkpoint_chain() -> None: + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.base import DiffChainValue + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + spec = DiffChannel(add_messages) + chain = DiffChainValue( + base=None, + deltas=[ + [HumanMessage(content="hi", id="h1")], + [AIMessage(content="hello", id="a1")], + [HumanMessage(content="bye", id="h2")], + ], + ) + ch = spec.from_checkpoint(chain) + msgs = ch.get() + assert len(msgs) == 3 + assert msgs[0].content == "hi" + assert msgs[1].content == "hello" + assert msgs[2].content == "bye" + + +def test_diff_channel_from_checkpoint_backwards_compat() -> None: + from langchain_core.messages import HumanMessage + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + # Old BinaryOperatorAggregate checkpoint: plain list + spec = DiffChannel(add_messages) + old_value = [HumanMessage(content="old", id="h1")] + ch = spec.from_checkpoint(old_value) + assert ch.get() == old_value + + +def test_diff_channel_overwrite_resets_chain() -> None: + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.base import DiffDelta + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + from langgraph.types import Overwrite + + ch = DiffChannel(add_messages).from_checkpoint(MISSING) + ch.after_checkpoint(None) + ch.update([HumanMessage(content="old", id="h1")]) + ch.after_checkpoint("v1") + + # Overwrite should create a root blob (prev_version=None) + ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) + d = ch.checkpoint() + assert isinstance(d, DiffDelta) + assert d.prev_version is None # chain root + assert len(d.delta) == 1 + assert d.delta[0].content == "new" + + +def test_diff_channel_unsupported_saver_raises() -> None: + from langgraph.checkpoint.base import DiffDelta + + from langgraph.channels.diff import DiffChannel + from langgraph.graph.message import add_messages + + # If a saver returns a raw DiffDelta (unsupported), from_checkpoint raises + spec = DiffChannel(add_messages) + raw_delta = DiffDelta(delta=[], prev_version=None) + with pytest.raises(ValueError, match="DiffChannel received a raw DiffDelta"): + spec.from_checkpoint(raw_delta)