fix(langgraph): keep a DeltaChannel's value when its field changes to a plain reducer (#9302)

Changing a field from a `DeltaChannel` to a plain reducer, for example
to `Annotated[list, add_messages]`, made threads written under the
`DeltaChannel` read that field as empty. The next run saved its writes
over the empty value, so continuing an old thread replaced its history.
The writes were still stored: the read only rebuilt fields whose current
annotation is a `DeltaChannel`.

- A `BinaryOperatorAggregate` stores its value at every version, so one
with a version and nothing stored was written as a `DeltaChannel`. It's
now rebuilt from the same history walk, one `update` per write, since a
single call would reject a second `Overwrite` and drop the writes after
one.
- `BinaryOperatorAggregate.from_checkpoint` unwraps a stored
`_DeltaSnapshot`.
- The next run saves the full value, so later reads are plain reads, and
switching back to a `DeltaChannel` keeps everything.

Threads written under a plain reducer never walk. A rebuilt read costs
the walk plus one reducer call per write since the last snapshot: with
`add_messages` on `InMemorySaver`, 25 ms at 2,048 messages with a
snapshot every 50 writes, and 0.3 s at 1,998 or 0.9 s at 4,998 messages
with the default 1000. After one run, reads take 6 to 16 ms.

`test_delta_channel_to_plain.py` covers every saver; all but the
plain-thread test fail on main.

JS: langchain-ai/langgraphjs#3017

libs/langgraph/langgraph/channels/binop.py

@@ -2,6 +2,7 @@ import collections.abc
 from collections.abc import Callable, Sequence
 from typing import Any, Generic
 
+from langgraph.checkpoint.serde.types import _DeltaSnapshot
 from typing_extensions import NotRequired, Required, Self
 
 from langgraph._internal._constants import OVERWRITE
@@ -116,6 +117,9 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     def from_checkpoint(self, checkpoint: Value) -> Self:
         empty = self.__class__(self.typ, self.operator)
         empty.key = self.key
+        # A thread written while this field was a DeltaChannel stores snapshots.
+        if isinstance(checkpoint, _DeltaSnapshot):
+            checkpoint = checkpoint.value
         if checkpoint is not MISSING:
             empty.value = checkpoint
         return empty

libs/langgraph/langgraph/pregel/_checkpoint.py

@@ -4,7 +4,7 @@ import uuid
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from datetime import datetime, timezone
 from inspect import signature
