fix(langgraph): keep a fork's DeltaChannel snapshot from starting nodes that never ran (#9264)

A fork from a checkpoint with pending writes to a DeltaChannel snapshots
the channel on the fork's first checkpoint, so the fork doesn't replay
the other branch's writes. If the channel had never been written there,
the snapshot also gives it its first version, and `_mark_bumps_seen`
only marked the bump seen for nodes that already have a `versions_seen`
entry. A node subscribed to the channel that never ran has none, so the
bump started it, and it read the empty value before the node that writes
the channel had run:

```python
graph.invoke("go", config)  # writer writes [1] to d, reader reads [1]
fork = graph.update_state(first_checkpoint, {"a": "go"}, as_node="__input__")
graph.get_state(fork).next  # ("writer", "reader"): resuming runs reader on [] first
```

A replay from that checkpoint saves the same kind of fork from the loop.
A channel bumped from no version was never written, so
`create_checkpoint` now also marks the bump seen for the nodes the
channel triggers, from the graph's `trigger_to_nodes`, at every call
that can bump. Only raw `Pregel` graphs hit this: `StateGraph` nodes
trigger on their `branch:to:*` channels, not on state keys. An
`as_node="__copy__"` fork still starts such a node, with a plain list
channel as well, so that case isn't DeltaChannel-specific and isn't
changed here.

Raised in review of langchain-ai/langgraphjs#2974. JS twin:
langchain-ai/langgraphjs#2983

## Tests

`test_delta_channel_seal_subscribers.py`: an update as input (sync and
async) and a replay from the checkpoint before a DeltaChannel's first
write leave only the writer next, and the reader then reads the written
value once. All three fail on main. The `langgraph` suite passes
(without Redis), and `lint_package` and `lint_tests` are clean.

libs/langgraph/langgraph/pregel/_checkpoint.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import uuid
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from datetime import datetime, timezone
 from inspect import signature
 from typing import Any, Literal, cast
@@ -256,6 +256,7 @@ def create_checkpoint(
     get_next_version: GetNextVersion | None = None,
     channels_to_snapshot: set[str] | None = None,
     stored_versions: ChannelVersions | None = None,
+    trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
 ) -> Checkpoint:
     """Build a new Checkpoint from the previous one and live channel state.
 
@@ -314,7 +315,9 @@ def create_checkpoint(
         id=id or str(uuid6(clock_seq=step)),
         channel_values=values,
         channel_versions=channel_versions,
-        versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped),
+        versions_seen=_mark_bumps_seen(
+            checkpoint["versions_seen"], bumped, trigger_to_nodes or {}
+        ),
         updated_channels=None if updated_channels is None else sorted(updated_channels),
     )
 
@@ -322,19 +325,27 @@ def create_checkpoint(
 def _mark_bumps_seen(
     versions_seen: dict[str, ChannelVersions],
     bumped: Mapping[str, tuple[Any, Any]],
+    trigger_to_nodes: Mapping[str, Sequence[str]],
 ) -> dict[str, ChannelVersions]:
     """Advance whoever had seen a bumped channel's old version to the new one.
 
     A bump that only stores a snapshot is not a write. Left unseen, it would
