From 4b0e98113b40d6271c4d917060a7cc7a12db0884 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Thu, 6 Aug 2026 11:37:10 -0400 Subject: [PATCH] fix(langgraph): seal a fork whose delta channel has no value yet A DeltaChannel that was never written on the branch being forked has no value to snapshot and no entry in channel_versions, so create_checkpoint skipped it and the fork's first checkpoint recorded no boundary at all. The walk then ran past the fork into the shared base and collected the abandoned branch's writes, the same failure this branch already fixes for channels that do have a value. Two shapes leaked. A run forking off a checkpoint older than the channel's first value and never writing that channel returned ['in-1'] where the plain-channel oracle returned []. A bulk update writing the delta key only in its second superstep returned ['in-1', 's2'] against ['s2']. No new blob type is needed. _DeltaSnapshot already carries the value and is already serialized by every saver, and from_checkpoint turns MISSING into typ(), so _DeltaSnapshot(typ()) reconstructs to the same empty value the channel would have had. What was missing is a version: without one, put drops the blob as not-a-new-version, so mint a first one. Deferring the seal to a later superstep does not work. That superstep reconstructs through the still-unsealed checkpoint and would only bake the corrupted value into its own snapshot. Checked that minting a version does not fire nodes that subscribe to the channel: a raw Pregel node subscribed directly to the delta channel stays silent across the fork. Reported by the Open SWE review bot on #8548. --- .../langgraph/langgraph/pregel/_checkpoint.py | 27 +++++- libs/langgraph/langgraph/pregel/_loop.py | 2 +- .../tests/test_delta_channel_fork.py | 94 ++++++++++++++++++- 3 files changed, 116 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index ba0c6178a..127700728 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -80,6 +80,8 @@ def get_updated_channels_from_tasks( def get_delta_channels_from_all_channels( channels: Mapping[str, BaseChannel], + *, + include_unavailable: bool = False, ) -> set[str]: """Every available DeltaChannel. @@ -91,7 +93,7 @@ def get_delta_channels_from_all_channels( 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()) } @@ -146,7 +148,9 @@ def create_checkpoint_plan_for_update_state_api( "parents": parents, } if is_fresh_thread or is_fork: - return get_delta_channels_from_all_channels(channels), metadata + return get_delta_channels_from_all_channels( + channels, include_unavailable=is_fork + ), metadata new_counters = create_metadata_for_update_state_api( channels, @@ -193,7 +197,9 @@ def create_fork_checkpoint( channels, step, get_next_version=get_next_version, - channels_to_snapshot=get_delta_channels_from_all_channels(channels), + channels_to_snapshot=get_delta_channels_from_all_channels( + channels, include_unavailable=True + ), ) @@ -225,9 +231,20 @@ def create_checkpoint( values = {} channel_versions = dict(checkpoint["channel_versions"]) for k in channels: - if k not in channel_versions: - continue 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 if k in channels_to_snapshot: # Callers force a full snapshot blob here: exit mode when a # delta channel reaches its snapshot cadence, update_state on diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 8bf03f680..21989fdf3 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -1174,7 +1174,7 @@ class PregelLoop: 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() + if k in self.channels ) # create new checkpoint channels_to_snapshot = ( diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index bb991e3f8..d53a86689 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -12,7 +12,8 @@ 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 (sync/async), -fork by ``update_state`` / ``aupdate_state``, and guards that neither an +fork by ``update_state`` / ``aupdate_state``, fork before the delta channel +ever had a value, and guards that neither an unaddressed run nor an unaddressed multi-superstep ``bulk_update_state`` departs from the normal ``snapshot_frequency`` cadence. """ @@ -46,6 +47,9 @@ def _append(current: list | None, writes: Sequence[Any]) -> list: class _State(TypedDict): log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)] plain: Annotated[list, add] + # Written only by the tests that fork before `log` ever has a value, so the + # fork advances without touching the delta channel. + other: Annotated[list, add] def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: @@ -65,6 +69,19 @@ def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: return builder.compile(checkpointer=checkpointer) +def _build_without_delta_writes(checkpointer: BaseCheckpointSaver, tag: str) -> Any: + """Compile a graph whose node writes only ``other``, never the delta channel.""" + + def node(state: _State) -> dict: + return {"other": [f"{tag}-other"]} + + 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}} @@ -257,6 +274,81 @@ def test_unaddressed_run_keeps_snapshot_cadence( assert not _snapshotted_checkpoints(sync_checkpointer, config) +def test_fork_before_first_value_when_fork_never_writes_the_channel( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Seal the fork even when the channel has no value to snapshot. + + Forking before ``log`` was ever written leaves nothing to copy into the + fork's first checkpoint, so without a minted version and an empty blob the + boundary goes unrecorded and the walk runs into the base. The fork here + never writes ``log`` at all, so no later superstep can seal it either. + """ + 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_without_delta_writes(sync_checkpointer, "third").invoke( + {"other": ["in-9"]}, _at(config, root), durability=durability + ) + + state = graph.get_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == [] + + +async def test_afork_before_first_value_when_fork_never_writes_the_channel( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Async twin of ``test_fork_before_first_value_when_fork_never_writes_the_channel``.""" + 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_without_delta_writes(async_checkpointer, "third").ainvoke( + {"other": ["in-9"]}, _at(config, root), durability=durability + ) + + state = await graph.aget_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == [] + + +def test_fork_before_first_value_by_bulk_update( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """The fork's first superstep touches another key, the delta key comes later. + + The first checkpoint has to seal the boundary on its own. Deferring until + the superstep that finally writes ``log`` is too late, because that + superstep reconstructs through the unsealed checkpoint first. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config) + + root = list(graph.get_state_history(config))[-1] + assert root.values["log"] == [] + + forked = graph.bulk_update_state( + _at(config, root), + [ + [StateUpdate({"other": ["s1"]}, "n")], + [StateUpdate(_input("s2"), "n")], + ], + ) + + state = graph.get_state(forked) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["s2"] + + @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