WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 09:32:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(langgraph): seal a fork on the first checkpoint it writes

Browse Source 
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.
This commit is contained in:
Elior Nataf Lackritz
2026-08-05 23:42:35 -04:00
parent c3f4947151
commit 834e53df19
3 changed files with 162 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/_checkpoint.py
+35
View File
@@ -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,
libs/langgraph/langgraph/pregel/main.py
+83 -22
View File
@@ -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(
libs/langgraph/tests/test_delta_channel_fork.py
+44 -1
View File
@@ -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: