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

libs/langgraph/langgraph/pregel/_checkpoint.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
 import uuid
@@ -80,12 +80,20 @@
 
 def get_delta_channels_from_all_channels(
     channels: Mapping[str, BaseChannel],
+    *,
+    include_unavailable: bool = False,
 ) -> 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()
-        if isinstance(ch, DeltaChannel) and ch.is_available()
+        if isinstance(ch, DeltaChannel) and (include_unavailable or ch.is_available())
     }
 
 
@@ -122,15 +130,27 @@
     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:
-        return get_delta_channels_from_all_channels(channels), metadata
+    if is_fresh_thread or is_fork:
+        return get_delta_channels_from_all_channels(
+            channels, include_unavailable=is_fork
+        ), metadata
 
     new_counters = create_metadata_for_update_state_api(
         channels,
@@ -146,6 +166,43 @@
     return channels_to_snapshot, metadata
 
 
+def create_fork_checkpoint(
+    checkpoint: Checkpoint,
+    channels: Mapping[str, BaseChannel],
+    step: int,
+    *,
+    is_fork: bool,
+    get_next_version: GetNextVersion,
+) -> Checkpoint:
+    """``create_checkpoint`` for an update_state path that bypasses the plan.
+
+    The ``as_node`` INPUT and END paths write the fork's first checkpoint and
+    return before ``create_checkpoint_plan_for_update_state_api`` runs. Left
+    without a snapshot that checkpoint does not seal the fork, and the next
+    superstep reconstructs its delta channels by walking through the shared
+    base, picking up the abandoned branch's writes and then baking them into
+    whatever it snapshots. Sealing has to happen on the fork's *first*
+    checkpoint, which is this one.
+
+    ``get_next_version`` is required for the same reason exit mode needs it:
+    these paths apply writes to the input channel, not to the delta channel,
+    so nothing bumps the delta channel's version and ``put`` would drop the
+    blob as not-a-new-version. Callers must derive ``new_versions`` from the
+    returned checkpoint rather than the one they passed in.
+    """
+    if not is_fork:
+        return create_checkpoint(checkpoint, channels, step)
+    return create_checkpoint(
+        checkpoint,
+        channels,
+        step,
+        get_next_version=get_next_version,
+        channels_to_snapshot=get_delta_channels_from_all_channels(
+            channels, include_unavailable=True
+        ),
+    )
+
+
 def create_checkpoint(
     checkpoint: Checkpoint,
     channels: Mapping[str, BaseChannel] | None,
@@ -174,14 +231,27 @@
         values = {}
         channel_versions = dict(checkpoint["channel_versions"])
         for k in channels:
+            ch = channels[k]
             if k not in channel_versions:
+                # Nothing was ever written to this channel on this branch, so
+                # it has no version and `put` would drop any blob stored for
+                # it. A *forced* snapshot still has to land: it is the only
+                # thing that stops the ancestor walk running past this
+                # checkpoint into a fork base that holds another branch's
+                # writes. Mint a first version so the blob survives.
+                if k in channels_to_snapshot and get_next_version is not None:
+                    channel_versions[k] = get_next_version(None, None)
+                    values[k] = _DeltaSnapshot(
+                        ch.get() if ch.is_available() else ch.typ()
+                    )
                 continue
-            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 @@ def __init__(
             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 @@ def tick(self) -> bool:
     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 @@ def _first(
                 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 @@ def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
         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 k in self.channels
+            )
         # 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 @@ def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
         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 @@ def _put_exit_delta_writes(self) -> None:
         )
         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 @@ def __enter__(self) -> Self:
         )
         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 @@ async def __aenter__(self) -> Self:
         )
         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
         )