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.

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:

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

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]