diff --git a/docs/superpowers/plans/2026-04-17-diff-channel.md b/docs/superpowers/plans/2026-04-17-diff-channel.md deleted file mode 100644 index 371a98599..000000000 --- a/docs/superpowers/plans/2026-04-17-diff-channel.md +++ /dev/null @@ -1,1350 +0,0 @@ -# DiffChannel: Incremental Checkpoint Storage Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Implement `DiffChannel`, a new opt-in channel type that stores only per-step write deltas instead of the full accumulated list, reducing checkpoint storage from O(N²) to O(N) for append-style reducers like `add_messages`. - -**Architecture:** A `DiffChannel` wraps a binary operator (e.g. `add_messages`) and returns a `DiffDelta` from `checkpoint()` instead of the full accumulated list. Savers detect the `"diff"` type tag, follow the `prev_version` chain via dict lookups (InMemorySaver) or a SQL range query (PostgresSaver), assemble a `DiffChainValue`, and pass it to `from_checkpoint` which replays deltas through the operator to reconstruct the full list. A new `after_checkpoint(version)` hook on `BaseChannel` allows `DiffChannel` to track its chain pointer without changing the saver public interface. - -**Tech Stack:** Python 3.11+, ormsgpack, langchain-core messages, pytest/anyio, psycopg (Postgres tests) - -**Branch:** `diff-channel-incremental-checkpointing` - ---- - -## File Map - -| File | Action | Responsibility | -|---|---|---| -| `libs/checkpoint/langgraph/checkpoint/base/__init__.py` | Modify | Add `DiffDelta`, `DiffChainValue` dataclasses | -| `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` | Modify | Add `"diff"` in `dumps_typed` + `loads_typed` | -| `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` | Modify | Chain traversal in `_load_blobs` | -| `libs/checkpoint/tests/test_jsonplus.py` | Modify | Serde round-trip tests | -| `libs/checkpoint/tests/test_memory.py` | Modify | InMemorySaver diff-chain tests | -| `libs/langgraph/langgraph/channels/base.py` | Modify | Add no-op `after_checkpoint` | -| `libs/langgraph/langgraph/channels/diff.py` | Create | `DiffChannel` implementation | -| `libs/langgraph/langgraph/channels/__init__.py` | Modify | Export `DiffChannel` | -| `libs/langgraph/tests/test_channels.py` | Modify | `DiffChannel` unit tests | -| `libs/langgraph/tests/test_pregel.py` | Modify | End-to-end graph integration tests | -| `libs/langgraph/langgraph/pregel/_checkpoint.py` | Modify | Call `after_checkpoint` in `channels_from_checkpoint` | -| `libs/langgraph/langgraph/pregel/_loop.py` | Modify | Call `after_checkpoint` after `create_checkpoint` | -| `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` | Modify | Range-query chain reconstruction in `_load_blobs` | -| `libs/checkpoint-postgres/tests/test_postgres.py` | Modify | Postgres diff-chain integration tests | - ---- - -## Task 1: Add Protocol Types `DiffDelta` and `DiffChainValue` - -**Files:** -- Modify: `libs/checkpoint/langgraph/checkpoint/base/__init__.py` - -These are the shared contract types between `DiffChannel` and savers. `DiffDelta` is what `channel.checkpoint()` returns; `DiffChainValue` is what savers pass to `channel.from_checkpoint()`. - -- [ ] **Step 1: Open `libs/checkpoint/langgraph/checkpoint/base/__init__.py` and locate the `PendingWrite` line (currently around line 30)** - -```python -PendingWrite = tuple[str, str, Any] -``` - -- [ ] **Step 2: Add `dataclasses` import and the two new types immediately after `PendingWrite`** - -Add to the imports at the top of the file (after the existing imports block): -```python -import dataclasses -``` - -Then directly after `PendingWrite = tuple[str, str, Any]`: -```python -@dataclasses.dataclass -class DiffDelta: - """Returned by DiffChannel.checkpoint(). Represents one step's writes.""" - - delta: list[Any] - prev_version: str | None # version of previous diff blob; None = chain root - - -@dataclasses.dataclass -class DiffChainValue: - """Passed to DiffChannel.from_checkpoint(). Assembled by saver _load_blobs().""" - - base: list[Any] | None # starting accumulated value; None = start from empty - deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest -``` - -- [ ] **Step 3: Add `DiffDelta` and `DiffChainValue` to the module's `__all__` (if one exists) or confirm they are importable** - -Run: -```bash -cd /Users/sydney_runkle/oss/langgraph && python -c "from langgraph.checkpoint.base import DiffDelta, DiffChainValue; print('ok')" -``` -Expected: `ok` - -- [ ] **Step 4: Commit** - -```bash -git add libs/checkpoint/langgraph/checkpoint/base/__init__.py -git commit -m "feat(checkpoint): add DiffDelta and DiffChainValue protocol types" -``` - ---- - -## Task 2: Extend Serde for `"diff"` Type Tag - -**Files:** -- Modify: `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` -- Modify: `libs/checkpoint/tests/test_jsonplus.py` - -The serde must serialize `DiffDelta` as `("diff", bytes)` and deserialize `("diff", bytes)` back to `{"d": [...], "p": version_or_none}`. Savers check the `"diff"` type tag and call `serde.loads_typed` to decode — no direct `ormsgpack` import needed in savers. - -- [ ] **Step 1: Write the failing serde tests in `libs/checkpoint/tests/test_jsonplus.py`** - -Add at the end of the file: -```python -def test_diff_delta_serde_round_trip() -> None: - from langchain_core.messages import HumanMessage - from langgraph.checkpoint.base import DiffDelta - - serde = JsonPlusSerializer() - prev = "00000000000000000000000000000001.1234567890123456" - delta = DiffDelta( - delta=[HumanMessage(content="hello", id="msg-1")], - prev_version=prev, - ) - type_tag, blob = serde.dumps_typed(delta) - assert type_tag == "diff" - - result = serde.loads_typed(("diff", blob)) - assert isinstance(result, dict) - assert result["p"] == prev - assert len(result["d"]) == 1 - assert result["d"][0].content == "hello" - - -def test_diff_delta_serde_root_blob() -> None: - from langgraph.checkpoint.base import DiffDelta - - serde = JsonPlusSerializer() - delta = DiffDelta(delta=[], prev_version=None) - type_tag, blob = serde.dumps_typed(delta) - assert type_tag == "diff" - - result = serde.loads_typed(("diff", blob)) - assert result["p"] is None - assert result["d"] == [] -``` - -- [ ] **Step 2: Run tests to confirm they fail** - -```bash -cd libs/checkpoint && TEST=tests/test_jsonplus.py::test_diff_delta_serde_round_trip make test -``` -Expected: `FAILED` — `NotImplementedError: Unknown serialization type: diff` (or similar) - -- [ ] **Step 3: Add the `"diff"` branch to `dumps_typed` in `jsonplus.py`** - -In `JsonPlusSerializer.dumps_typed`, locate this block (around line 235): -```python -def dumps_typed(self, obj: Any) -> tuple[str, bytes]: - if obj is None: - return "null", EMPTY_BYTES - elif isinstance(obj, bytes): - return "bytes", obj - elif isinstance(obj, bytearray): - return "bytearray", obj - else: - try: - return "msgpack", _msgpack_enc(obj) -``` - -Add the `DiffDelta` branch **before** the `else` block: -```python -def dumps_typed(self, obj: Any) -> tuple[str, bytes]: - if obj is None: - return "null", EMPTY_BYTES - elif isinstance(obj, bytes): - return "bytes", obj - elif isinstance(obj, bytearray): - return "bytearray", obj - elif isinstance(obj, DiffDelta): - return "diff", _msgpack_enc({"d": obj.delta, "p": obj.prev_version}) - else: - try: - return "msgpack", _msgpack_enc(obj) -``` - -Add the import of `DiffDelta` at the top of `jsonplus.py` (alongside existing imports): -```python -from langgraph.checkpoint.base import DiffDelta -``` - -- [ ] **Step 4: Add the `"diff"` branch to `loads_typed` in `jsonplus.py`** - -In `JsonPlusSerializer.loads_typed`, locate the dispatch (around line 250): -```python -def loads_typed(self, data: tuple[str, bytes]) -> Any: - type_, data_ = data - if type_ == "null": - return None - elif type_ == "bytes": - return data_ - elif type_ == "bytearray": - return bytearray(data_) - elif type_ == "json": - return json.loads(data_, object_hook=self._reviver) - elif type_ == "msgpack": - return ormsgpack.unpackb( - data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS - ) - elif self.pickle_fallback and type_ == "pickle": - return pickle.loads(data_) - else: - raise NotImplementedError(f"Unknown serialization type: {type_}") -``` - -Add the `"diff"` branch **before** the `else` raise: -```python - elif type_ == "diff": - return ormsgpack.unpackb( - data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS - ) - elif self.pickle_fallback and type_ == "pickle": - return pickle.loads(data_) - else: - raise NotImplementedError(f"Unknown serialization type: {type_}") -``` - -- [ ] **Step 5: Run the serde tests to confirm they pass** - -```bash -cd libs/checkpoint && TEST=tests/test_jsonplus.py::test_diff_delta_serde_round_trip\ tests/test_jsonplus.py::test_diff_delta_serde_root_blob make test -``` -Expected: both `PASSED` - -- [ ] **Step 6: Run the full serde test suite to check for regressions** - -```bash -cd libs/checkpoint && TEST=tests/test_jsonplus.py make test -``` -Expected: all existing tests still pass - -- [ ] **Step 7: Commit** - -```bash -git add libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py libs/checkpoint/tests/test_jsonplus.py -git commit -m "feat(checkpoint/serde): serialize DiffDelta as 'diff' type tag" -``` - ---- - -## Task 3: Add `after_checkpoint` Hook to `BaseChannel` - -**Files:** -- Modify: `libs/langgraph/langgraph/channels/base.py` - -This is a no-op default method. All existing channels inherit it silently. `DiffChannel` will override it to advance `_base_version` and clear `_pending`. - -- [ ] **Step 1: Open `libs/langgraph/langgraph/channels/base.py` and locate the `finish` method (currently the last method, around line 112)** - -```python -def finish(self) -> bool: - ... - return False -``` - -- [ ] **Step 2: Add `after_checkpoint` after `finish`** - -```python - def after_checkpoint(self, version: Any) -> 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. DiffChannel). - """ - pass -``` - -- [ ] **Step 3: Verify the method is accessible on existing channel types** - -```bash -cd libs/langgraph && python -c " -from langgraph.channels.last_value import LastValue -from langgraph.channels.binop import BinaryOperatorAggregate -import operator -ch = LastValue(int).from_checkpoint(3) -ch.after_checkpoint('v1') # must not raise -ch2 = BinaryOperatorAggregate(int, operator.add).from_checkpoint(0) -ch2.after_checkpoint('v2') # must not raise -print('ok') -" -``` -Expected: `ok` - -- [ ] **Step 4: Commit** - -```bash -git add libs/langgraph/langgraph/channels/base.py -git commit -m "feat(channels): add no-op after_checkpoint hook to BaseChannel" -``` - ---- - -## Task 4: Implement `DiffChannel` - -**Files:** -- Create: `libs/langgraph/langgraph/channels/diff.py` -- Modify: `libs/langgraph/langgraph/channels/__init__.py` -- Modify: `libs/langgraph/tests/test_channels.py` - -`DiffChannel` wraps a binary operator, accumulates incoming writes in `_pending`, and returns a `DiffDelta` from `checkpoint()`. `from_checkpoint` accepts a `DiffChainValue` and replays write-sets through the operator. - -- [ ] **Step 1: Write failing unit tests in `libs/langgraph/tests/test_channels.py`** - -Add at the end of the file: -```python -def test_diff_channel_basic_two_steps() -> None: - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.channels.diff import DiffChannel - from langgraph.checkpoint.base import DiffDelta - 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.channels.diff import DiffChannel - from langgraph.checkpoint.base import DiffChainValue - 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.channels.diff import DiffChannel - from langgraph.checkpoint.base import DiffDelta - 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.channels.diff import DiffChannel - from langgraph.checkpoint.base import DiffDelta - 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) -``` - -- [ ] **Step 2: Run to confirm tests fail** - -```bash -cd libs/langgraph && TEST=tests/test_channels.py::test_diff_channel_basic_two_steps make test -``` -Expected: `FAILED` — `ImportError: cannot import name 'DiffChannel'` - -- [ ] **Step 3: Create `libs/langgraph/langgraph/channels/diff.py`** - -```python -from __future__ import annotations - -import collections.abc -from collections.abc import Callable, Sequence -from typing import Any, Generic - -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.checkpoint.base import DiffChainValue, DiffDelta -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 -``` - -- [ ] **Step 4: Export `DiffChannel` from `libs/langgraph/langgraph/channels/__init__.py`** - -Open `libs/langgraph/langgraph/channels/__init__.py` and add `DiffChannel` to the imports and `__all__`. The file currently exports `BinaryOperatorAggregate`, `EphemeralValue`, `LastValue`, `Topic`, `UntrackedValue`. Add: - -```python -from langgraph.channels.diff import DiffChannel -``` - -And add `"DiffChannel"` to `__all__` if present. - -- [ ] **Step 5: Run all DiffChannel unit tests** - -```bash -cd libs/langgraph && TEST="tests/test_channels.py::test_diff_channel_basic_two_steps tests/test_channels.py::test_diff_channel_after_checkpoint_no_op_when_unchanged tests/test_channels.py::test_diff_channel_from_checkpoint_chain tests/test_channels.py::test_diff_channel_from_checkpoint_backwards_compat tests/test_channels.py::test_diff_channel_overwrite_resets_chain tests/test_channels.py::test_diff_channel_unsupported_saver_raises" make test -``` -Expected: all 6 `PASSED` - -- [ ] **Step 6: Run full channel test suite for regressions** - -```bash -cd libs/langgraph && TEST=tests/test_channels.py make test -``` -Expected: all tests pass - -- [ ] **Step 7: Commit** - -```bash -git add libs/langgraph/langgraph/channels/diff.py libs/langgraph/langgraph/channels/__init__.py libs/langgraph/tests/test_channels.py -git commit -m "feat(channels): implement DiffChannel for incremental checkpoint storage" -``` - ---- - -## Task 5: Extend `InMemorySaver._load_blobs` for Chain Traversal - -**Files:** -- Modify: `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` -- Modify: `libs/checkpoint/tests/test_memory.py` - -When `_load_blobs` encounters a blob with type `"diff"`, it follows the `prev_version` chain backwards through `self.blobs`, collects all deltas, and returns a `DiffChainValue` instead of a plain deserialized value. - -- [ ] **Step 1: Write failing integration tests in `libs/checkpoint/tests/test_memory.py`** - -Add at the end of the file (after the `TestMemorySaver` class): - -```python -class TestInMemorySaverDiffChannel: - def test_diff_channel_chain_reconstruction(self) -> None: - """_load_blobs follows the diff chain and returns DiffChainValue.""" - from langgraph.checkpoint.base import DiffChainValue, DiffDelta - - saver = InMemorySaver() - serde = JsonPlusSerializer() - - thread_id = "t1" - ns = "" - - # Simulate two steps: v1 (root) and v2 (chained to v1) - v1 = "00000000000000000000000000000001.1234567890000000" - v2 = "00000000000000000000000000000002.1234567890000000" - - delta1 = DiffDelta(delta=["msg1"], prev_version=None) - delta2 = DiffDelta(delta=["msg2"], prev_version=v1) - - saver.blobs[(thread_id, ns, "messages", v1)] = serde.dumps_typed(delta1) - saver.blobs[(thread_id, ns, "messages", v2)] = serde.dumps_typed(delta2) - - channel_values = saver._load_blobs(thread_id, ns, {"messages": v2}) - - assert "messages" in channel_values - result = channel_values["messages"] - assert isinstance(result, DiffChainValue) - assert result.base is None - assert result.deltas == [["msg1"], ["msg2"]] - - def test_diff_channel_mixed_old_and_new_blobs(self) -> None: - """When chain hits an old non-diff blob, it becomes base.""" - from langgraph.checkpoint.base import DiffChainValue, DiffDelta - - saver = InMemorySaver() - serde = JsonPlusSerializer() - - thread_id = "t2" - ns = "" - - v_old = "00000000000000000000000000000001.0000000000000000" - v_new = "00000000000000000000000000000002.0000000000000000" - - # Old-style full-list blob - saver.blobs[(thread_id, ns, "messages", v_old)] = serde.dumps_typed(["old_msg"]) - # New diff blob chained to old - delta = DiffDelta(delta=["new_msg"], prev_version=v_old) - saver.blobs[(thread_id, ns, "messages", v_new)] = serde.dumps_typed(delta) - - channel_values = saver._load_blobs(thread_id, ns, {"messages": v_new}) - result = channel_values["messages"] - assert isinstance(result, DiffChainValue) - assert result.base == ["old_msg"] - assert result.deltas == [["new_msg"]] -``` - -- [ ] **Step 2: Run to confirm tests fail** - -```bash -cd libs/checkpoint && TEST=tests/test_memory.py::TestInMemorySaverDiffChannel make test -``` -Expected: `FAILED` — diff channel values not wrapped in `DiffChainValue` - -- [ ] **Step 3: Update `_load_blobs` in `libs/checkpoint/langgraph/checkpoint/memory/__init__.py`** - -Locate `_load_blobs` (around line 123): -```python -def _load_blobs( - self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions -) -> dict[str, Any]: - channel_values: dict[str, Any] = {} - for k, v in versions.items(): - kk = (thread_id, checkpoint_ns, k, v) - if kk in self.blobs: - vv = self.blobs[kk] - if vv[0] != "empty": - channel_values[k] = self.serde.loads_typed(vv) - return channel_values -``` - -Replace with: -```python -def _load_blobs( - self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions -) -> dict[str, Any]: - from langgraph.checkpoint.base import DiffChainValue - - channel_values: dict[str, Any] = {} - diff_channels: dict[str, Any] = {} - - for k, v in versions.items(): - kk = (thread_id, checkpoint_ns, k, v) - if kk not in self.blobs: - continue - vv = self.blobs[kk] - if vv[0] == "diff": - diff_channels[k] = v - elif vv[0] != "empty": - channel_values[k] = self.serde.loads_typed(vv) - - for k, current_version in diff_channels.items(): - chain_deltas: list[list[Any]] = [] - base: list[Any] | None = None - version: str | None = current_version - while version is not None: - kk = (thread_id, checkpoint_ns, k, version) - if kk not in self.blobs: - break - vv = self.blobs[kk] - if vv[0] == "diff": - payload = self.serde.loads_typed(vv) # {"d": [...], "p": version|None} - chain_deltas.append(payload["d"]) - version = payload["p"] - else: - base = self.serde.loads_typed(vv) - break - chain_deltas.reverse() - channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas) - - return channel_values -``` - -- [ ] **Step 4: Run the new tests** - -```bash -cd libs/checkpoint && TEST=tests/test_memory.py::TestInMemorySaverDiffChannel make test -``` -Expected: both `PASSED` - -- [ ] **Step 5: Run the full memory test suite for regressions** - -```bash -cd libs/checkpoint && TEST=tests/test_memory.py make test -``` -Expected: all tests pass - -- [ ] **Step 6: Commit** - -```bash -git add libs/checkpoint/langgraph/checkpoint/memory/__init__.py libs/checkpoint/tests/test_memory.py -git commit -m "feat(checkpoint/memory): chain-traverse diff blobs in _load_blobs" -``` - ---- - -## Task 6: Update Pregel Layer (`channels_from_checkpoint` + `_put_checkpoint`) - -**Files:** -- Modify: `libs/langgraph/langgraph/pregel/_checkpoint.py` -- Modify: `libs/langgraph/langgraph/pregel/_loop.py` -- Modify: `libs/langgraph/tests/test_pregel.py` - -Two small changes: call `channel.after_checkpoint(version)` after `from_checkpoint` (so `DiffChannel` knows its starting version), and after `create_checkpoint` (so `DiffChannel` clears `_pending` and advances `_base_version`). - -- [ ] **Step 1: Write the failing end-to-end integration test in `libs/langgraph/tests/test_pregel.py`** - -Find the end of the test file and add: - -```python -async def test_diff_channel_end_to_end_inmemory() -> None: - """Full graph run: DiffChannel accumulates correctly across multiple turns.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.channels.diff import DiffChannel - from langgraph.checkpoint.memory import InMemorySaver - from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages - - class State(TypedDict): - messages: Annotated[list, DiffChannel(add_messages)] - - def respond(state: State) -> dict: - n = len(state["messages"]) - return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]} - - builder = StateGraph(State) - builder.add_node("respond", respond) - builder.add_edge(START, "respond") - graph = builder.compile(checkpointer=InMemorySaver()) - - config = {"configurable": {"thread_id": "diff-test-1"}} - - # Turn 1 - graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) - # Turn 2 - graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config) - # Turn 3 - graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config) - - state = graph.get_state(config) - msgs = state.values["messages"] - # 3 human + 3 AI = 6 total - assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}" - assert msgs[0].content == "hello" - assert msgs[2].content == "world" - assert msgs[4].content == "bye" - - -async def test_diff_channel_time_travel() -> None: - """Time-travel to an earlier checkpoint reconstructs the correct partial history.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.channels.diff import DiffChannel - from langgraph.checkpoint.memory import InMemorySaver - from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages - - class State(TypedDict): - messages: Annotated[list, DiffChannel(add_messages)] - - 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") - saver = InMemorySaver() - graph = builder.compile(checkpointer=saver) - - config = {"configurable": {"thread_id": "diff-time-travel"}} - - # Run 2 turns - graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config) - graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config) - - # Collect checkpoint history - history = list(graph.get_state_history(config)) - # history[0] = latest; find checkpoint after first turn (4 msgs: input + respond + input + respond) - # We want the state after first turn = 2 messages - after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2) - - assert len(after_turn1.values["messages"]) == 2 - assert after_turn1.values["messages"][0].content == "h1" -``` - -- [ ] **Step 2: Run to confirm they fail** - -```bash -cd libs/langgraph && TEST="tests/test_pregel.py::test_diff_channel_end_to_end_inmemory tests/test_pregel.py::test_diff_channel_time_travel" make test -``` -Expected: `FAILED` — `DiffChannel` does not receive `after_checkpoint` so `_base_version` is never set, causing each step to emit a root blob and reconstruction only returns the last step's messages. - -- [ ] **Step 3: Update `channels_from_checkpoint` in `libs/langgraph/langgraph/pregel/_checkpoint.py`** - -Locate `channels_from_checkpoint` (around line 58): -```python - return ( - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - for k, v in channel_specs.items() - }, - managed_specs, - ) -``` - -Replace the return statement with: -```python - 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)) - channels[k] = ch - return channels, managed_specs -``` - -- [ ] **Step 4: Update `_put_checkpoint` in `libs/langgraph/langgraph/pregel/_loop.py`** - -Find `_put_checkpoint` and locate where `self.checkpoint` is assigned from `create_checkpoint`. The relevant block (around line 877) is: - -```python -self.checkpoint = create_checkpoint( - self.checkpoint, - self.channels if do_checkpoint else None, - self.step, - id=self.checkpoint["id"] if exiting else None, - updated_channels=self.updated_channels, -) -``` - -Immediately after that assignment, add the `after_checkpoint` notification for all channels that were actually checkpointed (`do_checkpoint and self.channels is not None`): - -```python -self.checkpoint = create_checkpoint( - self.checkpoint, - self.channels if do_checkpoint else None, - self.step, - 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)) -``` - -- [ ] **Step 5: Run the integration tests** - -```bash -cd libs/langgraph && TEST="tests/test_pregel.py::test_diff_channel_end_to_end_inmemory tests/test_pregel.py::test_diff_channel_time_travel" make test -``` -Expected: both `PASSED` - -- [ ] **Step 6: Run the full pregel test suite for regressions (this is a large suite — may take several minutes)** - -```bash -cd libs/langgraph && make test -``` -Expected: all existing tests pass - -- [ ] **Step 7: Commit** - -```bash -git add libs/langgraph/langgraph/pregel/_checkpoint.py libs/langgraph/langgraph/pregel/_loop.py libs/langgraph/tests/test_pregel.py -git commit -m "feat(pregel): call after_checkpoint hook when loading and saving channels" -``` - ---- - -## Task 7: Extend `PostgresSaver._load_blobs` for Range-Query Chain Reconstruction - -**Files:** -- Modify: `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` -- Modify: `libs/checkpoint-postgres/tests/test_postgres.py` (or equivalent test file) - -After the existing JOIN fetches one blob per channel, detect channels with `type = "diff"` and issue one additional SQL range query per diff channel (typically just `messages`) to retrieve the full chain. - -- [ ] **Step 1: Find the Postgres test file** - -```bash -ls libs/checkpoint-postgres/tests/ -``` - -Use whatever test file exists (likely `test_postgres.py` or `test_async_postgres.py`). - -- [ ] **Step 2: Write failing Postgres diff-chain tests** - -These tests require a live Postgres instance. Add to the existing test class/file (check how the Postgres fixture is set up in the existing tests and reuse it): - -```python -async def test_diff_channel_postgres_chain_reconstruction(postgres_url: str) -> None: - """PostgresSaver reconstructs DiffChannel chain via range query.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.channels.diff import DiffChannel - from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver - from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages - - class State(TypedDict): - messages: Annotated[list, DiffChannel(add_messages)] - - def respond(state: State) -> dict: - n = len(state["messages"]) - return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]} - - builder = StateGraph(State) - builder.add_node("respond", respond) - builder.add_edge(START, "respond") - - async with AsyncPostgresSaver.from_conn_string(postgres_url) as saver: - await saver.setup() - graph = builder.compile(checkpointer=saver) - config = {"configurable": {"thread_id": "pg-diff-test-1"}} - - await graph.ainvoke( - {"messages": [HumanMessage(content="hi", id="h1")]}, config - ) - await graph.ainvoke( - {"messages": [HumanMessage(content="there", id="h2")]}, config - ) - - state = await graph.aget_state(config) - msgs = state.values["messages"] - assert len(msgs) == 4, f"expected 4, got {len(msgs)}" - assert msgs[0].content == "hi" - assert msgs[2].content == "there" -``` - -- [ ] **Step 3: Run to confirm test fails** - -```bash -cd libs/checkpoint-postgres && TEST=tests/test_postgres.py::test_diff_channel_postgres_chain_reconstruction make test -``` -Expected: `FAILED` — diff chain not assembled; only last delta returned - -- [ ] **Step 4: Update `_load_blobs` in `base.py` to accept context kwargs and detect diff channels** - -Locate `_load_blobs` in `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` (around line 187): -```python -def _load_blobs( - self, blob_values: list[tuple[bytes, bytes, bytes]] -) -> dict[str, Any]: - if not blob_values: - return {} - return { - k.decode(): self.serde.loads_typed((t.decode(), v)) - for k, t, v in blob_values - if t.decode() != "empty" - } -``` - -Replace with: -```python -def _load_blobs( - self, - blob_values: list[tuple[bytes, bytes, bytes]], - *, - thread_id: str = "", - checkpoint_ns: str = "", -) -> dict[str, Any]: - from langgraph.checkpoint.base import DiffChainValue - - if not blob_values: - return {} - - result: dict[str, Any] = {} - diff_channel_payloads: dict[str, dict[str, Any]] = {} - - for k, t, v in blob_values: - channel = k.decode() - type_tag = t.decode() - if type_tag == "diff": - diff_channel_payloads[channel] = self.serde.loads_typed((type_tag, v)) - elif type_tag != "empty": - result[channel] = self.serde.loads_typed((type_tag, v)) - - if diff_channel_payloads: - result.update( - self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads) - ) - - return result - -def _load_diff_chains( - self, - thread_id: str, - checkpoint_ns: str, - diff_channel_payloads: dict[str, dict[str, Any]], -) -> dict[str, Any]: - """Override in sync/async subclasses. Resolves diff-chain blobs to DiffChainValue.""" - raise NotImplementedError -``` - -- [ ] **Step 5: Override `_load_diff_chains` in sync `PostgresSaver` and update `_load_checkpoint_tuple` in both subclasses** - -**5a — Sync `PostgresSaver`** (`libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py`): - -Add this method to the `PostgresSaver` class (after `_load_checkpoint_tuple`): -```python -def _load_diff_chains( - self, - thread_id: str, - checkpoint_ns: str, - diff_channel_payloads: dict[str, dict[str, Any]], -) -> dict[str, Any]: - from langgraph.checkpoint.base import DiffChainValue - - result: dict[str, Any] = {} - with self._cursor() as cur: - for channel, current_payload in diff_channel_payloads.items(): - # Walk the prev_version chain backwards collecting payloads. - payloads: list[dict[str, Any]] = [current_payload] - version_cursor: str | None = current_payload["p"] - base: list[Any] | None = None - - while version_cursor is not None: - cur.execute( - "SELECT type, blob FROM checkpoint_blobs " - "WHERE thread_id = %s AND checkpoint_ns = %s " - "AND channel = %s AND version = %s", - (thread_id, checkpoint_ns, channel, version_cursor), - ) - row = cur.fetchone() - if row is None: - break - # row is a dict (dict_row factory): {"type": str, "blob": bytes} - if row["type"] == "diff": - payload = self.serde.loads_typed(("diff", row["blob"])) - payloads.append(payload) - version_cursor = payload["p"] - else: - base = self.serde.loads_typed((row["type"], row["blob"])) - break - - # payloads is newest→oldest; reverse for oldest→newest deltas - payloads.reverse() - result[channel] = DiffChainValue( - base=base, deltas=[p["d"] for p in payloads] - ) - return result -``` - -Update `_load_checkpoint_tuple` in sync `PostgresSaver` to pass `thread_id` and `checkpoint_ns`: -```python -def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: - return CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], - } - }, - { - **value["checkpoint"], - "channel_values": { - **(value["checkpoint"].get("channel_values") or {}), - **self._load_blobs( - value["channel_values"], - thread_id=value["thread_id"], - checkpoint_ns=value["checkpoint_ns"], - ), - }, - }, - value["metadata"], - ( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - self._load_writes(value["pending_writes"]), - ) -``` - -**5b — Async `AsyncPostgresSaver`** (`libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py`): - -Since `_load_checkpoint_tuple` is already `async`, override it to call an async version of the chain loader instead of `_load_blobs`. Add `_load_diff_chains_async` and update `_load_checkpoint_tuple`: - -```python -async def _load_diff_chains_async( - self, - thread_id: str, - checkpoint_ns: str, - diff_channel_payloads: dict[str, dict[str, Any]], -) -> dict[str, Any]: - from langgraph.checkpoint.base import DiffChainValue - - result: dict[str, Any] = {} - async with self._cursor() as cur: - for channel, current_payload in diff_channel_payloads.items(): - payloads: list[dict[str, Any]] = [current_payload] - version_cursor: str | None = current_payload["p"] - base: list[Any] | None = None - - while version_cursor is not None: - await cur.execute( - "SELECT type, blob FROM checkpoint_blobs " - "WHERE thread_id = %s AND checkpoint_ns = %s " - "AND channel = %s AND version = %s", - (thread_id, checkpoint_ns, channel, version_cursor), - ) - row = await cur.fetchone() - if row is None: - break - if row["type"] == "diff": - payload = self.serde.loads_typed(("diff", row["blob"])) - payloads.append(payload) - version_cursor = payload["p"] - else: - base = self.serde.loads_typed((row["type"], row["blob"])) - break - - payloads.reverse() - result[channel] = DiffChainValue( - base=base, deltas=[p["d"] for p in payloads] - ) - return result - -async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: - thread_id = value["thread_id"] - checkpoint_ns = value["checkpoint_ns"] - blob_values = value["channel_values"] - - # Load non-diff channels synchronously using the base _load_blobs, - # but intercept diff channels for async resolution below. - from langgraph.checkpoint.base import DiffChainValue - non_diff: dict[str, Any] = {} - diff_payloads: dict[str, dict[str, Any]] = {} - if blob_values: - for k, t, v in blob_values: - channel = k.decode() - type_tag = t.decode() - if type_tag == "diff": - diff_payloads[channel] = self.serde.loads_typed((type_tag, v)) - elif type_tag != "empty": - non_diff[channel] = self.serde.loads_typed((type_tag, v)) - - diff_values = ( - await self._load_diff_chains_async(thread_id, checkpoint_ns, diff_payloads) - if diff_payloads - else {} - ) - - return CheckpointTuple( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["checkpoint_id"], - } - }, - { - **value["checkpoint"], - "channel_values": { - **(value["checkpoint"].get("channel_values") or {}), - **non_diff, - **diff_values, - }, - }, - value["metadata"], - ( - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None - ), - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) -``` - -- [ ] **Step 6: Run Postgres diff-channel test** - -```bash -cd libs/checkpoint-postgres && TEST=tests/test_postgres.py::test_diff_channel_postgres_chain_reconstruction make test -``` -Expected: `PASSED` - -- [ ] **Step 7: Run full Postgres test suite for regressions** - -```bash -cd libs/checkpoint-postgres && make test -``` -Expected: all existing tests pass - -- [ ] **Step 8: Commit** - -```bash -git add libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py libs/checkpoint-postgres/tests/test_postgres.py -git commit -m "feat(checkpoint/postgres): range-query diff chain reconstruction in _load_blobs" -``` - ---- - -## Task 8: Format, Lint, and Final Integration Check - -**Files:** All modified libraries - -- [ ] **Step 1: Format and lint `libs/checkpoint`** - -```bash -cd libs/checkpoint && make format && make lint -``` -Fix any issues, then re-run until clean. - -- [ ] **Step 2: Format and lint `libs/langgraph`** - -```bash -cd libs/langgraph && make format && make lint -``` -Fix any issues. - -- [ ] **Step 3: Format and lint `libs/checkpoint-postgres`** - -```bash -cd libs/checkpoint-postgres && make format && make lint -``` -Fix any issues. - -- [ ] **Step 4: Run all tests across affected libraries** - -```bash -cd libs/checkpoint && make test -cd libs/langgraph && make test -``` -Expected: all green. - -- [ ] **Step 5: Verify `DiffChannel` is importable from the public API** - -```bash -python -c " -from langgraph.channels import DiffChannel -from langgraph.channels.diff import DiffChannel as DC2 -from langgraph.checkpoint.base import DiffDelta, DiffChainValue -from langgraph.graph.message import add_messages -from typing import Annotated, TypedDict -from langchain_core.messages import AnyMessage - -class State(TypedDict): - messages: Annotated[list[AnyMessage], DiffChannel(add_messages)] - -print('DiffChannel public API: ok') -" -``` -Expected: `DiffChannel public API: ok` - -- [ ] **Step 6: Commit final cleanup** - -```bash -git add -p # stage any remaining formatting changes -git commit -m "chore: format and lint DiffChannel implementation" -``` - ---- - -## Implementation Notes - -**PostgresSaver `_load_diff_chains` complexity:** Task 7 Step 5 describes two approaches. The single bulk range query (using `ANY(%s)` with channel list and `version <= %s` per channel) is preferred for production but requires passing `thread_id`, `checkpoint_ns`, and `channel_versions` down to `_load_blobs`. The per-link single-row query is simpler to implement first. Implement the bulk query for the async saver since async Postgres is the common production path. - -**AsyncPostgresSaver:** The async variant in `libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py` has its own `_load_blobs` and `get_tuple`. Apply the same changes there, using `await cur.execute` / `await cur.fetchall` instead of synchronous cursor calls. - -**`_pending` accumulation across tasks in one step:** Within one superstep, multiple tasks may write to the same channel. `apply_writes` calls `channel.update(values)` where `values` is the list of all writes from all tasks. `_pending` accumulates all of them. The `DiffDelta.delta` for that step is the full list of writes from all tasks — this is correct. - -**`DiffChannel` with `MISSING` initial value:** `from_checkpoint(MISSING)` sets `value = []`. The channel is "available" if `value` is not `MISSING` (the MISSING sentinel from `_internal._typing`). Check `is_available()` returns `True` even for an empty list — messages channel starts empty then gets populated. - -**`EmptyChannelError` on `get()`:** The current implementation raises if `self.value is MISSING`. But after `from_checkpoint(MISSING)`, `self.value = []` (not MISSING). This means `get()` returns `[]` rather than raising `EmptyChannelError` for a never-updated DiffChannel. This matches `add_messages` semantics where an empty messages list is valid state. If `EmptyChannelError` is needed (e.g. for graph routing), adjust `is_available` to also check `bool(self.value)`. diff --git a/docs/superpowers/specs/2026-04-17-diff-channel-design.md b/docs/superpowers/specs/2026-04-17-diff-channel-design.md deleted file mode 100644 index 7aa15b139..000000000 --- a/docs/superpowers/specs/2026-04-17-diff-channel-design.md +++ /dev/null @@ -1,405 +0,0 @@ -# DiffChannel: Incremental Checkpoint Storage for Append-Style Reducers - -**Date:** 2026-04-17 -**Status:** Approved for implementation -**Scope:** `libs/checkpoint`, `libs/langgraph`, `libs/checkpoint-postgres` - ---- - -## Motivation - -LangGraph checkpoints today store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, this means each checkpoint blob contains the entire conversation history. Storage cost grows O(N²) in the number of turns: step 1 stores 1 message, step 100 stores 100 messages, step 1000 stores 1000 messages. For long-running agentic conversations with high-token messages this is untenable. - -The fix is to store only the **delta** (new writes) per step, reconstructing the full accumulated value at load time by replaying the chain. This is an opt-in mechanism — existing graphs are unaffected. - ---- - -## Non-Goals - -- **Compaction / materialized snapshots**: deferred. Load cost stays O(N) blob fetches but those fetches are batched into a single query — acceptable for now. -- **SQLite saver support**: SQLite stores all channel values inline in one row (no per-channel blob table). Deferred to a follow-up. -- **Automatic migration** of existing `BinaryOperatorAggregate` channels: users opt in explicitly. Old checkpoints load correctly via the backwards-compatibility path in `from_checkpoint`. - ---- - -## Architecture Overview - -``` -User state definition - └── Annotated[list[AnyMessage], DiffChannel(add_messages)] - -Write path (per superstep) - DiffChannel.update() — apply operator, accumulate writes in _pending - DiffChannel.checkpoint() — return DiffDelta(delta=_pending, prev_version=_base_version) - serde.dumps_typed() — serialize DiffDelta as ("diff", msgpack_bytes) - saver.put() — store blob at (thread_id, ns, "messages", version_N) - DiffChannel.after_checkpoint(version_N) — advance _base_version, clear _pending - -Read path (on graph load or time-travel) - saver.get_tuple() — fetch current-version blob per channel - saver._load_blobs() — detect "diff" type → follow chain to reconstruct DiffChainValue - DiffChannel.from_checkpoint(DiffChainValue) — replay deltas with operator → full list - DiffChannel.after_checkpoint(version_N) — set _base_version for next write -``` - -The pregel layer (`_checkpoint.py`, `_loop.py`) is unchanged except for two small additions to call the new `after_checkpoint` hook. The saver public interface (`BaseCheckpointSaver`) gains no new methods. All chain-following logic lives inside each saver's private `_load_blobs`. - ---- - -## New Protocol Types - -**Location:** `libs/checkpoint/langgraph/checkpoint/base/__init__.py` - -Two dataclasses form the contract between `DiffChannel` and savers: - -```python -@dataclass -class DiffDelta: - """Returned by DiffChannel.checkpoint(). Written to the blob store.""" - delta: list[Any] # raw writes passed to update() this step - prev_version: str | None # version of the previous diff blob; None = chain root -``` - -```python -@dataclass -class DiffChainValue: - """Passed to DiffChannel.from_checkpoint(). Assembled by _load_blobs().""" - base: list[Any] | None # starting accumulated value (None = empty start) - deltas: list[list[Any]] # write-sets ordered oldest → newest -``` - -`DiffDelta` lives in the checkpoint base package (not the channel module) so savers can import it without creating a circular dependency. `DiffChainValue` is there for the same reason. - ---- - -## `BaseChannel.after_checkpoint()` Hook - -**Location:** `libs/langgraph/langgraph/channels/base.py` - -```python -def after_checkpoint(self, version: Any) -> None: - """Called after checkpoint() (with the new version) and after from_checkpoint() - (with the current version). No-op by default; DiffChannel overrides.""" - pass -``` - -This is a **non-abstract, no-op default** — fully backwards compatible. All existing channels inherit it silently. It is NOT in the abstract interface. - ---- - -## `DiffChannel[V]` - -**Location:** `libs/langgraph/langgraph/channels/diff.py` (new file) - -### Internal state - -| Attribute | Type | Description | -|---|---|---| -| `value` | `list[V]` | Full accumulated value (the reconstructed list) | -| `operator` | `Callable` | The binary reducer (e.g. `add_messages`) | -| `_pending` | `list[Any]` | Raw writes accumulated since last `after_checkpoint` call | -| `_base_version` | `str \| None` | Version this channel was last checkpointed at (= `prev_version` for next delta) | -| `_overwritten` | `bool` | True if an `Overwrite` was applied since last `after_checkpoint`; makes next blob a chain root | - -### `update(values)` - -Mirrors `BinaryOperatorAggregate.update()` with two additions: - -1. For each non-Overwrite value: apply `self.operator(self.value, value)` as before; **also append the raw incoming value to `self._pending`**. -2. For an `Overwrite(v)` value: set `self.value = v`; set `self._pending = list(v)` (full value becomes the new delta); set `self._overwritten = True`. - -The key: `_pending` stores the **incoming writes** (what was passed to `update()`), not the diff of `self.value`. This is important because `add_messages` handles removal and update-by-ID — replaying the writes with `operator` during reconstruction applies that logic correctly. - -### `checkpoint()` - -```python -def checkpoint(self) -> DiffDelta: - return DiffDelta( - delta=self._pending[:], - prev_version=None if self._overwritten else self._base_version, - ) -``` - -- Normal step: `prev_version = self._base_version` → chain link -- After Overwrite: `prev_version = None` → chain root (reconstruction stops here and uses `delta` as the full base value) - -Returns `DiffDelta`, never the raw accumulated list. The serde handles serialization. - -### `from_checkpoint(checkpoint)` - -```python -def from_checkpoint(self, checkpoint) -> Self: - new = DiffChannel(self.typ, self.operator) - new.key = self.key - if checkpoint is MISSING: - new.value = [] - elif isinstance(checkpoint, DiffChainValue): - accumulated = checkpoint.base or [] - for step_writes in checkpoint.deltas: - # Mirror update() exactly: apply each write individually so operator - # semantics (e.g. add_messages ID-based removal) are respected. - for write in step_writes: - accumulated = new.operator(accumulated, write) - new.value = accumulated - elif isinstance(checkpoint, DiffDelta): - # Unsupported saver: _load_blobs returned a raw DiffDelta instead of - # assembling a DiffChainValue. Raise rather than silently losing history. - 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 = checkpoint - new._pending = [] - new._base_version = None # set by the subsequent after_checkpoint() call - return new -``` - -The operator is available on `self` (the channel spec) so reconstruction is correct for any reducer — the saver never needs to know about `add_messages`. - -`_pending` stores **individual writes** (each `value` from `update()`'s `values` sequence), so each `step_writes` list in `DiffChainValue.deltas` is replayed write-by-write — identical to the `update()` loop. - -### `after_checkpoint(version)` - -```python -def after_checkpoint(self, version: Any) -> None: - if version != self._base_version: - self._base_version = version - self._pending = [] - self._overwritten = False -``` - -No-op when `version == self._base_version` (channel wasn't updated this step — blob was not written). Clears `_pending` and advances `_base_version` when the channel was actually checkpointed. - -### Opt-in API - -```python -from langgraph.channels.diff import DiffChannel - -class State(TypedDict): - messages: Annotated[list[AnyMessage], DiffChannel(add_messages)] -``` - -`StateGraph` already handles `BaseChannel` instances as annotation metadata — `DiffChannel` inherits this without any changes to `StateGraph`. - ---- - -## Serde Extension - -**Location:** `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` - -Add one branch to `dumps_typed` (before the `else` msgpack fallback), using the existing module-level `_msgpack_enc` so message ext-types (Pydantic v2, etc.) are handled correctly: - -```python -elif isinstance(obj, DiffDelta): - return "diff", _msgpack_enc({"d": obj.delta, "p": obj.prev_version}) -``` - -Add one branch to `loads_typed` so savers can decode diff blobs without importing `ormsgpack` directly: - -```python -elif type_ == "diff": - return ormsgpack.unpackb( - data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS - ) - # returns {"d": [writes...], "p": prev_version_str_or_none} -``` - -Savers call `serde.loads_typed(("diff", raw_bytes))` to decode a diff blob into `{"d": ..., "p": ...}`, then check `type_tag == "diff"` to trigger chain traversal. The serde layer is the only place that knows about `ormsgpack`. - ---- - -## Saver Changes - -### InMemorySaver - -**`put()` — `libs/checkpoint/langgraph/checkpoint/memory/__init__.py`** - -No change needed. The existing `self.serde.dumps_typed(values[k])` call already handles `DiffDelta` via the new serde branch above, storing it as `("diff", bytes)`. - -**`_load_blobs()` — same file** - -After checking `vv[0] != "empty"`, add a branch for `"diff"` before calling `serde.loads_typed`: - -```python -def _load_blobs(self, thread_id, checkpoint_ns, versions): - channel_values = {} - diff_channels = {} # channel_name -> current_version for diff channels - - for k, v in versions.items(): - kk = (thread_id, checkpoint_ns, k, v) - if kk not in self.blobs: - continue - type_tag, blob_bytes = self.blobs[kk] - if type_tag == "diff": - diff_channels[k] = v # handle below - elif type_tag != "empty": - channel_values[k] = self.serde.loads_typed((type_tag, blob_bytes)) - - for k, current_version in diff_channels.items(): - # Follow chain: newest → oldest, then reverse - chain_deltas = [] - base = None - version = current_version - while version is not None: - kk = (thread_id, checkpoint_ns, k, version) - if kk not in self.blobs: - break - type_tag, blob_bytes = self.blobs[kk] - if type_tag == "diff": - # Use serde so we don't need to import ormsgpack directly - payload = self.serde.loads_typed((type_tag, blob_bytes)) - chain_deltas.append(payload["d"]) - version = payload["p"] # prev_version; None = root - else: - # Old non-diff blob encountered: treat as base accumulated value - base = self.serde.loads_typed((type_tag, blob_bytes)) - break - chain_deltas.reverse() - channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas) - - return channel_values -``` - -Each blob lookup is O(1) on the dict. Total: N dict lookups for a chain of depth N. Memory usage is identical to loading a single full-list blob (same total bytes, split across N entries). - -### PostgresSaver - -**`_load_blobs()` — `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py`** - -The existing `SELECT_SQL` fetches one blob per channel via a JOIN. After running that query, detect any `"diff"` channels in the result and issue one additional range query: - -```python -def _load_blobs(self, blob_values): - if not blob_values: - return {} - - result = {} - diff_channels = {} # channel_name -> current_version (as str) - - for k, t, v in blob_values: - channel = k.decode() - type_tag = t.decode() - if type_tag == "diff": - # Decode via serde — no direct ormsgpack import needed - payload = self.serde.loads_typed((type_tag, v)) - diff_channels[channel] = payload # store for chain fetch - elif type_tag != "empty": - result[channel] = self.serde.loads_typed((type_tag, v)) - - if diff_channels: - result.update(self._load_diff_chains(diff_channels)) - - return result -``` - -`_load_diff_chains` issues one SQL query per diff channel (typically just `messages`): - -```sql -SELECT version, type, blob -FROM checkpoint_blobs -WHERE thread_id = %s - AND checkpoint_ns = %s - AND channel = %s - AND version <= %s -ORDER BY version ASC -``` - -In Python, iterate rows in ascending version order: if `type = "diff"`, accumulate the delta; if any other type is encountered, treat it as the base accumulated value and stop. Return `DiffChainValue(base=..., deltas=[...])`. - -This results in **at most 2 queries total** for a graph with one `DiffChannel` — existing behaviour for all other channels is unchanged. - -**`put()` / `_dump_blobs()`** - -No change needed. `_dump_blobs` calls `self.serde.dumps_typed(v)` for each channel value in `new_versions`. When `v` is a `DiffDelta`, the serde produces `("diff", bytes)` which is stored as `type = "diff"` in `checkpoint_blobs`. The `ON CONFLICT DO NOTHING` semantics are preserved. - -### SQLite - -Deferred. `SqliteSaver` stores the entire checkpoint as a single serialized row — it has no per-channel blob table. Supporting `DiffChannel` on SQLite would require adding a new blobs table, which is a separate migration tracked separately. - ---- - -## Pregel Layer Changes - -### `channels_from_checkpoint` — `libs/langgraph/langgraph/pregel/_checkpoint.py` - -After constructing each channel from its checkpoint value, call `after_checkpoint` so the channel records its current version: - -```python -channels = {} -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)) - channels[k] = ch -return channels, managed_specs -``` - -Existing channels get the no-op `after_checkpoint`. `DiffChannel` uses it to set `_base_version`. - -### `PregelLoop._put_checkpoint` — `libs/langgraph/langgraph/pregel/_loop.py` - -After `create_checkpoint(self.checkpoint, self.channels, self.step, ...)` returns and `do_checkpoint is True` and `self.channels is not None`, iterate channels and notify: - -```python -if do_checkpoint and self.channels: - for k, ch in self.channels.items(): - ch.after_checkpoint(self.checkpoint["channel_versions"].get(k)) -``` - -This is called after `create_checkpoint` updates `self.checkpoint["channel_versions"]`, so `get(k)` returns the new version for updated channels and the old version for unchanged ones. `DiffChannel.after_checkpoint` only clears `_pending` when `version != _base_version`, so unchanged channels are no-ops. - ---- - -## Backwards Compatibility - -| Scenario | Behaviour | -|---|---| -| Existing graph using `add_messages` (BinaryOperatorAggregate) | Unaffected — no code changes, no data migration | -| New graph with `DiffChannel`, loading old checkpoint blobs | `from_checkpoint` receives a plain `list` → used directly as accumulated value | -| `DiffChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported | -| `DiffChannel` with `SqliteSaver` | `from_checkpoint` receives a raw `DiffDelta` (SqliteSaver stores channel_values inline), raises `ValueError` with a clear message pointing to supported savers | -| Time-travel / fork to past checkpoint | Chain traversal uses the version at that checkpoint → reconstruction is correct | -| `update_state` | Treated as a normal step: writes are deltas chained to history | -| `Overwrite` value | Resets chain: next blob has `prev_version=None`; reconstruction starts fresh | - ---- - -## Testing Strategy - -1. **Unit tests for `DiffChannel`** (`libs/langgraph/tests/`): - - `update` → `checkpoint` → `after_checkpoint` → `checkpoint` lifecycle (2 steps, verify delta isolation) - - `from_checkpoint(DiffChainValue)` correctly replays multi-step chains using the operator - - `from_checkpoint(plain_list)` backwards-compat path - - `Overwrite` creates a root blob (`prev_version=None`) and reconstruction ignores prior chain - - `after_checkpoint` no-ops when version is unchanged - -2. **Integration tests with `InMemorySaver`** (`libs/langgraph/tests/`): - - 10-step conversation: verify final loaded state equals full accumulated messages - - Time-travel: fork to step 5, verify only messages 1–5 are present - - Mixed graph: some channels `BinaryOperatorAggregate`, one `DiffChannel` — both reconstruct correctly - -3. **Serde tests** (`libs/checkpoint/tests/`): - - `DiffDelta` round-trips through `dumps_typed` / saver storage - - Old `"msgpack"` blob for a channel → `DiffChannel.from_checkpoint` handles it - -4. **Postgres integration tests** (`libs/checkpoint-postgres/tests/`): - - Range query reconstructs correct full list after N steps - - Time-travel to checkpoint M reconstructs correct list of M messages - ---- - -## Files Changed - -| File | Change | -|---|---| -| `libs/checkpoint/langgraph/checkpoint/base/__init__.py` | Add `DiffDelta`, `DiffChainValue` dataclasses | -| `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` | Add `"diff"` branch in `dumps_typed` | -| `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` | Chain traversal in `_load_blobs` | -| `libs/langgraph/langgraph/channels/base.py` | Add no-op `after_checkpoint` method | -| `libs/langgraph/langgraph/channels/diff.py` | **New file** — `DiffChannel` implementation | -| `libs/langgraph/langgraph/channels/__init__.py` | Export `DiffChannel` | -| `libs/langgraph/langgraph/pregel/_checkpoint.py` | Call `after_checkpoint` in `channels_from_checkpoint` | -| `libs/langgraph/langgraph/pregel/_loop.py` | Call `after_checkpoint` after `create_checkpoint` | -| `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` | Range-query chain reconstruction in `_load_blobs` |