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
@@ -81,7 +81,13 @@
 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 @@
     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 @@
             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 (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 @@ 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
         )

libs/langgraph/langgraph/pregel/main.py

@@ -1638,7 +1638,10 @@ def bulk_update_state(
                 raise ValueError(f"Subgraph {recast} not found")
 
         def perform_superstep(
-            input_config: RunnableConfig, updates: Sequence[StateUpdate]
+            input_config: RunnableConfig,
+            updates: Sequence[StateUpdate],
+            *,
+            is_fork: bool,
         ) -> RunnableConfig:
             # get last checkpoint
             config = ensure_config(self.config, input_config)
@@ -1873,6 +1876,9 @@ def perform_superstep(
                     return perform_superstep(
                         patch_checkpoint_map(next_config, saved.metadata),
                         [item for lst in user_group_by.values() for item in lst],
+                        # The checkpoint just written clears tasks and carries
+                        # no delta snapshot, so a fork is still unsealed here.
+                        is_fork=is_fork,
                     )
 
                 return patch_checkpoint_map(next_config, saved.metadata)
@@ -2020,6 +2026,7 @@ def perform_superstep(
                     parents=saved.metadata.get("parents", {}) if saved else {},
                     saved_metadata=saved.metadata if saved else None,
                     is_fresh_thread=saved is None,
+                    is_fork=is_fork,
                 )
             )
             checkpoint = create_checkpoint(
@@ -2049,8 +2056,16 @@ def perform_superstep(
         current_config = patch_configurable(
             config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
         )
+        # Only the first superstep can fork. `perform_superstep` returns the
+        # config of the checkpoint it just wrote and the loop feeds that back
+        # in, so from the second superstep on the incoming config always names
+        # a checkpoint whether or not the caller addressed one.
+        is_fork = bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID))
         for superstep in supersteps:
-            current_config = perform_superstep(current_config, superstep)
+            current_config = perform_superstep(
+                current_config, superstep, is_fork=is_fork
+            )
+            is_fork = False
         return current_config
 
     async def abulk_update_state(
@@ -2104,7 +2119,10 @@ async def abulk_update_state(
                 raise ValueError(f"Subgraph {recast} not found")
 
         async def aperform_superstep(
-            input_config: RunnableConfig, updates: Sequence[StateUpdate]
+            input_config: RunnableConfig,
+            updates: Sequence[StateUpdate],
+            *,
+            is_fork: bool,
         ) -> RunnableConfig:
             # get last checkpoint
             config = ensure_config(self.config, input_config)
@@ -2335,6 +2353,9 @@ async def aperform_superstep(
                     return await aperform_superstep(
                         patch_checkpoint_map(next_config, saved.metadata),
                         [item for lst in user_group_by.values() for item in lst],
+                        # The checkpoint just written clears tasks and carries
+                        # no delta snapshot, so a fork is still unsealed here.
+                        is_fork=is_fork,
                     )
 
                 return patch_checkpoint_map(
@@ -2480,6 +2501,7 @@ async def aperform_superstep(
                     parents=saved.metadata.get("parents", {}) if saved else {},
                     saved_metadata=saved.metadata if saved else None,
                     is_fresh_thread=saved is None,
+                    is_fork=is_fork,
                 )
             )
             checkpoint = create_checkpoint(
@@ -2508,8 +2530,16 @@ async def aperform_superstep(
         current_config = patch_configurable(
             config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
         )
+        # Only the first superstep can fork. `aperform_superstep` returns the
+        # config of the checkpoint it just wrote and the loop feeds that back
+        # in, so from the second superstep on the incoming config always names
+        # a checkpoint whether or not the caller addressed one.
+        is_fork = bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID))
         for superstep in supersteps:
-            current_config = await aperform_superstep(current_config, superstep)
+            current_config = await aperform_superstep(
+                current_config, superstep, is_fork=is_fork
+            )
+            is_fork = False
         return current_config
 
     def update_state(