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.
This commit is contained in:
Elior Nataf Lackritz
2026-08-06 11:37:10 -04:00
parent 6b3d0dc314
commit 4b0e98113b
3 changed files with 116 additions and 7 deletions
+22 -5
View File
@@ -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
+1 -1
View File
@@ -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 = (
@@ -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