diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index c336f75a6..91d7b5bb2 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -81,7 +81,13 @@ def get_updated_channels_from_tasks( def get_delta_channels_from_all_channels( channels: Mapping[str, BaseChannel], ) -> set[str]: - """DeltaChannels to snapshot on the first update_state of a fresh thread.""" + """Every available DeltaChannel. + + The set to snapshot whenever no ancestor walk can reconstruct these + channels: the first update_state of a fresh thread (no ancestors at all), + and the first checkpoint of a fork (whose base also holds the writes of the + branch the fork abandons). + """ return { k for k, ch in channels.items() @@ -122,14 +128,24 @@ def create_checkpoint_plan_for_update_state_api( parents: dict[str, Any], saved_metadata: Mapping[str, Any] | None, is_fresh_thread: bool, + is_fork: bool, ) -> tuple[set[str], dict[str, Any]]: - """Return ``(channels_to_snapshot, metadata)`` for an update_state head.""" + """Return ``(channels_to_snapshot, metadata)`` for an update_state head. + + ``is_fork`` (the update was addressed at an explicit checkpoint) forces a + full snapshot for the same reason ``is_fresh_thread`` does: the ancestor + walk cannot reconstruct this head. The base a fork branches off keeps the + pending writes of the branch being abandoned, and nothing records which + child consumed which write, so the walk would replay them here too. + Snapshotting terminates the walk at this checkpoint. Every delta channel + snapshots, so no counters carry over. + """ metadata: dict[str, Any] = { "source": "update", "step": step, "parents": parents, } - if is_fresh_thread: + if is_fresh_thread or is_fork: return get_delta_channels_from_all_channels(channels), metadata new_counters = create_metadata_for_update_state_api( @@ -179,9 +195,11 @@ def create_checkpoint( ch = channels[k] if k in channels_to_snapshot: # Callers force a full snapshot blob here: exit mode when a - # delta channel reaches its snapshot cadence, and update_state - # on a fresh thread (no ancestor to replay writes from). The - # manual version-bump below only applies to the exit-mode case. + # delta channel reaches its snapshot cadence, update_state on + # a fresh thread (no ancestor to replay writes from), and a + # fork (whose base also holds the abandoned branch's writes). + # The manual version-bump below only applies to the exit-mode + # case. # # In exit mode, the snapshot decision is deferred to exit # time (intermediate steps have do_checkpoint=False). The diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 371037768..8bf03f680 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -222,10 +222,32 @@ class PregelLoop: # under the saver's `ORDER BY task_id, idx` sorting. _exit_delta_writes: list[tuple[int, str, str, Any]] | None = None - # Delta channels that saw an Overwrite since the last checkpoint. These - # channels must snapshot after live update applies overwrite semantics so - # sparse replay starts from the same post-overwrite value. - _delta_channels_with_overwrite: set[str] + # Delta channels that must snapshot at the next checkpoint, whatever their + # cadence counters say. Two sources: + # * an Overwrite arrived since the last checkpoint, so the snapshot has to + # happen after live update applied overwrite semantics and sparse replay + # starts from the same post-overwrite value; + # * this run forked off an explicitly addressed checkpoint, see + # `_delta_channels_awaiting_fork_snapshot`. + _delta_channels_forced_snapshot: set[str] + + # Delta channels still owed a fork snapshot, when this run was launched + # against an explicitly addressed checkpoint (time travel / fork). That + # base keeps the pending writes of the branch the fork abandons, and + # nothing records which child consumed which write, so the ancestor walk + # would replay them into this branch too. Snapshotting terminates the walk + # inside the fork instead of at the shared base. Names drop out once the + # blob has landed; a channel with no value yet has nothing to snapshot, so + # it waits for the superstep that gives it one. + # + # The trigger is deliberately coarse: any addressed checkpoint, not only + # one that turns out to have abandoned writes on it. Which writes belong to + # which child is exactly what is not recorded, so a narrower test would + # have to trust the base's `pending_writes` to be complete, and a saver + # that leaves them out would silently go back to leaking. Snapshotting when + # it was not needed costs one blob per addressed run; not snapshotting when + # it was needed is silent corruption. + _delta_channels_awaiting_fork_snapshot: set[str] # The checkpoint_config that points at the parent loaded at `__enter__` # (or the synthetic-empty checkpoint, on first run). We capture it @@ -369,6 +391,16 @@ class PregelLoop: if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) else () ) + # Checks the value, not just key presence like `is_replaying` above: + # subgraph task configs always carry an explicit `None` here, and only + # a real id means the caller addressed one specific checkpoint. Read + # off `checkpoint_config` so subgraphs resolved through a checkpoint + # map during time travel are covered too, matching `__enter__`. + self._delta_channels_awaiting_fork_snapshot = ( + {k for k, spec in specs.items() if isinstance(spec, DeltaChannel)} + if self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) + else set() + ) self.prev_checkpoint_config = None runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME) self.control = runtime.control if isinstance(runtime, Runtime) else None @@ -683,7 +715,7 @@ class PregelLoop: def after_tick(self) -> None: # finish superstep writes = [w for t in self.tasks.values() for w in t.writes] - self._delta_channels_with_overwrite.update( + self._delta_channels_forced_snapshot.update( ch for ch, v in writes if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0] @@ -991,7 +1023,7 @@ class PregelLoop: manager=None, updated_channels=updated_channels, ) - self._delta_channels_with_overwrite.update( + self._delta_channels_forced_snapshot.update( c for c, v in input_writes if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0] @@ -1133,10 +1165,21 @@ class PregelLoop: do_checkpoint = self._checkpointer_put_after_previous is not None and ( exiting or self.durability != "exit" ) + # Fork: make this checkpoint self-contained, so the ancestor walk stops + # inside the fork instead of reaching the base this run forked off and + # collecting the abandoned branch's writes from it. Resolved here + # rather than in `_first` so channels that only got a value this + # superstep are covered too. + if self._delta_channels_awaiting_fork_snapshot: + self._delta_channels_forced_snapshot.update( + k + for k in self._delta_channels_awaiting_fork_snapshot + if (ch := self.channels.get(k)) is not None and ch.is_available() + ) # create new checkpoint channels_to_snapshot = ( delta_channels_to_snapshot(self.channels, new_counters) - | self._delta_channels_with_overwrite + | self._delta_channels_forced_snapshot if do_checkpoint else set() ) @@ -1154,7 +1197,13 @@ class PregelLoop: for k in channels_to_snapshot: new_counters[k] = (0, 0) if do_checkpoint: - self._delta_channels_with_overwrite.difference_update(channels_to_snapshot) + self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot) + # `create_checkpoint` drops a requested snapshot for a channel with + # no version in this checkpoint yet (nothing was ever written to it + # on this branch), so keep asking until the blob really landed. + self._delta_channels_awaiting_fork_snapshot.difference_update( + self.checkpoint["channel_values"] + ) non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)} if non_zero: self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero @@ -1239,7 +1288,7 @@ class PregelLoop: ) channels_to_snapshot = ( delta_channels_to_snapshot(self.channels, counters) - | self._delta_channels_with_overwrite + | self._delta_channels_forced_snapshot ) pending = [ @@ -1684,7 +1733,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_with_overwrite = set() + self._delta_channels_forced_snapshot = set() self._exit_delta_writes = ( [] if self.durability == "exit" and self.checkpointer is not None else None ) @@ -1942,7 +1991,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_with_overwrite = set() + self._delta_channels_forced_snapshot = set() self._exit_delta_writes = ( [] if self.durability == "exit" and self.checkpointer is not None else None ) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index f4ca48024..eb66e531e 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -2020,6 +2020,7 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, + is_fork=bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)), ) ) checkpoint = create_checkpoint( @@ -2480,6 +2481,7 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, + is_fork=bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)), ) ) checkpoint = create_checkpoint( diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py new file mode 100644 index 000000000..ea228284b --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -0,0 +1,248 @@ +"""Forking a thread must not replay the abandoned branch into the fork. + +Regression suite for #8443. Addressing an older checkpoint creates a fork: the +shared base ends up with two children, and it keeps the ``checkpoint_writes`` +of the branch the fork abandons. Nothing in the stored data records which child +consumed which write, so the ``DeltaChannel`` ancestor walk used to collect the +abandoned branch's writes as well. + +Every graph here carries a ``DeltaChannel`` and a plain reducer channel fed the +same values. ``full`` channels store complete ``channel_values`` and need no +replay, so the plain channel is the oracle: after a fork the two must agree. + +Coverage: fork by ``invoke`` with new input (sync/async, all durabilities), +fork off the checkpoint that predates the thread's first input, fork by +``update_state`` / ``aupdate_state``, and a guard that an unaddressed run still +follows the normal ``snapshot_frequency`` cadence. +""" + +from collections.abc import Sequence +from operator import add +from typing import Annotated, Any + +import pytest +from langchain_core.runnables import RunnableConfig +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import StateGraph +from langgraph.types import Durability, StateSnapshot + +pytestmark = pytest.mark.anyio + + +def _append(current: list | None, writes: Sequence[Any]) -> list: + """DeltaChannel reducer: extend the list with every batched write.""" + out = list(current or []) + for write in writes: + out.extend(write if isinstance(write, list) else [write]) + return out + + +class _State(TypedDict): + log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)] + plain: Annotated[list, add] + + +def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: + """Compile a one-node graph whose node appends ``{tag}-out`` to both channels. + + ``snapshot_frequency=1000`` keeps the cadence from masking the bug: without + a forced snapshot the fork's ancestor walk always runs past the fork base. + """ + + def node(state: _State) -> dict: + return {"log": [f"{tag}-out"], "plain": [f"{tag}-out"]} + + builder = StateGraph(_State) + builder.add_node("n", node) + builder.set_entry_point("n") + builder.set_finish_point("n") + return builder.compile(checkpointer=checkpointer) + + +def _thread(thread_id: str) -> RunnableConfig: + return {"configurable": {"thread_id": thread_id}} + + +def _at(config: RunnableConfig, snapshot: StateSnapshot) -> RunnableConfig: + """Config addressing one specific checkpoint of ``config``'s thread.""" + return { + "configurable": { + **config["configurable"], + "checkpoint_ns": "", + "checkpoint_id": snapshot.config["configurable"]["checkpoint_id"], + } + } + + +def _input(marker: str) -> dict: + return {"log": [marker], "plain": [marker]} + + +def _assert_fork_is_clean(state: StateSnapshot, abandoned: str) -> None: + """The delta channel must match the plain channel and drop ``abandoned``.""" + assert state.values["log"] == state.values["plain"], ( + f"delta channel diverged from the plain channel: " + f"{state.values['log']} != {state.values['plain']}" + ) + assert abandoned not in state.values["log"], ( + f"{abandoned!r} belongs to the branch the fork replaced, " + f"but was replayed into {state.values['log']}" + ) + + +def test_fork_by_invoke( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke( + _input("in-1"), config, durability=durability + ) + graph = _build(sync_checkpointer, "second") + graph.invoke(_input("in-2"), config, durability=durability) + + # The last checkpoint that predates "in-2" entering state: forking here + # abandons the "in-2" branch, whose writes still hang off this checkpoint. + base = next( + snapshot + for snapshot in graph.get_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + _build(sync_checkpointer, "third").invoke( + _input("in-3"), _at(config, base), durability=durability + ) + + state = graph.get_state(config) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "in-3", "third-out"] + + +async def test_afork_by_invoke( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + await _build(async_checkpointer, "first").ainvoke( + _input("in-1"), config, durability=durability + ) + graph = _build(async_checkpointer, "second") + await graph.ainvoke(_input("in-2"), config, durability=durability) + + base = await anext( + snapshot + async for snapshot in graph.aget_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + await _build(async_checkpointer, "third").ainvoke( + _input("in-3"), _at(config, base), durability=durability + ) + + state = await graph.aget_state(config) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "in-3", "third-out"] + + +def test_fork_off_checkpoint_before_first_input( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Fork off the root checkpoint, which predates any value for ``log``. + + ``create_checkpoint`` drops a requested snapshot for a channel absent from + ``channel_versions``, so the fork's own first checkpoint cannot carry the + blob. The request has to stay queued until a superstep gives the channel a + value, otherwise the root's ``in-1`` write still leaks into the fork. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config, durability=durability) + + root = list(graph.get_state_history(config))[-1] + assert root.values["log"] == [] + + _build(sync_checkpointer, "third").invoke( + _input("in-9"), _at(config, root), durability=durability + ) + + state = graph.get_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["in-9", "third-out"] + + +async def test_afork_off_checkpoint_before_first_input( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Async twin of ``test_fork_off_checkpoint_before_first_input``.""" + config = _thread("t") + graph = _build(async_checkpointer, "first") + await graph.ainvoke(_input("in-1"), config, durability=durability) + + root = [snapshot async for snapshot in graph.aget_state_history(config)][-1] + assert root.values["log"] == [] + + await _build(async_checkpointer, "third").ainvoke( + _input("in-9"), _at(config, root), durability=durability + ) + + state = await graph.aget_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["in-9", "third-out"] + + +def test_fork_by_update_state(sync_checkpointer: BaseCheckpointSaver) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke(_input("in-1"), config) + graph = _build(sync_checkpointer, "second") + graph.invoke(_input("in-2"), config) + + base = next( + snapshot + for snapshot in graph.get_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + forked = graph.update_state(_at(config, base), _input("patched")) + + state = graph.get_state(forked) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "patched"] + + +async def test_afork_by_update_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + await _build(async_checkpointer, "first").ainvoke(_input("in-1"), config) + graph = _build(async_checkpointer, "second") + await graph.ainvoke(_input("in-2"), config) + + base = await anext( + snapshot + async for snapshot in graph.aget_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + forked = await graph.aupdate_state(_at(config, base), _input("patched")) + + state = await graph.aget_state(forked) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "patched"] + + +def test_unaddressed_run_keeps_snapshot_cadence( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """A run with no explicitly addressed checkpoint writes no snapshot blob. + + Guards the cost of the fix: the forced snapshot is one per addressed run, + not a change to the normal ``snapshot_frequency`` cadence. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config, durability=durability) + graph.invoke(_input("in-2"), config, durability=durability) + + assert not [ + tuple_.config["configurable"]["checkpoint_id"] + for tuple_ in sync_checkpointer.list(config) + if isinstance(tuple_.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + ]