diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index eb66e531e..6705c9c17 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1638,7 +1638,10 @@ class Pregel( 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 @@ class Pregel( 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,7 +2026,7 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, - is_fork=bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)), + is_fork=is_fork, ) ) checkpoint = create_checkpoint( @@ -2050,8 +2056,16 @@ 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)) 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( @@ -2105,7 +2119,10 @@ class Pregel( 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) @@ -2336,6 +2353,9 @@ class Pregel( 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( @@ -2481,7 +2501,7 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, - is_fork=bool(config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)), + is_fork=is_fork, ) ) checkpoint = create_checkpoint( @@ -2510,8 +2530,16 @@ 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)) 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( diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index ea228284b..c7d56f5d5 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -11,9 +11,10 @@ same values. ``full`` channels store complete ``channel_values`` and need no 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, fork by -``update_state`` / ``aupdate_state``, and a guard that an unaddressed run still -follows the normal ``snapshot_frequency`` cadence. +fork off the checkpoint that predates the thread's first input (sync/async), +fork by ``update_state`` / ``aupdate_state``, and guards that neither an +unaddressed run nor an unaddressed multi-superstep ``bulk_update_state`` +departs from the normal ``snapshot_frequency`` cadence. """ from collections.abc import Sequence @@ -28,7 +29,7 @@ from typing_extensions import TypedDict from langgraph.channels.delta import DeltaChannel from langgraph.graph import StateGraph -from langgraph.types import Durability, StateSnapshot +from langgraph.types import Durability, StateSnapshot, StateUpdate pytestmark = pytest.mark.anyio @@ -82,6 +83,17 @@ def _input(marker: str) -> dict: return {"log": [marker], "plain": [marker]} +def _snapshotted_checkpoints( + checkpointer: BaseCheckpointSaver, config: RunnableConfig +) -> list[str]: + """Ids of this thread's checkpoints carrying a ``log`` snapshot blob.""" + return [ + tuple_.config["configurable"]["checkpoint_id"] + for tuple_ in checkpointer.list(config) + if isinstance(tuple_.checkpoint["channel_values"].get("log"), _DeltaSnapshot) + ] + + def _assert_fork_is_clean(state: StateSnapshot, abandoned: str) -> None: """The delta channel must match the plain channel and drop ``abandoned``.""" assert state.values["log"] == state.values["plain"], ( @@ -241,8 +253,27 @@ def test_unaddressed_run_keeps_snapshot_cadence( graph.invoke(_input("in-1"), config, durability=durability) graph.invoke(_input("in-2"), config, durability=durability) - assert not [ - tuple_.config["configurable"]["checkpoint_id"] - for tuple_ in sync_checkpointer.list(config) - if isinstance(tuple_.checkpoint["channel_values"].get("log"), _DeltaSnapshot) - ] + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def test_unaddressed_bulk_update_keeps_snapshot_cadence( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """A multi-superstep ``bulk_update_state`` forks at most once, at the head. + + ``perform_superstep`` returns the config of the checkpoint it just wrote + and the driver feeds that back in, so every superstep after the first + receives a config naming a checkpoint even when the caller addressed none. + Deriving the fork flag from that config snapshots the whole growing value + once per superstep. + """ + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_input("in-1"), config) + + graph.bulk_update_state( + config, + [[StateUpdate(_input(f"u{i}"), "n")] for i in range(4)], + ) + + assert not _snapshotted_checkpoints(sync_checkpointer, config)