From a6731f436700c7d46a9224ccd3a0e0ce3a5bfdb5 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Tue, 13 Oct 2026 20:47:19 -0400 Subject: [PATCH] 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 | 4 + .../langgraph/langgraph/pregel/_checkpoint.py | 41 +++-- .../tests/test_delta_channel_to_plain.py | 148 ++++++++++++++++++ 3 files changed, 177 insertions(+), 16 deletions(-) create mode 100644 libs/langgraph/tests/test_delta_channel_to_plain.py diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index 58346198c..6d4b20b1e 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 1583d5a96..a0cf7809f 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/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 diff --git a/libs/langgraph/tests/test_delta_channel_to_plain.py b/libs/langgraph/tests/test_delta_channel_to_plain.py new file mode 100644 index 000000000..c01085239 --- /dev/null +++ b/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()