From 834e53df192af45ff251d2ff7404876e6e77a665 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Wed, 5 Aug 2026 23:42:35 -0400 Subject: [PATCH] fix(langgraph): seal a fork on the first checkpoint it writes The as_node INPUT, END and __copy__ paths write a checkpoint and return before create_checkpoint_plan_for_update_state_api runs, so a bulk update whose first superstep took one of them left the branch unsealed. Only INPUT actually leaked: END absorbs the base's already-run task writes, so its delta and plain channels agree. Sealing on a later superstep does not help. By then that superstep has reconstructed its value by walking through the unsealed checkpoint into the shared base, so it snapshots an already-corrupted list. The fork's first checkpoint is the one that has to carry the blob, which is what create_fork_checkpoint does. That snapshot was still being dropped by put: these paths apply writes to the input channel, not the delta channel, so nothing bumped the delta channel's version and it never entered new_versions. Pass get_next_version for the manual bump, the same reason exit mode needs it, and derive new_versions from the returned checkpoint. fork_pending tracks what is still owed, mirroring _delta_channels_awaiting_fork_snapshot in _loop.py. Caught by the Open SWE review bot on #8548. --- .../langgraph/langgraph/pregel/_checkpoint.py | 35 ++++++ libs/langgraph/langgraph/pregel/main.py | 105 ++++++++++++++---- .../tests/test_delta_channel_fork.py | 45 +++++++- 3 files changed, 162 insertions(+), 23 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 91d7b5bb2..ba0c6178a 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -162,6 +162,41 @@ def create_checkpoint_plan_for_update_state_api( 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), + ) + + def create_checkpoint( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel] | None, diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 6705c9c17..c32636d39 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -108,6 +108,7 @@ from langgraph.callbacks import ( get_sync_graph_callback_manager_for_config, ) from langgraph.channels.base import BaseChannel +from langgraph.channels.delta import DeltaChannel from langgraph.channels.topic import Topic from langgraph.config import get_config from langgraph.constants import END @@ -133,6 +134,7 @@ from langgraph.pregel._checkpoint import ( copy_checkpoint, create_checkpoint, create_checkpoint_plan_for_update_state_api, + create_fork_checkpoint, empty_checkpoint, get_updated_channels_from_tasks, ) @@ -1637,6 +1639,19 @@ class Pregel( else: raise ValueError(f"Subgraph {recast} not found") + # Delta channels still owed a fork snapshot. Mirrors + # `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only + # sealed once a checkpoint actually carries the blob. Several superstep + # paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and + # return before reaching the plan, so a flag cleared after the first + # superstep would leave the branch unsealed and the next superstep would + # reconstruct through the shared base. + fork_pending: set[str] = ( + {k for k, v in self.channels.items() if isinstance(v, DeltaChannel)} + if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) + else set() + ) + def perform_superstep( input_config: RunnableConfig, updates: Sequence[StateUpdate], @@ -1729,9 +1744,17 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + next_checkpoint = create_fork_checkpoint( + checkpoint, + channels, + step, + is_fork=is_fork, + get_next_version=checkpointer.get_next_version, + ) + fork_pending.difference_update(next_checkpoint["channel_values"]) next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + next_checkpoint, { "source": "update", "step": step + 1, @@ -1739,7 +1762,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) return patch_checkpoint_map( @@ -1768,9 +1791,17 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + next_checkpoint = create_fork_checkpoint( + checkpoint, + channels, + next_step, + is_fork=is_fork, + get_next_version=checkpointer.get_next_version, + ) + fork_pending.difference_update(next_checkpoint["channel_values"]) next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + next_checkpoint, { "source": "input", "step": next_step, @@ -1780,7 +1811,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) @@ -2039,6 +2070,8 @@ class Pregel( else None, channels_to_snapshot=channels_to_snapshot, ) + if is_fork: + fork_pending.difference_update(checkpoint["channel_values"]) next_config = checkpointer.put( checkpoint_config, checkpoint, @@ -2056,16 +2089,14 @@ class Pregel( 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)) + # The flag cannot be derived from `current_config`: `perform_superstep` + # returns the config of the checkpoint it just wrote and the loop feeds + # that back in, so from the second superstep on it always names a + # checkpoint whether or not the caller addressed one. for superstep in supersteps: current_config = perform_superstep( - current_config, superstep, is_fork=is_fork + current_config, superstep, is_fork=bool(fork_pending) ) - is_fork = False return current_config async def abulk_update_state( @@ -2118,6 +2149,19 @@ class Pregel( else: raise ValueError(f"Subgraph {recast} not found") + # Delta channels still owed a fork snapshot. Mirrors + # `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only + # sealed once a checkpoint actually carries the blob. Several superstep + # paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and + # return before reaching the plan, so a flag cleared after the first + # superstep would leave the branch unsealed and the next superstep would + # reconstruct through the shared base. + fork_pending: set[str] = ( + {k for k, v in self.channels.items() if isinstance(v, DeltaChannel)} + if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) + else set() + ) + async def aperform_superstep( input_config: RunnableConfig, updates: Sequence[StateUpdate], @@ -2208,16 +2252,25 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + next_checkpoint = create_fork_checkpoint( + checkpoint, + channels, + step, + is_fork=is_fork, + get_next_version=checkpointer.get_next_version, + ) + fork_pending.difference_update(next_checkpoint["channel_values"]) next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + next_checkpoint, { "source": "update", "step": step + 1, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( - checkpoint_previous_versions, checkpoint["channel_versions"] + checkpoint_previous_versions, + next_checkpoint["channel_versions"], ), ) return patch_checkpoint_map( @@ -2246,9 +2299,17 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + next_checkpoint = create_fork_checkpoint( + checkpoint, + channels, + next_step, + is_fork=is_fork, + get_next_version=checkpointer.get_next_version, + ) + fork_pending.difference_update(next_checkpoint["channel_values"]) next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + next_checkpoint, { "source": "input", "step": next_step, @@ -2258,7 +2319,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) @@ -2514,6 +2575,8 @@ class Pregel( else None, channels_to_snapshot=channels_to_snapshot, ) + if is_fork: + fork_pending.difference_update(checkpoint["channel_values"]) next_config = await checkpointer.aput( checkpoint_config, checkpoint, @@ -2530,16 +2593,14 @@ class Pregel( 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)) + # The flag cannot be derived from `current_config`: `aperform_superstep` + # returns the config of the checkpoint it just wrote and the loop feeds + # that back in, so from the second superstep on it always names a + # checkpoint whether or not the caller addressed one. for superstep in supersteps: current_config = await aperform_superstep( - current_config, superstep, is_fork=is_fork + current_config, superstep, is_fork=bool(fork_pending) ) - is_fork = False return current_config def update_state( diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index c7d56f5d5..bb991e3f8 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -27,8 +27,9 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.serde.types import _DeltaSnapshot from typing_extensions import TypedDict +from langgraph._internal._constants import INPUT from langgraph.channels.delta import DeltaChannel -from langgraph.graph import StateGraph +from langgraph.graph import END, StateGraph from langgraph.types import Durability, StateSnapshot, StateUpdate pytestmark = pytest.mark.anyio @@ -256,6 +257,48 @@ def test_unaddressed_run_keeps_snapshot_cadence( assert not _snapshotted_checkpoints(sync_checkpointer, config) +@pytest.mark.parametrize("first_as_node", [INPUT, END, "__copy__"]) +def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan( + sync_checkpointer: BaseCheckpointSaver, first_as_node: str +) -> None: + """The fork must be sealed by whichever checkpoint the fork writes first. + + ``as_node`` of INPUT, END or ``__copy__`` writes a checkpoint and returns + before ``create_checkpoint_plan_for_update_state_api`` runs. If that + checkpoint carries no snapshot the branch is still unsealed, so the next + superstep reconstructs through the shared base, picks up the abandoned + writes, and bakes them into whatever it snapshots. Sealing later is too + late: by then the in-memory value is already wrong. + """ + 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"] + ) + first = ( + StateUpdate(_input("first-step"), first_as_node) + if first_as_node == INPUT + else StateUpdate(None, first_as_node) + ) + forked = graph.bulk_update_state( + _at(config, base), + [[first], [StateUpdate(_input("second-step"), "n")]], + ) + + state = graph.get_state(forked) + # END legitimately absorbs the base's already-run task writes, so "in-2" + # belongs there; the plain channel is the oracle for which is which. + assert state.values["log"] == state.values["plain"], ( + f"delta channel diverged from the plain channel: " + f"{state.values['log']} != {state.values['plain']}" + ) + + def test_unaddressed_bulk_update_keeps_snapshot_cadence( sync_checkpointer: BaseCheckpointSaver, ) -> None: