fix(langgraph): store an update_state input's writes on the checkpoint it updates

libs/langgraph/langgraph/pregel/main.py

@@ -1788,6 +1788,25 @@ class Pregel(
                     )
 
                 if input_writes := deque(map_input(self.input_channels, values)):
+                    # A DeltaChannel reads the writes stored on a checkpoint's
+                    # ancestors, not its own, so they go on the checkpoint this
+                    # update builds on, as a node's writes do.
+                    if (
+                        is_first
+                        and saved is not None
+                        and checkpoint_superseded(checkpointer, config, saved)
+                    ):
+                        fork_pending.update(
+                            ch
+                            for ch, _ in input_writes
+                            if isinstance(self.channels.get(ch), DeltaChannel)
+                        )
+                    elif saved is not None:
+                        checkpointer.put_writes(
+                            checkpoint_config,
+                            input_writes,
+                            str(uuid5(UUID(checkpoint["id"]), INPUT)),
+                        )
                     updated_channels = apply_writes(
                         checkpoint,
                         channels,
@@ -1833,13 +1852,6 @@ class Pregel(
                         ),
                     )
 
-                    # store the writes
-                    checkpointer.put_writes(
-                        next_config,
-                        input_writes,
-                        str(uuid5(UUID(checkpoint["id"]), INPUT)),
-                    )
-
                     return patch_checkpoint_map(
                         next_config, saved.metadata if saved else None
                     )
@@ -2280,6 +2292,25 @@ class Pregel(
                     )
 
                 if input_writes := deque(map_input(self.input_channels, values)):
+                    # A DeltaChannel reads the writes stored on a checkpoint's
+                    # ancestors, not its own, so they go on the checkpoint this
+                    # update builds on, as a node's writes do.
+                    if (
+                        is_first
+                        and saved is not None
+                        and await acheckpoint_superseded(checkpointer, config, saved)
+                    ):
+                        fork_pending.update(
+                            ch
+                            for ch, _ in input_writes
+                            if isinstance(self.channels.get(ch), DeltaChannel)
+                        )
+                    elif saved is not None:
+                        await checkpointer.aput_writes(
+                            checkpoint_config,
+                            input_writes,
+                            str(uuid5(UUID(checkpoint["id"]), INPUT)),
+                        )
                     updated_channels = apply_writes(
                         checkpoint,
                         channels,
@@ -2325,13 +2356,6 @@ class Pregel(
                         ),
                     )
 
-                    # store the writes
-                    await checkpointer.aput_writes(
-                        next_config,
-                        input_writes,
-                        str(uuid5(UUID(checkpoint["id"]), INPUT)),
-                    )
-
                     return patch_checkpoint_map(
                         next_config, saved.metadata if saved else None
                     )

libs/langgraph/tests/test_delta_channel_update_state.py

@@ -25,9 +25,12 @@ 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.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 StateSnapshot, StateUpdate
 
 pytestmark = pytest.mark.anyio
@@ -538,3 +541,94 @@ def test_update_state_that_snapshots_keeps_a_deferred_node_pending() -> None:
 
     assert [m.content for m in final["messages"]] == ["s", "a", "u", "b"]
     assert graph.get_state(config).next == ()
+
+
+def _sorted_extend(current: list, writes: list) -> list:
+    return sorted([*current, *(item for write in writes for item in write)])
+
+
+def _delta_input_graph(snapshot_frequency: int = 1000) -> Any:
+    node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
+    return Pregel(
+        nodes={"n": node},
+        channels={
+            "log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency),
+            "plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
+            "go": LastValue(int),
+        },
+        input_channels=["log", "plain", "go"],
+        output_channels=["log", "plain"],
+        checkpointer=InMemorySaver(),
+    )
+
+
+@pytest.mark.parametrize("snapshot_frequency", [1, 2])
+def test_update_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
+    snapshot_frequency: int,
+) -> None:
+    graph = _delta_input_graph(snapshot_frequency)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"log": [0], "plain": [0], "go": 1}, config)
+
+    graph.update_state(config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__")
+    after_update = graph.get_state(config).values
+    graph.invoke(None, config)
+    after_run = graph.get_state(config).values
+
+    assert after_update["log"] == after_update["plain"]
+    assert after_run["log"] == after_run["plain"]
+
+
+@pytest.mark.parametrize("snapshot_frequency", [1, 2])
+async def test_aupdate_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
+    snapshot_frequency: int,
+) -> None:
+    graph = _delta_input_graph(snapshot_frequency)
+    config = {"configurable": {"thread_id": "t"}}
+    await graph.ainvoke({"log": [0], "plain": [0], "go": 1}, config)
+
+    await graph.aupdate_state(
+        config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
+    )
+    after_update = (await graph.aget_state(config)).values
+    await graph.ainvoke(None, config)
+    after_run = (await graph.aget_state(config)).values
+
+    assert after_update["log"] == after_update["plain"]
+    assert after_run["log"] == after_run["plain"]
+
+
+def test_update_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> None:
+    graph = _delta_input_graph()
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"go": 1}, config)
+    older = graph.get_state(config).config
+    graph.invoke({"go": 1}, config)
+    other_branch = graph.get_state(config)
+
+    edited = graph.update_state(
+        older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
+    )
+
+    values = graph.get_state(edited).values
+    assert values["log"] == values["plain"]
+    assert graph.get_state(other_branch.config).values == other_branch.values
+
+
+async def test_aupdate_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> (
+    None
+):
+    graph = _delta_input_graph()
+    config = {"configurable": {"thread_id": "t"}}
+    await graph.ainvoke({"go": 1}, config)
+    older = (await graph.aget_state(config)).config
+    await graph.ainvoke({"go": 1}, config)
+    other_branch = await graph.aget_state(config)
+
+    edited = await graph.aupdate_state(
+        older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
+    )
+
+    values = (await graph.aget_state(edited)).values
+    assert values["log"] == values["plain"]
+    assert (await graph.aget_state(other_branch.config)).values == other_branch.values