-    re-fire `interrupt_before` and rerun the channel's subscribers. For each
-    entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the
-    node really read, so `versions_seen_without_bumps` can put the read back.
+    re-fire `interrupt_before` and rerun the channel's subscribers. A channel
+    bumped from no version was never written, so it also goes to the
+    subscribers that never ran: they have no entry, and would start on the bump.
+    For each entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the
+    one the node really read, so `versions_seen_without_bumps` can put the read
+    back.
     """
     if not bumped:
         return versions_seen
     out = dict(versions_seen)
+    for k, (old, _) in bumped.items():
+        if old is None:
+            for node in trigger_to_nodes.get(k, ()):
+                out.setdefault(node, {})
     marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {}))
-    for node, seen in versions_seen.items():
+    for node, seen in out.items():
         if node == SNAPSHOT_BUMPS:
             continue
         for k, (old, new) in bumped.items():

libs/langgraph/langgraph/pregel/_loop.py

@@ -1256,6 +1256,7 @@ class PregelLoop:
             else None,
             channels_to_snapshot=channels_to_snapshot,
             stored_versions=self.checkpoint_previous_versions,
+            trigger_to_nodes=self.trigger_to_nodes,
         )
         for k in channels_to_snapshot:
             new_counters[k] = (0, 0)

libs/langgraph/langgraph/pregel/main.py

@@ -1766,6 +1766,7 @@ class Pregel(
                     get_next_version=checkpointer.get_next_version,
                     channels_to_snapshot=channels_to_snapshot,
                     stored_versions=checkpoint_previous_versions,
+                    trigger_to_nodes=self.trigger_to_nodes,
                 )
                 next_config = checkpointer.put(
                     checkpoint_config,
@@ -1822,6 +1823,7 @@ class Pregel(
                         get_next_version=checkpointer.get_next_version,
                         channels_to_snapshot=channels_to_snapshot,
                         stored_versions=checkpoint_previous_versions,
+                        trigger_to_nodes=self.trigger_to_nodes,
                     )
                     next_config = checkpointer.put(
                         checkpoint_config,
@@ -2086,6 +2088,7 @@ class Pregel(
                 else None,
                 channels_to_snapshot=channels_to_snapshot,
                 stored_versions=checkpoint_previous_versions,
+                trigger_to_nodes=self.trigger_to_nodes,
             )
             next_config = checkpointer.put(
                 checkpoint_config,
@@ -2258,6 +2261,7 @@ class Pregel(
                     get_next_version=checkpointer.get_next_version,
                     channels_to_snapshot=channels_to_snapshot,
                     stored_versions=checkpoint_previous_versions,
+                    trigger_to_nodes=self.trigger_to_nodes,
                 )
                 next_config = await checkpointer.aput(
                     checkpoint_config,
@@ -2314,6 +2318,7 @@ class Pregel(
                         get_next_version=checkpointer.get_next_version,
                         channels_to_snapshot=channels_to_snapshot,
                         stored_versions=checkpoint_previous_versions,
+                        trigger_to_nodes=self.trigger_to_nodes,
                     )
                     next_config = await checkpointer.aput(
                         checkpoint_config,
@@ -2577,6 +2582,7 @@ class Pregel(
                 else None,
                 channels_to_snapshot=channels_to_snapshot,
                 stored_versions=checkpoint_previous_versions,
+                trigger_to_nodes=self.trigger_to_nodes,
             )
             next_config = await checkpointer.aput(
                 checkpoint_config,

libs/langgraph/tests/test_delta_channel_seal_subscribers.py

@@ -0,0 +1,78 @@
+import pytest
+from langgraph.checkpoint.memory import InMemorySaver
+
+from langgraph.channels.delta import DeltaChannel
+from langgraph.channels.last_value import LastValue
+from langgraph.pregel import NodeBuilder, Pregel
+
+pytestmark = pytest.mark.anyio
+
+CONFIG = {"configurable": {"thread_id": "t"}}
+
+
+def _extend(current: list, writes: list) -> list:
+    return [*current, *(item for write in writes for item in write)]
+
+
+def _graph(reads: list) -> Pregel:
+    writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
+    reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
+    return Pregel(
+        nodes={"writer": writer, "reader": reader},
+        channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
+        input_channels="a",
+        output_channels=["d"],
+        checkpointer=InMemorySaver(),
+    )
+
+
+def test_an_input_update_from_before_the_first_write_starts_only_the_writer() -> None:
+    reads: list = []
+    graph = _graph(reads)
+    graph.invoke("go", CONFIG)
+    first = next(
+        s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
+    )
+
+    fork = graph.update_state(first, {"a": "go"}, as_node="__input__")
+
+    assert graph.get_state(fork).next == ("writer",)
+    graph.invoke(None, fork)
+    assert reads == [[1], [1]]
+
+
+def test_a_replay_from_before_the_first_write_forks_with_only_the_writer_next() -> None:
+    reads: list = []
+    graph = _graph(reads)
+    graph.invoke("go", CONFIG)
+    first = next(
+        s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
+    )
+
+    graph.invoke(None, first, durability="sync")
+
+    fork = next(
+        s for s in graph.get_state_history(CONFIG) if s.metadata["source"] == "fork"
+    )
+    assert fork.next == ("writer",)
+    graph.invoke(None, fork.config)
+    assert reads == [[1], [1], [1]]
+
+
+async def test_an_ainput_update_from_before_the_first_write_starts_only_the_writer() -> (
+    None
+):
+    reads: list = []
+    graph = _graph(reads)
+    await graph.ainvoke("go", CONFIG)
+    first = [
+        s.config
+        async for s in graph.aget_state_history(CONFIG)
+        if s.metadata["step"] == -1
+    ][0]
+
+    fork = await graph.aupdate_state(first, {"a": "go"}, as_node="__input__")
+
+    assert (await graph.aget_state(fork)).next == ("writer",)
+    await graph.ainvoke(None, fork)
+    assert reads == [[1], [1]]