fix(langgraph): don't replay an abandoned branch into a DeltaChannel fork

Addressing an older checkpoint creates a fork: the shared base ends up
with two children and keeps the checkpoint_writes of the branch the fork
abandons. Nothing records which child consumed which write, so the
DeltaChannel ancestor walk collected the abandoned branch's writes too.
Live execution was correct; only the reconstruction after a reload was
wrong, and it was wrong on every saver.

Fixed on the write side, so no saver changes are needed. A run launched
against an explicitly addressed checkpoint forces every DeltaChannel to
snapshot into its first checkpoint, terminating the walk inside the fork
instead of at the shared base. This mirrors the existing force-snapshot
for Overwrite writes, hence the rename to _delta_channels_forced_snapshot.
update_state against an older checkpoint takes the same path, for the
same reason is_fresh_thread already does.

A channel with no value at the fork base cannot carry a snapshot blob
yet, so the request stays queued until the first superstep that gives it
one. Cost is one snapshot per addressed run, not per superstep.

Fixes #8443

Co-Authored-By: AnnaSuSu <64579968+AnnaSuSu@users.noreply.github.com>
Co-Authored-By: UditDewan <194863456+UditDewan@users.noreply.github.com>

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

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
         )

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(

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