-from typing import Any, Literal, cast
+from typing import Any, Literal
 
 from langchain_core.runnables import RunnableConfig
 from langgraph.checkpoint.base import (
@@ -12,6 +12,7 @@ from langgraph.checkpoint.base import (
     ChannelVersions,
     Checkpoint,
     CheckpointTuple,
+    DeltaChannelHistory,
     PendingWrite,
 )
 from langgraph.checkpoint.base.id import uuid6
@@ -32,7 +33,7 @@ from langgraph._internal._constants import (
 )
 from langgraph._internal._typing import MISSING
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.binop import _get_overwrite
+from langgraph.channels.binop import BinaryOperatorAggregate, _get_overwrite
 from langgraph.channels.delta import DeltaChannel
 from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
 
@@ -409,23 +410,39 @@ def versions_seen_without_bumps(
 def _delta_channels_to_replay(
     specs: Mapping[str, BaseChannel], checkpoint: Checkpoint
 ) -> list[str]:
-    """DeltaChannels whose value at `checkpoint` the ancestor walk rebuilds.
+    """Channels whose value at `checkpoint` the ancestor walk rebuilds.
 
     A `_DeltaSnapshot` blob or a plain value (migration) resolves directly via
     `from_checkpoint`, so only a channel with nothing stored here needs the
     walk. A channel with no version was never written, so it is empty without
     one; a walk for it would find no snapshot to stop at and read every
     ancestor, every time the thread is loaded.
+
+    A `BinaryOperatorAggregate` stores its value at every version, so one with
+    nothing stored was written as a `DeltaChannel` before its annotation
+    changed, and is rebuilt the same way.
     """
     return [
         k
         for k, spec in specs.items()
-        if isinstance(spec, DeltaChannel)
+        if isinstance(spec, (DeltaChannel, BinaryOperatorAggregate))
         and k in checkpoint["channel_versions"]
         and checkpoint["channel_values"].get(k, MISSING) is MISSING
     ]
 
 
+def _replay(spec: BaseChannel, history: DeltaChannelHistory) -> BaseChannel:
+    ch = spec.from_checkpoint(history.get("seed", MISSING))
+    if isinstance(ch, DeltaChannel):
+        ch.replay_writes(history["writes"])
+    else:
+        # One update per write: the writes span many supersteps, and a single
+        # call would reject a second Overwrite and drop the writes after one.
+        for _, _, value in history["writes"]:
+            ch.update([value])
+    return ch
+
+
 def _require_saver_for_history(
     delta_channels: list[str],
     saver: BaseCheckpointSaver | None,
@@ -465,7 +482,7 @@ def channels_from_checkpoint(
 
     delta_channels = _delta_channels_to_replay(channel_specs, checkpoint)
     _require_saver_for_history(delta_channels, saver, config)
-    histories: Mapping[str, Any] = {}
+    histories: Mapping[str, DeltaChannelHistory] = {}
     if delta_channels and saver is not None and config is not None:
         histories = saver.get_delta_channel_history(
             config=config, channels=delta_channels
@@ -475,11 +492,7 @@ def channels_from_checkpoint(
     for k, spec in channel_specs.items():
         ch: BaseChannel
         if k in histories:
-            delta_spec = cast(DeltaChannel, spec)
-            history = histories[k]
-            replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING))
-            replay_ch.replay_writes(history["writes"])
-            ch = replay_ch
+            ch = _replay(spec, histories[k])
         else:
             ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
         channels[k] = ch
@@ -504,7 +517,7 @@ async def achannels_from_checkpoint(
 
     delta_channels = _delta_channels_to_replay(channel_specs, checkpoint)
     _require_saver_for_history(delta_channels, saver, config)
-    histories: Mapping[str, Any] = {}
+    histories: Mapping[str, DeltaChannelHistory] = {}
     if delta_channels and saver is not None and config is not None:
         histories = await saver.aget_delta_channel_history(
             config=config, channels=delta_channels
@@ -514,11 +527,7 @@ async def achannels_from_checkpoint(
     for k, spec in channel_specs.items():
         ch: BaseChannel
         if k in histories:
-            delta_spec = cast(DeltaChannel, spec)
-            history = histories[k]
-            replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING))
-            replay_ch.replay_writes(history["writes"])
-            ch = replay_ch
+            ch = _replay(spec, histories[k])
         else:
             ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
         channels[k] = ch

libs/langgraph/tests/test_delta_channel_to_plain.py

@@ -0,0 +1,148 @@
+import operator
+from typing import Annotated, Any
+from unittest.mock import patch
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.serde.types import _DeltaSnapshot
+from typing_extensions import TypedDict
+
+from langgraph.channels.binop import BinaryOperatorAggregate
+from langgraph.channels.delta import DeltaChannel
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import _messages_delta_reducer, add_messages
+from langgraph.pregel._checkpoint import channels_from_checkpoint, empty_checkpoint
+from langgraph.types import Overwrite
+
+pytestmark = pytest.mark.anyio
+
+
+class DeltaState(TypedDict):
+    messages: Annotated[
+        list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3)
+    ]
+
+
+class PlainState(TypedDict):
+    messages: Annotated[list, add_messages]
+
+
+def _app(schema: type, checkpointer: BaseCheckpointSaver) -> Any:
+    builder = StateGraph(schema)
+    builder.add_node(
+        "chat",
+        lambda state: {"messages": [AIMessage(f"reply {len(state['messages'])}")]},
+    )
+    builder.add_edge(START, "chat")
+    builder.add_edge("chat", END)
+    return builder.compile(checkpointer=checkpointer)
+
+
+def _delta_thread(checkpointer: BaseCheckpointSaver) -> RunnableConfig:
+    config: RunnableConfig = {"configurable": {"thread_id": "1"}}
+    delta = _app(DeltaState, checkpointer)
+    for i in range(4):
+        delta.invoke({"messages": [HumanMessage(f"q{i}")]}, config)
+    assert any(
+        isinstance(t.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
+        for t in checkpointer.list(config)
+    ), "the thread should have a snapshot to read through"
+    return config
+
+
+def test_plain_reducer_reads_every_delta_written_checkpoint(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    config = _delta_thread(sync_checkpointer)
+    delta = _app(DeltaState, sync_checkpointer)
+    plain = _app(PlainState, sync_checkpointer)
+
+    assert [s.values for s in plain.get_state_history(config)] == [
+        s.values for s in delta.get_state_history(config)
+    ]
+
+
+def test_plain_reducer_continues_delta_written_thread(
+    sync_checkpointer: BaseCheckpointSaver,
+) -> None:
+    config = _delta_thread(sync_checkpointer)
+    delta = _app(DeltaState, sync_checkpointer)
+    before = delta.get_state(config).values["messages"]
+
+    after = _app(PlainState, sync_checkpointer).invoke(
+        {"messages": [HumanMessage("q4")]}, config
+    )["messages"]
+
+    assert after[:-2] == before
+    assert [m.content for m in after[-2:]] == ["q4", f"reply {len(before) + 1}"]
+    head = sync_checkpointer.get_tuple(config)
+    assert head is not None
+    assert head.checkpoint["channel_values"]["messages"] == after
+    assert delta.get_state(config).values["messages"] == after
+
+
+async def test_plain_reducer_reads_and_continues_delta_written_thread_async(
+    async_checkpointer: BaseCheckpointSaver,
+) -> None:
+    config: RunnableConfig = {"configurable": {"thread_id": "1"}}
+    delta = _app(DeltaState, async_checkpointer)
+    for i in range(4):
+        await delta.ainvoke({"messages": [HumanMessage(f"q{i}")]}, config)
+    plain = _app(PlainState, async_checkpointer)
+
+    expected = [s.values async for s in delta.aget_state_history(config)]
+    assert [s.values async for s in plain.aget_state_history(config)] == expected
+
+    after = (await plain.ainvoke({"messages": [HumanMessage("q4")]}, config))[
+        "messages"
+    ]
+    assert after[:-2] == expected[0]["messages"]
+
+
+@pytest.mark.parametrize(
+    ("seed", "writes", "expected"),
+    [
+        (_DeltaSnapshot(["a"]), [["b"], ["c"]], ["a", "b", "c"]),
+        (["a"], [["b"], Overwrite(["c"]), ["d"]], ["c", "d"]),
+        (None, [["b"], Overwrite(["c"]), Overwrite(["d"]), ["e"]], ["d", "e"]),
+    ],
+)
+def test_plain_reducer_replays_delta_history_one_write_at_a_time(
+    seed: Any, writes: list[Any], expected: list[str]
+) -> None:
+    history: dict[str, Any] = {"writes": [("t", "items", w) for w in writes]}
+    if seed is not None:
+        history["seed"] = seed
+    checkpoint = empty_checkpoint()
+    checkpoint["channel_versions"]["items"] = 1
+    saver = InMemorySaver()
+
+    with patch.object(
+        saver, "get_delta_channel_history", return_value={"items": history}
+    ):
+        channels, _ = channels_from_checkpoint(
+            {"items": BinaryOperatorAggregate(list, operator.add)},
+            checkpoint,
+            saver=saver,
+            config={"configurable": {"thread_id": "1"}},
+        )
+
+    assert channels["items"].get() == expected
+
+
+def test_plain_thread_reads_never_walk_history() -> None:
+    saver = InMemorySaver()
+    config: RunnableConfig = {"configurable": {"thread_id": "1"}}
+    plain = _app(PlainState, saver)
+
+    with patch.object(saver, "get_delta_channel_history") as walk:
+        plain.invoke({"messages": [HumanMessage("q0")]}, config)
+        plain.update_state(config, {"messages": [HumanMessage("note")]})
+        plain.invoke({"messages": [HumanMessage("q1")]}, config)
+        plain.get_state(config)
+        list(plain.get_state_history(config))
+
+    walk.assert_not_called()