diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 127700728..a2454c805 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -83,13 +83,7 @@ def get_delta_channels_from_all_channels( *, include_unavailable: bool = False, ) -> set[str]: - """Every available DeltaChannel. - - The set to snapshot whenever no ancestor walk can reconstruct these - channels: the first update_state of a fresh thread (no ancestors at all), - and the first checkpoint of a fork (whose base also holds the writes of the - branch the fork abandons). - """ + """DeltaChannels to snapshot on the first update_state of a fresh thread or fork.""" return { k for k, ch in channels.items() @@ -134,13 +128,8 @@ def create_checkpoint_plan_for_update_state_api( ) -> tuple[set[str], dict[str, Any]]: """Return ``(channels_to_snapshot, metadata)`` for an update_state head. - ``is_fork`` (the update was addressed at an explicit checkpoint) forces a - full snapshot for the same reason ``is_fresh_thread`` does: the ancestor - walk cannot reconstruct this head. The base a fork branches off keeps the - pending writes of the branch being abandoned, and nothing records which - child consumed which write, so the walk would replay them here too. - Snapshotting terminates the walk at this checkpoint. Every delta channel - snapshots, so no counters carry over. + A fork snapshots everything, like a fresh thread: its base also holds the + writes of the branch it abandons, so the ancestor walk must stop here. """ metadata: dict[str, Any] = { "source": "update", @@ -174,21 +163,12 @@ def create_fork_checkpoint( is_fork: bool, get_next_version: GetNextVersion, ) -> Checkpoint: - """``create_checkpoint`` for an update_state path that bypasses the plan. + """``create_checkpoint`` for the update_state paths that skip 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. + The fork has to be sealed by its first checkpoint: any later superstep + has already rebuilt its delta channels through the shared base. These + paths never write the delta channel, so its version must be bumped here + or ``put`` drops the blob; derive ``new_versions`` from the result. """ if not is_fork: return create_checkpoint(checkpoint, channels, step) @@ -233,12 +213,9 @@ def create_checkpoint( for k in channels: 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. + # A forced snapshot of a never-written channel still has to + # land to stop the ancestor walk, and `put` only stores blobs + # for versioned channels. if k in channels_to_snapshot and get_next_version is not None: channel_versions[k] = get_next_version(None, None) values[k] = _DeltaSnapshot( @@ -249,9 +226,8 @@ def create_checkpoint( # Callers force a full snapshot blob here: exit mode when a # delta channel reaches its snapshot cadence, update_state on # a fresh thread (no ancestor to replay writes from), and a - # fork (whose base also holds the abandoned branch's writes). - # The manual version-bump below only applies to the exit-mode - # case. + # fork. The manual version-bump below only applies to the + # exit-mode case. # # In exit mode, the snapshot decision is deferred to exit # time (intermediate steps have do_checkpoint=False). The diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 21989fdf3..ab1cf7a35 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -223,32 +223,16 @@ class PregelLoop: _exit_delta_writes: list[tuple[int, str, str, Any]] | None = None # Delta channels that must snapshot at the next checkpoint, whatever their - # cadence counters say. Two sources: - # * an Overwrite arrived since the last checkpoint, so the snapshot has to - # happen after live update applied overwrite semantics and sparse replay - # starts from the same post-overwrite value; - # * this run forked off an explicitly addressed checkpoint, see - # `_delta_channels_awaiting_fork_snapshot`. + # cadence counters say: + # * an Overwrite arrived since the last checkpoint, so sparse replay has to + # start from the post-overwrite value; + # * this run forked off an explicitly addressed checkpoint. That base also + # holds the writes of the branch the fork abandons, and nothing records + # which child consumed which, so the ancestor walk must stop inside the + # fork. Any addressed checkpoint counts, because telling a real fork + # apart would mean trusting the base's `pending_writes` to be complete. _delta_channels_forced_snapshot: set[str] - # Delta channels still owed a fork snapshot, when this run was launched - # against an explicitly addressed checkpoint (time travel / fork). That - # base keeps the pending writes of the branch the fork abandons, and - # nothing records which child consumed which write, so the ancestor walk - # would replay them into this branch too. Snapshotting terminates the walk - # inside the fork instead of at the shared base. Names drop out once the - # blob has landed; a channel with no value yet has nothing to snapshot, so - # it waits for the superstep that gives it one. - # - # The trigger is deliberately coarse: any addressed checkpoint, not only - # one that turns out to have abandoned writes on it. Which writes belong to - # which child is exactly what is not recorded, so a narrower test would - # have to trust the base's `pending_writes` to be complete, and a saver - # that leaves them out would silently go back to leaking. Snapshotting when - # it was not needed costs one blob per addressed run; not snapshotting when - # it was needed is silent corruption. - _delta_channels_awaiting_fork_snapshot: set[str] - # The checkpoint_config that points at the parent loaded at `__enter__` # (or the synthetic-empty checkpoint, on first run). We capture it # eagerly because every `_put_checkpoint` advances `self.checkpoint_config` @@ -391,12 +375,9 @@ class PregelLoop: if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) else () ) - # Checks the value, not just key presence like `is_replaying` above: - # subgraph task configs always carry an explicit `None` here, and only - # a real id means the caller addressed one specific checkpoint. Read - # off `checkpoint_config` so subgraphs resolved through a checkpoint - # map during time travel are covered too, matching `__enter__`. - self._delta_channels_awaiting_fork_snapshot = ( + # Value, not key presence like `is_replaying`: subgraph task configs + # always carry an explicit `None` checkpoint_id. + self._delta_channels_forced_snapshot = ( {k for k, spec in specs.items() if isinstance(spec, DeltaChannel)} if self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) else set() @@ -1165,17 +1146,6 @@ class PregelLoop: do_checkpoint = self._checkpointer_put_after_previous is not None and ( exiting or self.durability != "exit" ) - # Fork: make this checkpoint self-contained, so the ancestor walk stops - # inside the fork instead of reaching the base this run forked off and - # collecting the abandoned branch's writes from it. Resolved here - # rather than in `_first` so channels that only got a value this - # superstep are covered too. - if self._delta_channels_awaiting_fork_snapshot: - self._delta_channels_forced_snapshot.update( - k - for k in self._delta_channels_awaiting_fork_snapshot - if k in self.channels - ) # create new checkpoint channels_to_snapshot = ( delta_channels_to_snapshot(self.channels, new_counters) @@ -1198,12 +1168,6 @@ class PregelLoop: new_counters[k] = (0, 0) if do_checkpoint: self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot) - # `create_checkpoint` drops a requested snapshot for a channel with - # no version in this checkpoint yet (nothing was ever written to it - # on this branch), so keep asking until the blob really landed. - self._delta_channels_awaiting_fork_snapshot.difference_update( - self.checkpoint["channel_values"] - ) non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)} if non_zero: self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero @@ -1733,7 +1697,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_forced_snapshot = set() self._exit_delta_writes = ( [] if self.durability == "exit" and self.checkpointer is not None else None ) @@ -1991,7 +1954,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_forced_snapshot = set() self._exit_delta_writes = ( [] if self.durability == "exit" and self.checkpointer is not None else None ) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index c32636d39..8a39d6946 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1639,13 +1639,10 @@ 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. + # Read once from the caller's config: every later superstep receives + # the config of the checkpoint just written, which always names one. + # Cleared by the first checkpoint that carries the snapshots, which + # `__copy__` does not write. fork_pending: set[str] = ( {k for k, v in self.channels.items() if isinstance(v, DeltaChannel)} if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) @@ -1907,8 +1904,6 @@ 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, ) @@ -2089,10 +2084,6 @@ class Pregel( current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_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=bool(fork_pending) @@ -2149,13 +2140,10 @@ 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. + # Read once from the caller's config: every later superstep receives + # the config of the checkpoint just written, which always names one. + # Cleared by the first checkpoint that carries the snapshots, which + # `__copy__` does not write. fork_pending: set[str] = ( {k for k, v in self.channels.items() if isinstance(v, DeltaChannel)} if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID) @@ -2414,8 +2402,6 @@ 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, ) @@ -2593,10 +2579,6 @@ class Pregel( current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_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=bool(fork_pending) diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index bc406e918..136fcc34c 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -85,15 +85,9 @@ class MemorySaverAssertImmutable(InMemorySaver): ) == saved ), config["configurable"]["checkpoint_ns"] - # call super to write checkpoint next_config = super().put(config, checkpoint, metadata, new_versions) - # Record the checkpoint as the saver reads it back, not the object it - # was handed. `channel_values` are stored per (channel, version), so a - # channel a step did not write is refilled from the blob its inherited - # version still points at. A `DeltaChannel` omits its value except at a - # snapshot, which makes the two representations differ for reasons that - # are not mutation. Comparing read-back against read-back still catches - # a checkpoint whose stored data actually changed. + # Read back, not the object handed in: a DeltaChannel a step did not + # write is refilled on read from the blob its inherited version points at. self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = ( self.serde.dumps_typed(super().get(next_config)) ) diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py index d53a86689..b55fd3563 100644 --- a/libs/langgraph/tests/test_delta_channel_fork.py +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -1,21 +1,7 @@ """Forking a thread must not replay the abandoned branch into the fork. -Regression suite for #8443. Addressing an older checkpoint creates a fork: the -shared base ends up with two children, and it keeps the ``checkpoint_writes`` -of the branch the fork abandons. Nothing in the stored data records which child -consumed which write, so the ``DeltaChannel`` ancestor walk used to collect the -abandoned branch's writes as well. - -Every graph here carries a ``DeltaChannel`` and a plain reducer channel fed the -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 (sync/async), -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. +Every graph carries a ``DeltaChannel`` and a plain reducer channel fed the same +values; the plain channel needs no replay, so it is the oracle. """ from collections.abc import Sequence @@ -37,7 +23,6 @@ pytestmark = pytest.mark.anyio def _append(current: list | None, writes: Sequence[Any]) -> list: - """DeltaChannel reducer: extend the list with every batched write.""" out = list(current or []) for write in writes: out.extend(write if isinstance(write, list) else [write]) @@ -47,18 +32,10 @@ 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: - """Compile a one-node graph whose node appends ``{tag}-out`` to both channels. - - ``snapshot_frequency=1000`` keeps the cadence from masking the bug: without - a forced snapshot the fork's ancestor walk always runs past the fork base. - """ - def node(state: _State) -> dict: return {"log": [f"{tag}-out"], "plain": [f"{tag}-out"]} @@ -70,8 +47,6 @@ def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: 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"]} @@ -87,7 +62,6 @@ def _thread(thread_id: str) -> RunnableConfig: def _at(config: RunnableConfig, snapshot: StateSnapshot) -> RunnableConfig: - """Config addressing one specific checkpoint of ``config``'s thread.""" return { "configurable": { **config["configurable"], @@ -104,7 +78,6 @@ def _input(marker: str) -> dict: 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) @@ -113,7 +86,6 @@ def _snapshotted_checkpoints( 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"], ( f"delta channel diverged from the plain channel: " f"{state.values['log']} != {state.values['plain']}" @@ -133,9 +105,8 @@ def test_fork_by_invoke( ) graph = _build(sync_checkpointer, "second") graph.invoke(_input("in-2"), config, durability=durability) + abandoned_head = graph.get_state(config) - # The last checkpoint that predates "in-2" entering state: forking here - # abandons the "in-2" branch, whose writes still hang off this checkpoint. base = next( snapshot for snapshot in graph.get_state_history(config) @@ -149,6 +120,9 @@ def test_fork_by_invoke( _assert_fork_is_clean(state, "in-2") assert state.values["log"] == [*base.values["log"], "in-3", "third-out"] + abandoned = graph.get_state(abandoned_head.config).values + assert abandoned["log"] == abandoned["plain"] == abandoned_head.values["log"] + async def test_afork_by_invoke( async_checkpointer: BaseCheckpointSaver, durability: Durability @@ -159,6 +133,7 @@ async def test_afork_by_invoke( ) graph = _build(async_checkpointer, "second") await graph.ainvoke(_input("in-2"), config, durability=durability) + abandoned_head = await graph.aget_state(config) base = await anext( snapshot @@ -173,17 +148,13 @@ async def test_afork_by_invoke( _assert_fork_is_clean(state, "in-2") assert state.values["log"] == [*base.values["log"], "in-3", "third-out"] + abandoned = (await graph.aget_state(abandoned_head.config)).values + assert abandoned["log"] == abandoned["plain"] == abandoned_head.values["log"] + def test_fork_off_checkpoint_before_first_input( sync_checkpointer: BaseCheckpointSaver, durability: Durability ) -> None: - """Fork off the root checkpoint, which predates any value for ``log``. - - ``create_checkpoint`` drops a requested snapshot for a channel absent from - ``channel_versions``, so the fork's own first checkpoint cannot carry the - blob. The request has to stay queued until a superstep gives the channel a - value, otherwise the root's ``in-1`` write still leaks into the fork. - """ config = _thread("t") graph = _build(sync_checkpointer, "first") graph.invoke(_input("in-1"), config, durability=durability) @@ -203,7 +174,6 @@ def test_fork_off_checkpoint_before_first_input( async def test_afork_off_checkpoint_before_first_input( async_checkpointer: BaseCheckpointSaver, durability: Durability ) -> None: - """Async twin of ``test_fork_off_checkpoint_before_first_input``.""" config = _thread("t") graph = _build(async_checkpointer, "first") await graph.ainvoke(_input("in-1"), config, durability=durability) @@ -261,11 +231,6 @@ async def test_afork_by_update_state( def test_unaddressed_run_keeps_snapshot_cadence( sync_checkpointer: BaseCheckpointSaver, durability: Durability ) -> None: - """A run with no explicitly addressed checkpoint writes no snapshot blob. - - Guards the cost of the fix: the forced snapshot is one per addressed run, - not a change to the normal ``snapshot_frequency`` cadence. - """ config = _thread("t") graph = _build(sync_checkpointer, "first") graph.invoke(_input("in-1"), config, durability=durability) @@ -277,13 +242,6 @@ def test_unaddressed_run_keeps_snapshot_cadence( 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) @@ -303,7 +261,6 @@ def test_fork_before_first_value_when_fork_never_writes_the_channel( 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) @@ -323,12 +280,6 @@ async def test_afork_before_first_value_when_fork_never_writes_the_channel( 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) @@ -353,15 +304,6 @@ def test_fork_before_first_value_by_bulk_update( 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") @@ -383,8 +325,6 @@ def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan( ) 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']}" @@ -394,14 +334,6 @@ def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan( 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)