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]]