From 9d92f33cca143926c563deb87ce377397883cd92 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Fri, 9 Oct 2026 14:36:27 -0400 Subject: [PATCH] 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. --- .../langgraph/langgraph/pregel/_checkpoint.py | 23 ++++-- libs/langgraph/langgraph/pregel/_loop.py | 1 + libs/langgraph/langgraph/pregel/main.py | 6 ++ .../test_delta_channel_seal_subscribers.py | 78 +++++++++++++++++++ 4 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 libs/langgraph/tests/test_delta_channel_seal_subscribers.py diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index b2cddc19d..7225483da 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/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(): diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2c46108c3..0d20b5bde 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/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) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1cf1aa4bb..d45c2ddff 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/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, diff --git a/libs/langgraph/tests/test_delta_channel_seal_subscribers.py b/libs/langgraph/tests/test_delta_channel_seal_subscribers.py new file mode 100644 index 000000000..8a46e6be8 --- /dev/null +++ b/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]]