From 03234e34b6a20d4b5d4da5df36ed607edb63afa9 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Fri, 9 Oct 2026 12:02:21 -0400 Subject: [PATCH] fix(langgraph): snapshot a DeltaChannel that an update_state Overwrite resets A node's Overwrite makes the loop snapshot the channel on the checkpoint its superstep saves. The same write through update_state, as a node or as input, didn't snapshot, so reads walked back past the reset to the last snapshot. update_state now snapshots it too. --- .../langgraph/langgraph/pregel/_checkpoint.py | 16 +++ libs/langgraph/langgraph/pregel/main.py | 13 +++ .../test_delta_channel_update_overwrite.py | 110 ++++++++++++++++++ 3 files changed, 139 insertions(+) create mode 100644 libs/langgraph/tests/test_delta_channel_update_overwrite.py diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index b2cddc19d..f5a436b0e 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -32,6 +32,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.delta import DeltaChannel from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec @@ -152,6 +153,21 @@ def delta_channels_with_pending_writes( } +def delta_channels_overwritten( + specs: Mapping[str, Any], writes: Iterable[tuple[str, Any]] +) -> set[str]: + """DeltaChannels an `Overwrite` among `writes` resets. + + They snapshot on the checkpoint the writes produce, as the loop does for a + node's `Overwrite`, so no read replays across the reset. + """ + return { + ch + for ch, value in writes + if isinstance(specs.get(ch), DeltaChannel) and _get_overwrite(value)[0] + } + + def checkpoint_superseded( saver: BaseCheckpointSaver, config: RunnableConfig, saved: CheckpointTuple ) -> bool: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1cf1aa4bb..e45185998 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -137,6 +137,7 @@ from langgraph.pregel._checkpoint import ( copy_checkpoint, create_checkpoint, create_checkpoint_plan_for_update_state_api, + delta_channels_overwritten, delta_channels_with_pending_writes, empty_checkpoint, get_updated_channels_from_tasks, @@ -1795,6 +1796,9 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) + fork_pending |= delta_channels_overwritten( + self.channels, input_writes + ) # apply input write to channels next_step = ( @@ -2033,6 +2037,9 @@ class Pregel( ), ) updated_channels = get_updated_channels_from_tasks(run_tasks) + fork_pending |= delta_channels_overwritten( + self.channels, (w for t in run_tasks for w in t.writes) + ) # The base's other children replay whatever is stored on it, so an # edit of an older checkpoint stores none of its writes there: the # checkpoint written here carries them, its delta channels @@ -2287,6 +2294,9 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) + fork_pending |= delta_channels_overwritten( + self.channels, input_writes + ) # apply input write to channels next_step = ( @@ -2524,6 +2534,9 @@ class Pregel( ), ) updated_channels = get_updated_channels_from_tasks(run_tasks) + fork_pending |= delta_channels_overwritten( + self.channels, (w for t in run_tasks for w in t.writes) + ) # The base's other children replay whatever is stored on it, so an # edit of an older checkpoint stores none of its writes there: the # checkpoint written here carries them, its delta channels diff --git a/libs/langgraph/tests/test_delta_channel_update_overwrite.py b/libs/langgraph/tests/test_delta_channel_update_overwrite.py new file mode 100644 index 000000000..88687326a --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_update_overwrite.py @@ -0,0 +1,110 @@ +"""An `Overwrite` through `update_state` snapshots its DeltaChannel on the +checkpoint the update saves, as a node's `Overwrite` does on the loop's.""" + +from typing import Annotated, Any + +import pytest +from langchain_core.messages import HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.channels.last_value import LastValue +from langgraph.graph import START, StateGraph +from langgraph.graph.message import _messages_delta_reducer +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.types import Overwrite + +pytestmark = pytest.mark.anyio + + +class _State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + +def _messages_graph(saver: InMemorySaver) -> Any: + builder = StateGraph(_State) + builder.add_node("model", lambda state: {}) + builder.add_edge(START, "model") + return builder.compile(checkpointer=saver) + + +def _extend(current: list, writes: list) -> list: + return [*current, *(item for write in writes for item in write)] + + +def _delta_input_graph(saver: InMemorySaver) -> Pregel: + node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log") + return Pregel( + nodes={"n": node}, + channels={"log": DeltaChannel(_extend), "go": LastValue(int)}, + input_channels=["log", "go"], + output_channels=["log"], + checkpointer=saver, + ) + + +def test_update_state_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _messages_graph(saver) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"messages": [HumanMessage(content="a", id="1")]}, config) + + graph.update_state( + config, + {"messages": Overwrite([HumanMessage(content="b", id="2")])}, + as_node="model", + ) + + head = saver.get_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot) + assert [m.content for m in graph.get_state(config).values["messages"]] == ["b"] + + +async def test_aupdate_state_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _messages_graph(saver) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"messages": [HumanMessage(content="a", id="1")]}, config) + + await graph.aupdate_state( + config, + {"messages": Overwrite([HumanMessage(content="b", id="2")])}, + as_node="model", + ) + + head = await saver.aget_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot) + values = (await graph.aget_state(config)).values + assert [m.content for m in values["messages"]] == ["b"] + + +def test_update_state_as_input_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _delta_input_graph(saver) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"log": [0], "go": 1}, config) + + graph.update_state(config, {"log": Overwrite([1])}, as_node="__input__") + + head = saver.get_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + assert graph.get_state(config).values["log"] == [1] + + +async def test_aupdate_state_as_input_with_an_overwrite_snapshots_the_channel() -> None: + saver = InMemorySaver() + graph = _delta_input_graph(saver) + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"log": [0], "go": 1}, config) + + await graph.aupdate_state(config, {"log": Overwrite([1])}, as_node="__input__") + + head = await saver.aget_tuple(config) + assert head is not None + assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + assert (await graph.aget_state(config)).values["log"] == [1]