diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index 58a4e7917..7a369b891 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -23,6 +23,9 @@ RETURN = sys.intern("__return__") # for writes of a task where we simply record the return value PREVIOUS = sys.intern("__previous__") # the implicit branch that handles each node's Control values +SNAPSHOT_BUMPS = sys.intern("__snapshot_bumps__") +# `versions_seen` key recording which entries a snapshot-only version bump +# advanced, and the versions those nodes really read # --- Reserved cache namespaces --- @@ -116,6 +119,7 @@ RESERVED = { ERROR, ERROR_SOURCE_NODE, NO_WRITES, + SNAPSHOT_BUMPS, # reserved config.configurable keys CONFIG_KEY_SEND, CONFIG_KEY_READ, diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 69ca88708..42c8f51a0 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -8,13 +8,15 @@ from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, + ChannelVersions, Checkpoint, + PendingWrite, ) from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.types import _DeltaSnapshot from langgraph._internal._config import DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT -from langgraph._internal._constants import PUSH +from langgraph._internal._constants import NS_END, NS_SEP, PUSH, SNAPSHOT_BUMPS from langgraph._internal._typing import MISSING from langgraph.channels.base import BaseChannel from langgraph.channels.delta import DeltaChannel @@ -50,17 +52,23 @@ def exit_delta_task_id(step: int, task_id: str) -> str: def delta_channels_to_snapshot( channels: Mapping[str, BaseChannel], counters_since_delta_snapshot: Mapping[str, tuple[int, int]], + channel_versions: ChannelVersions, ) -> set[str]: """Return the set of DeltaChannel names that should snapshot now. A channel snapshots when EITHER its accumulated update count reaches `snapshot_frequency` OR the total supersteps since its last snapshot - reaches `DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT`. This is a pure - predicate — no mutation. + reaches `DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT`. A channel without a version + was never written on this branch, so it has nothing to snapshot. This is a + pure predicate — no mutation. """ result: set[str] = set() for name, ch in channels.items(): - if not isinstance(ch, DeltaChannel) or not ch.is_available(): + if ( + not isinstance(ch, DeltaChannel) + or not ch.is_available() + or name not in channel_versions + ): continue updates, supersteps = counters_since_delta_snapshot.get(name, (0, 0)) if ( @@ -80,12 +88,31 @@ def get_updated_channels_from_tasks( def get_delta_channels_from_all_channels( channels: Mapping[str, BaseChannel], + channel_versions: ChannelVersions, ) -> set[str]: - """DeltaChannels to snapshot on the first update_state of a fresh thread.""" + """DeltaChannels to snapshot on the first update_state of a fresh thread: + the ones it wrote, as the rest have no version and nothing to store.""" return { k for k, ch in channels.items() - if isinstance(ch, DeltaChannel) and ch.is_available() + if isinstance(ch, DeltaChannel) and ch.is_available() and k in channel_versions + } + + +def delta_channels_with_pending_writes( + specs: Mapping[str, Any], + pending_writes: Iterable[PendingWrite] | None, +) -> set[str]: + """DeltaChannels a branch starting from this checkpoint must snapshot. + + A checkpoint's pending writes belong to the child that consumed them, and + nothing records which child that was. A new branch snapshots every delta + channel they touch, so its ancestor walk never replays them. + """ + return { + ch + for _, ch, _ in pending_writes or () + if isinstance(specs.get(ch), DeltaChannel) } @@ -122,6 +149,8 @@ def create_checkpoint_plan_for_update_state_api( parents: dict[str, Any], saved_metadata: Mapping[str, Any] | None, is_fresh_thread: bool, + fork_channels: set[str], + channel_versions: ChannelVersions, ) -> tuple[set[str], dict[str, Any]]: """Return ``(channels_to_snapshot, metadata)`` for an update_state head.""" metadata: dict[str, Any] = { @@ -130,14 +159,19 @@ def create_checkpoint_plan_for_update_state_api( "parents": parents, } if is_fresh_thread: - return get_delta_channels_from_all_channels(channels), metadata + return get_delta_channels_from_all_channels( + channels, channel_versions + ), metadata new_counters = create_metadata_for_update_state_api( channels, updated_channels, prev_metadata=saved_metadata, ) - channels_to_snapshot = delta_channels_to_snapshot(channels, new_counters) + channels_to_snapshot = ( + delta_channels_to_snapshot(channels, new_counters, channel_versions) + | fork_channels + ) for k in channels_to_snapshot: new_counters[k] = (0, 0) non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)} @@ -155,6 +189,7 @@ def create_checkpoint( updated_channels: set[str] | None = None, get_next_version: GetNextVersion | None = None, channels_to_snapshot: set[str] | None = None, + stored_versions: ChannelVersions | None = None, ) -> Checkpoint: """Build a new Checkpoint from the previous one and live channel state. @@ -164,9 +199,14 @@ def create_checkpoint( from `checkpoint_writes`. Callers compute the set via `delta_channels_to_snapshot(channels, counters)`; defaults to empty (no snapshots) when not provided. + + `stored_versions` are the channel versions of the last checkpoint the + saver stored. A snapshotted channel whose version has not moved since + then is bumped. """ ts = datetime.now(timezone.utc).isoformat() channels_to_snapshot = channels_to_snapshot or set() + bumped: dict[str, tuple[Any, Any]] = {} if channels is None: values = checkpoint["channel_values"] channel_versions = checkpoint["channel_versions"] @@ -174,30 +214,29 @@ 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: + # 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) + bumped[k] = (None, channel_versions[k]) + values[k] = _DeltaSnapshot(ch.get()) + continue if k in channels_to_snapshot: - # Callers force a full snapshot blob here: exit mode when a - # delta channel reaches its snapshot cadence, and update_state - # on a fresh thread (no ancestor to replay writes from). 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 - # channel's count may have reached snapshot_frequency over - # several supersteps, but the LAST superstep may not have - # written to this channel. In that case apply_writes() - # (in _algo.py) didn't bump this channel's version, so - # saver.put() wouldn't include it in new_versions and - # the snapshot blob would be silently dropped. The manual - # bump below closes the gap. In sync/async durability this - # branch is effectively dead code (the step that pushes - # the count to freq always writes the channel). - if get_next_version is not None and ( - updated_channels is None or k not in updated_channels + # `put` only stores a blob for a channel whose version moved, + # so snapshotting a channel this step did not write needs a + # bump: exit mode reaching the cadence on a superstep that + # skipped the channel, and a fork's first checkpoint. + if ( + get_next_version is not None + and stored_versions is not None + and channel_versions[k] == stored_versions.get(k) ): - channel_versions[k] = get_next_version(channel_versions[k], None) + old = channel_versions[k] + channel_versions[k] = get_next_version(old, None) + bumped[k] = (old, channel_versions[k]) values[k] = _DeltaSnapshot(ch.get()) else: v = ch.checkpoint() @@ -209,11 +248,69 @@ def create_checkpoint( id=id or str(uuid6(clock_seq=step)), channel_values=values, channel_versions=channel_versions, - versions_seen=checkpoint["versions_seen"], + versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped), updated_channels=None if updated_channels is None else sorted(updated_channels), ) +def _mark_bumps_seen( + versions_seen: dict[str, ChannelVersions], + bumped: Mapping[str, tuple[Any, Any]], +) -> dict[str, ChannelVersions]: + """Advance whoever had seen a bumped channel's old version to the new one. + + A bump that only stores a snapshot is not a write. Left unseen, it would + re-fire `interrupt_before` and rerun the channel's subscribers. For each + entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the + node really read, so `versions_seen_without_bumps` can put the read back. + """ + if not bumped: + return versions_seen + out = dict(versions_seen) + marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {})) + for node, seen in versions_seen.items(): + if node == SNAPSHOT_BUMPS: + continue + for k, (old, new) in bumped.items(): + if seen.get(k) != old: + continue + advanced, read = _bump_keys(node, k) + # If an earlier bump set `old`, the real read is already recorded. + if old is not None and marks.get(advanced) != old: + marks[read] = old + marks[advanced] = new + out[node] = {**out[node], k: new} + if marks: + out[SNAPSHOT_BUMPS] = marks + return out + + +def _bump_keys(node: str, channel: str) -> tuple[str, str]: + # Node names can't contain either separator, so the keys can't collide. + return f"{node}{NS_SEP}{channel}", f"{node}{NS_END}{channel}" + + +def versions_seen_without_bumps( + versions_seen: dict[str, ChannelVersions], +) -> dict[str, ChannelVersions]: + """`versions_seen` as the nodes read it: an entry a bump advanced goes back + to the version the node really read, or away if it never read one.""" + if not (marks := versions_seen.get(SNAPSHOT_BUMPS)): + return versions_seen + out: dict[str, ChannelVersions] = {} + for node, seen in versions_seen.items(): + if node == SNAPSHOT_BUMPS: + continue + out[node] = {} + for k, v in seen.items(): + advanced, read = _bump_keys(node, k) + if marks.get(advanced) != v: + out[node][k] = v + elif read in marks: + out[node][k] = marks[read] + return out + + def _needs_replay(spec: BaseChannel, stored: object) -> bool: """True if `spec` is a `DeltaChannel` and no value is stored at this checkpoint, requiring an ancestor walk to reconstruct. diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 796fb5efc..a079fa4b0 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -102,6 +102,7 @@ from langgraph.pregel._checkpoint import ( copy_checkpoint, create_checkpoint, delta_channels_to_snapshot, + delta_channels_with_pending_writes, empty_checkpoint, exit_delta_task_id, ) @@ -223,10 +224,18 @@ class PregelLoop: # under the saver's `ORDER BY task_id, idx` sorting. _exit_delta_writes: list[tuple[int, str, str, Any]] | None = None - # Delta channels that saw an Overwrite since the last checkpoint. These - # channels must snapshot after live update applies overwrite semantics so - # sparse replay starts from the same post-overwrite value. - _delta_channels_with_overwrite: set[str] + # Delta channels that must snapshot at the next checkpoint, whatever their + # cadence counters say: + # * an Overwrite arrived since the last checkpoint, so sparse replay has to + # start from the post-overwrite value; + # * the checkpoint this run starts from has pending writes to them; see + # `delta_channels_with_pending_writes`. + _delta_channels_forced_snapshot: set[str] + # Set by `_first` for a resume: the writes it loaded with the checkpoint, + # which `after_tick` checks against the ones reapply handed back. + _resume_loaded_writes: Sequence[PendingWrite] = () + # Tasks `_reapply_writes_to_succeeded_nodes` handed loaded writes back to. + _reapplied_task_ids: set[str] # The checkpoint_config that points at the parent loaded at `__enter__` # (or the synthetic-empty checkpoint, on first run). We capture it @@ -582,7 +591,7 @@ class PregelLoop: # save the new task self.tasks[pushed.id] = pushed # match any pending writes to the new task - if not self.is_replaying: + if self._reapplies_pending_writes: self._reapply_writes_to_succeeded_nodes({pushed.id: pushed}) # return the new task, to be started if not run before return pushed @@ -660,7 +669,7 @@ class PregelLoop: return False # if there are pending writes from a previous loop, apply them - if not self.is_replaying and self.checkpoint_pending_writes: + if self._reapplies_pending_writes and self.checkpoint_pending_writes: self._reapply_writes_to_succeeded_nodes(self.tasks) self._resume_error_handlers_if_applicable() @@ -684,11 +693,26 @@ class PregelLoop: def after_tick(self) -> None: # finish superstep writes = [w for t in self.tasks.values() for w in t.writes] - self._delta_channels_with_overwrite.update( + self._delta_channels_forced_snapshot.update( ch for ch, v in writes if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0] ) + if self._resume_loaded_writes: + # A loaded write reapply didn't hand back belongs to a task this run + # drops or reruns: a `Send` that `Command(goto=...)` replaced, or an + # error handler that runs again. + self._delta_channels_forced_snapshot.update( + delta_channels_with_pending_writes( + self.specs, + [ + w + for w in self._resume_loaded_writes + if w[0] not in self._reapplied_task_ids + ], + ) + ) + self._resume_loaded_writes = () # all tasks have finished self.updated_channels = apply_writes( self.checkpoint, @@ -734,6 +758,12 @@ class PregelLoop: # private + @property + def _reapplies_pending_writes(self) -> bool: + """Whether the writes loaded with the checkpoint go back to the tasks + that made them, instead of those tasks rerunning.""" + return not self.is_replaying + def _reapply_writes_to_succeeded_nodes( self, tasks: Mapping[str, PregelExecutableTask] ) -> None: @@ -745,6 +775,7 @@ class PregelLoop: for tid, status in read_task_statuses(self.checkpoint_pending_writes).items(): if task := tasks.get(tid): task.writes.extend(status.output) + self._reapplied_task_ids.add(tid) def _resume_error_handlers_if_applicable(self) -> None: """On resume, schedule error handlers for tasks that failed in a prior run. @@ -874,6 +905,23 @@ class PregelLoop: self.checkpoint_pending_writes = [ w for w in self.checkpoint_pending_writes if w[1] != RESUME ] + # A resume that reapplies the head's pending writes only learns which + # of them go back to their tasks once those are scheduled, so + # `after_tick` seals the rest. Kept apart from the writes this run adds. + reapplies = is_resuming and self._reapplies_pending_writes + self._resume_loaded_writes = ( + [w for w in self.checkpoint_pending_writes if w[0] != NULL_TASK_ID] + if reapplies + else () + ) + self._reapplied_task_ids = set() + self._delta_channels_forced_snapshot = ( + set() + if reapplies + else delta_channels_with_pending_writes( + self.specs, self.checkpoint_pending_writes + ) + ) # map command to writes if input_is_command: @@ -967,7 +1015,7 @@ class PregelLoop: manager=None, updated_channels=updated_channels, ) - self._delta_channels_with_overwrite.update( + self._delta_channels_forced_snapshot.update( c for c, v in input_writes if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0] @@ -1111,8 +1159,10 @@ class PregelLoop: ) # create new checkpoint channels_to_snapshot = ( - delta_channels_to_snapshot(self.channels, new_counters) - | self._delta_channels_with_overwrite + delta_channels_to_snapshot( + self.channels, new_counters, self.checkpoint["channel_versions"] + ) + | self._delta_channels_forced_snapshot if do_checkpoint else set() ) @@ -1126,11 +1176,12 @@ class PregelLoop: if do_checkpoint else None, channels_to_snapshot=channels_to_snapshot, + stored_versions=self.checkpoint_previous_versions, ) for k in channels_to_snapshot: new_counters[k] = (0, 0) if do_checkpoint: - self._delta_channels_with_overwrite.difference_update(channels_to_snapshot) + self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot) 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 @@ -1214,8 +1265,10 @@ class PregelLoop: self.checkpoint_metadata.get("counters_since_delta_snapshot") or {} ) channels_to_snapshot = ( - delta_channels_to_snapshot(self.channels, counters) - | self._delta_channels_with_overwrite + delta_channels_to_snapshot( + self.channels, counters, self.checkpoint["channel_versions"] + ) + | self._delta_channels_forced_snapshot ) pending = [ @@ -1576,7 +1629,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if handler_task is None: return None self.tasks[handler_task.id] = handler_task - if not self.is_replaying: + if self._reapplies_pending_writes: self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task}) for task in self.match_cached_writes(): self.output_writes(task.id, task.writes, cached=True) @@ -1660,7 +1713,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_with_overwrite = set() self._exit_delta_writes = ( [] if self.durability == "exit" and self.checkpointer is not None else None ) @@ -1831,7 +1883,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if handler_task is None: return None self.tasks[handler_task.id] = handler_task - if not self.is_replaying: + if self._reapplies_pending_writes: self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task}) for task in await self.amatch_cached_writes(): self.output_writes(task.id, task.writes, cached=True) @@ -1918,7 +1970,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): ) self._delta_write_futs = [] self._error_handler_write_futs = [] - self._delta_channels_with_overwrite = 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 c32f8cb37..ffb828cbc 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -46,6 +46,7 @@ from langchain_core.runnables.schema import StreamEvent from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( BaseCheckpointSaver, + ChannelVersions, Checkpoint, CheckpointTuple, ) @@ -132,8 +133,10 @@ from langgraph.pregel._checkpoint import ( copy_checkpoint, create_checkpoint, create_checkpoint_plan_for_update_state_api, + delta_channels_with_pending_writes, empty_checkpoint, get_updated_channels_from_tasks, + versions_seen_without_bumps, ) from langgraph.pregel._draw import draw_graph from langgraph.pregel._io import map_input, read_channels @@ -1585,6 +1588,27 @@ class Pregel( checkpoint_tuple.config, checkpoint_tuple, saver=checkpointer ) + def _infer_as_node(self, versions_seen: dict[str, ChannelVersions]) -> str | None: + if len(self.nodes) == 1: + return next(iter(self.nodes)) + seen = versions_seen_without_bumps(versions_seen) + if not any(v for vv in seen.values() for v in vv.values()): + if ( + isinstance(self.input_channels, str) + and self.input_channels in self.nodes + ): + return self.input_channels + return None + last_seen_by_node = sorted( + (v, n) for n, s in seen.items() if n in self.nodes for v in s.values() + ) + # if two nodes updated the state at the same time, it's ambiguous + if len(last_seen_by_node) == 1 or ( + last_seen_by_node and last_seen_by_node[-1][0] != last_seen_by_node[-2][0] + ): + return last_seen_by_node[-1][1] + return None + def bulk_update_state( self, config: RunnableConfig, @@ -1619,13 +1643,24 @@ class Pregel( return pregel.bulk_update_state(subgraph_config, supersteps) def perform_superstep( - input_config: RunnableConfig, updates: Sequence[StateUpdate] + input_config: RunnableConfig, + updates: Sequence[StateUpdate], + is_first: bool = False, ) -> RunnableConfig: # get last checkpoint config = self._own_checkpoint_config( ensure_config(self.config, input_config) ) saved = checkpointer.get_tuple(config) + # Later supersteps, including the one after a `__copy__` (stored + # under the base's parent), never walk through the base's writes. + fork_pending = ( + delta_channels_with_pending_writes( + self.channels, saved.pending_writes if saved else None + ) + if is_first + else set() + ) if saved is not None: self._migrate_checkpoint(saved.checkpoint) checkpoint = ( @@ -1705,9 +1740,17 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + next_checkpoint = create_checkpoint( + checkpoint, + channels, + step, + get_next_version=checkpointer.get_next_version, + channels_to_snapshot=fork_pending, + stored_versions=checkpoint_previous_versions, + ) next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + next_checkpoint, { "source": "update", "step": step + 1, @@ -1715,7 +1758,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) return patch_checkpoint_map( @@ -1744,9 +1787,17 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + next_checkpoint = create_checkpoint( + checkpoint, + channels, + next_step, + get_next_version=checkpointer.get_next_version, + channels_to_snapshot=fork_pending, + stored_versions=checkpoint_previous_versions, + ) next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + next_checkpoint, { "source": "input", "step": next_step, @@ -1756,7 +1807,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) @@ -1885,31 +1936,8 @@ class Pregel( if len(updates) == 1: values, as_node, task_id = updates[0] # find last node that updated the state, if not provided - if as_node is None and len(self.nodes) == 1: - as_node = tuple(self.nodes)[0] - elif as_node is None and not any( - v - for vv in checkpoint["versions_seen"].values() - for v in vv.values() - ): - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() - ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] + if as_node is None: + as_node = self._infer_as_node(checkpoint["versions_seen"]) if as_node is None: raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: @@ -1999,17 +2027,19 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], ) ) checkpoint = create_checkpoint( checkpoint, channels, step + 1, - updated_channels=updated_channels if channels_to_snapshot else None, get_next_version=checkpointer.get_next_version if channels_to_snapshot else None, channels_to_snapshot=channels_to_snapshot, + stored_versions=checkpoint_previous_versions, ) next_config = checkpointer.put( checkpoint_config, @@ -2028,8 +2058,8 @@ class Pregel( current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])} ) - for superstep in supersteps: - current_config = perform_superstep(current_config, superstep) + for i, superstep in enumerate(supersteps): + current_config = perform_superstep(current_config, superstep, i == 0) return current_config async def abulk_update_state( @@ -2066,13 +2096,24 @@ class Pregel( return await pregel.abulk_update_state(subgraph_config, supersteps) async def aperform_superstep( - input_config: RunnableConfig, updates: Sequence[StateUpdate] + input_config: RunnableConfig, + updates: Sequence[StateUpdate], + is_first: bool = False, ) -> RunnableConfig: # get last checkpoint config = self._own_checkpoint_config( ensure_config(self.config, input_config) ) saved = await checkpointer.aget_tuple(config) + # Later supersteps, including the one after a `__copy__` (stored + # under the base's parent), never walk through the base's writes. + fork_pending = ( + delta_channels_with_pending_writes( + self.channels, saved.pending_writes if saved else None + ) + if is_first + else set() + ) if saved is not None: self._migrate_checkpoint(saved.checkpoint) checkpoint = ( @@ -2150,16 +2191,25 @@ class Pregel( self.trigger_to_nodes, ) # save checkpoint + next_checkpoint = create_checkpoint( + checkpoint, + channels, + step, + get_next_version=checkpointer.get_next_version, + channels_to_snapshot=fork_pending, + stored_versions=checkpoint_previous_versions, + ) 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( @@ -2188,9 +2238,17 @@ class Pregel( if saved and saved.metadata.get("step") is not None else -1 ) + next_checkpoint = create_checkpoint( + checkpoint, + channels, + next_step, + get_next_version=checkpointer.get_next_version, + channels_to_snapshot=fork_pending, + stored_versions=checkpoint_previous_versions, + ) next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + next_checkpoint, { "source": "input", "step": next_step, @@ -2200,7 +2258,7 @@ class Pregel( }, get_new_channel_versions( checkpoint_previous_versions, - checkpoint["channel_versions"], + next_checkpoint["channel_versions"], ), ) @@ -2328,27 +2386,8 @@ class Pregel( if len(updates) == 1: values, as_node, task_id = updates[0] # find last node that updated the state, if not provided - if as_node is None and len(self.nodes) == 1: - as_node = tuple(self.nodes)[0] - elif as_node is None and not saved: - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() - ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] + if as_node is None: + as_node = self._infer_as_node(checkpoint["versions_seen"]) if as_node is None: raise InvalidUpdateError("Ambiguous update, specify as_node") if as_node not in self.nodes: @@ -2438,17 +2477,19 @@ class Pregel( parents=saved.metadata.get("parents", {}) if saved else {}, saved_metadata=saved.metadata if saved else None, is_fresh_thread=saved is None, + fork_channels=fork_pending, + channel_versions=checkpoint["channel_versions"], ) ) checkpoint = create_checkpoint( checkpoint, channels, step + 1, - updated_channels=updated_channels if channels_to_snapshot else None, get_next_version=checkpointer.get_next_version if channels_to_snapshot else None, channels_to_snapshot=channels_to_snapshot, + stored_versions=checkpoint_previous_versions, ) next_config = await checkpointer.aput( checkpoint_config, @@ -2466,8 +2507,8 @@ class Pregel( current_config = patch_configurable( config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])} ) - for superstep in supersteps: - current_config = await aperform_superstep(current_config, superstep) + for i, superstep in enumerate(supersteps): + current_config = await aperform_superstep(current_config, superstep, i == 0) return current_config def update_state( diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 3987d422b..136fcc34c 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -85,11 +85,13 @@ class MemorySaverAssertImmutable(InMemorySaver): ) == saved ), config["configurable"]["checkpoint_ns"] + next_config = super().put(config, checkpoint, metadata, new_versions) + # 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(checkpoint) + self.serde.dumps_typed(super().get(next_config)) ) - # call super to write checkpoint - return super().put(config, checkpoint, metadata, new_versions) + return next_config class MemorySaverNoPending(InMemorySaver): diff --git a/libs/langgraph/tests/test_delta_channel_fork.py b/libs/langgraph/tests/test_delta_channel_fork.py new file mode 100644 index 000000000..5736dc22e --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_fork.py @@ -0,0 +1,731 @@ +"""Forking a thread must not replay the abandoned branch into the fork. + +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. +""" + +import sys +import threading +from collections.abc import Sequence +from operator import add +from typing import Annotated, Any + +import pytest +from langchain_core.runnables import RunnableConfig +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.errors import NodeError +from langgraph.graph import END, START, StateGraph +from langgraph.types import ( + Command, + Durability, + Send, + StateSnapshot, + StateUpdate, + interrupt, +) + +pytestmark = pytest.mark.anyio +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +def _append(current: list | None, writes: Sequence[Any]) -> list: + out = list(current or []) + for write in writes: + out.extend(write if isinstance(write, list) else [write]) + return out + + +class _State(TypedDict): + log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)] + plain: Annotated[list, add] + other: Annotated[list, add] + + +def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any: + def node(state: _State) -> dict: + return {"log": [f"{tag}-out"], "plain": [f"{tag}-out"]} + + builder = StateGraph(_State) + builder.add_node("n", node) + builder.set_entry_point("n") + builder.set_finish_point("n") + return builder.compile(checkpointer=checkpointer) + + +def _build_without_delta_writes(checkpointer: BaseCheckpointSaver, tag: str) -> Any: + 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}} + + +def _at(config: RunnableConfig, snapshot: StateSnapshot) -> RunnableConfig: + return { + "configurable": { + **config["configurable"], + "checkpoint_ns": "", + "checkpoint_id": snapshot.config["configurable"]["checkpoint_id"], + } + } + + +def _both(marker: str) -> dict: + return {"log": [marker], "plain": [marker]} + + +def _snapshotted_checkpoints( + checkpointer: BaseCheckpointSaver, config: RunnableConfig +) -> list[str]: + 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: + assert state.values["log"] == state.values["plain"], ( + f"delta channel diverged from the plain channel: " + f"{state.values['log']} != {state.values['plain']}" + ) + assert abandoned not in state.values["log"], ( + f"{abandoned!r} belongs to the branch the fork replaced, " + f"but was replayed into {state.values['log']}" + ) + + +def test_fork_by_invoke( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke( + _both("in-1"), config, durability=durability + ) + graph = _build(sync_checkpointer, "second") + graph.invoke(_both("in-2"), config, durability=durability) + abandoned_head = graph.get_state(config) + + base = next( + snapshot + for snapshot in graph.get_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + _build(sync_checkpointer, "third").invoke( + _both("in-3"), _at(config, base), durability=durability + ) + + state = graph.get_state(config) + _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 +) -> None: + config = _thread("t") + await _build(async_checkpointer, "first").ainvoke( + _both("in-1"), config, durability=durability + ) + graph = _build(async_checkpointer, "second") + await graph.ainvoke(_both("in-2"), config, durability=durability) + abandoned_head = await graph.aget_state(config) + + base = await anext( + snapshot + async for snapshot in graph.aget_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + await _build(async_checkpointer, "third").ainvoke( + _both("in-3"), _at(config, base), durability=durability + ) + + state = await graph.aget_state(config) + _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: + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_both("in-1"), config, durability=durability) + + root = list(graph.get_state_history(config))[-1] + assert root.values["log"] == [] + + _build(sync_checkpointer, "third").invoke( + _both("in-9"), _at(config, root), durability=durability + ) + + state = graph.get_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["in-9", "third-out"] + + +async def test_afork_off_checkpoint_before_first_input( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build(async_checkpointer, "first") + await graph.ainvoke(_both("in-1"), config, durability=durability) + + root = [snapshot async for snapshot in graph.aget_state_history(config)][-1] + assert root.values["log"] == [] + + await _build(async_checkpointer, "third").ainvoke( + _both("in-9"), _at(config, root), durability=durability + ) + + state = await graph.aget_state(config) + _assert_fork_is_clean(state, "in-1") + assert state.values["log"] == ["in-9", "third-out"] + + +def test_fork_by_update_state(sync_checkpointer: BaseCheckpointSaver) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke(_both("in-1"), config) + graph = _build(sync_checkpointer, "second") + graph.invoke(_both("in-2"), config) + + base = next( + snapshot + for snapshot in graph.get_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + forked = graph.update_state(_at(config, base), _both("patched")) + + state = graph.get_state(forked) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "patched"] + + +async def test_afork_by_update_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + await _build(async_checkpointer, "first").ainvoke(_both("in-1"), config) + graph = _build(async_checkpointer, "second") + await graph.ainvoke(_both("in-2"), config) + + base = await anext( + snapshot + async for snapshot in graph.aget_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + forked = await graph.aupdate_state(_at(config, base), _both("patched")) + + state = await graph.aget_state(forked) + _assert_fork_is_clean(state, "in-2") + assert state.values["log"] == [*base.values["log"], "patched"] + + +def test_unaddressed_run_keeps_snapshot_cadence( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_both("in-1"), config, durability=durability) + graph.invoke(_both("in-2"), config, durability=durability) + + 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: + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_both("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: + config = _thread("t") + graph = _build(async_checkpointer, "first") + await graph.ainvoke(_both("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: + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_both("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(_both("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 +) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke(_both("in-1"), config) + graph = _build(sync_checkpointer, "second") + graph.invoke(_both("in-2"), config) + + base = next( + snapshot + for snapshot in graph.get_state_history(config) + if "in-2" not in snapshot.values["log"] + ) + first = ( + StateUpdate(_both("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(_both("second-step"), "n")]], + ) + + state = graph.get_state(forked) + assert state.values["log"] == state.values["plain"], ( + f"delta channel diverged from the plain channel: " + f"{state.values['log']} != {state.values['plain']}" + ) + + +def test_bulk_update_after_a_copy_stores_no_snapshot( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + _build(sync_checkpointer, "first").invoke(_both("in-1"), config) + graph = _build(sync_checkpointer, "second") + graph.invoke(_both("in-2"), config) + base = next( + s for s in graph.get_state_history(config) if "in-2" not in s.values["log"] + ) + + forked = graph.bulk_update_state( + _at(config, base), + [[StateUpdate(None, "__copy__")], [StateUpdate(_both("s2"), "n")]], + ) + + assert graph.get_state(forked).values["log"] == ["in-1", "first-out", "s2"] + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def test_unaddressed_bulk_update_keeps_snapshot_cadence( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + graph = _build(sync_checkpointer, "first") + graph.invoke(_both("in-1"), config) + + graph.bulk_update_state( + config, + [[StateUpdate(_both(f"u{i}"), "n")] for i in range(4)], + ) + + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def _build_paused_before_b(checkpointer: BaseCheckpointSaver) -> Any: + builder = StateGraph(_State) + builder.add_node("a", lambda state: _both("a")) + builder.add_node("b", lambda state: _both("b")) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", END) + return builder.compile(checkpointer=checkpointer, interrupt_before=["b"]) + + +def _build_parallel_interrupt(checkpointer: BaseCheckpointSaver) -> Any: + def ask(state: _State) -> dict: + interrupt("approve?") + return {"other": ["q"]} + + builder = StateGraph(_State) + builder.add_node("p", lambda state: _both("p")) + builder.add_node("q", ask) + builder.add_edge(START, "p") + builder.add_edge(START, "q") + return builder.compile(checkpointer=checkpointer) + + +def test_resume_at_interrupt_before_with_the_head_checkpoint_id_runs_the_node( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build_paused_before_b(sync_checkpointer) + graph.invoke(_both("in"), config, durability=durability) + + graph.invoke(None, graph.get_state(config).config, durability=durability) + + state = graph.get_state(config) + assert state.next == (), f"resume paused again before {state.next}" + assert state.values["log"] == state.values["plain"] == ["in", "a", "b"] + + +async def test_aresume_at_interrupt_before_with_the_head_checkpoint_id_runs_the_node( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build_paused_before_b(async_checkpointer) + await graph.ainvoke(_both("in"), config, durability=durability) + + await graph.ainvoke( + None, (await graph.aget_state(config)).config, durability=durability + ) + + state = await graph.aget_state(config) + assert state.next == (), f"resume paused again before {state.next}" + assert state.values["log"] == state.values["plain"] == ["in", "a", "b"] + + +def test_replay_from_a_paused_checkpoint_runs_the_node_once( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + graph = _build_paused_before_b(sync_checkpointer) + graph.invoke(_both("in"), config) + paused = graph.get_state(config).config + graph.invoke(None, config) + + graph.invoke(None, paused) + + state = graph.get_state(config) + assert state.next == (), f"replay paused again before {state.next}" + assert state.values["log"] == state.values["plain"] == ["in", "a", "b"] + + +@pytest.mark.parametrize("addressed", [False, True]) +def test_new_input_on_an_interrupted_head_does_not_replay_its_pending_writes( + sync_checkpointer: BaseCheckpointSaver, durability: Durability, addressed: bool +) -> None: + config = _thread("t") + graph = _build_parallel_interrupt(sync_checkpointer) + graph.invoke(_both("in-1"), config, durability=durability) + head = graph.get_state(config).config + + graph.invoke(_both("in-2"), head if addressed else config, durability=durability) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == ["in-1", "in-2", "p"] + + +def _build_deferred_after_interrupt(checkpointer: BaseCheckpointSaver) -> Any: + builder = StateGraph(_State) + builder.add_node("a", lambda state: _both("a")) + builder.add_node("b", lambda state: _both("b"), defer=True) + builder.add_node("c", lambda state: {}) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("a", "c") + return builder.compile(checkpointer=checkpointer, interrupt_after=["a"]) + + +@pytest.mark.parametrize( + "durability", + [ + "sync", + "async", + pytest.param( + "exit", + marks=pytest.mark.xfail( + reason="exit durability stores a resumed run's loaded writes twice", + strict=True, + ), + ), + ], +) +def test_resume_on_an_interrupted_head_consumes_its_writes_without_a_snapshot( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build_parallel_interrupt(sync_checkpointer) + graph.invoke(_both("in-1"), config, durability=durability) + + graph.invoke(Command(resume="yes"), config, durability=durability) + + state = graph.get_state(config) + assert state.next == () + assert state.values["log"] == state.values["plain"] == ["in-1", "p"] + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def test_resume_whose_node_writes_the_delta_channel_stores_no_snapshot( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + def ask(state: _State) -> dict: + return _both(interrupt("continue?")) + + graph = ( + StateGraph(_State) + .add_node("ask", ask) + .add_edge(START, "ask") + .compile(checkpointer=sync_checkpointer) + ) + config = _thread("t") + graph.invoke(_both("in"), config, durability=durability) + + graph.invoke(Command(resume="yes"), config, durability=durability) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == ["in", "yes"] + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def test_resume_after_interrupt_before_stores_no_snapshot( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build(sync_checkpointer, "n") + graph.invoke(_both("in"), config, interrupt_before=["n"], durability=durability) + + graph.invoke(None, config, durability=durability) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == ["in", "n-out"] + assert not _snapshotted_checkpoints(sync_checkpointer, config) + + +def _build_send_fan_out(checkpointer: BaseCheckpointSaver) -> Any: + def q(state: _State) -> dict: + interrupt("continue?") + return _both("q") + + builder = StateGraph(_State) + builder.add_node("p", lambda state: _both("p")) + builder.add_node("q", q) + builder.add_conditional_edges( + START, lambda state: [Send("p", state), Send("q", state)], ["p", "q"] + ) + return builder.compile(checkpointer=checkpointer) + + +def test_resume_that_replaces_the_pending_sends_drops_the_finished_task( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + graph = _build_send_fan_out(sync_checkpointer) + config = _thread("t") + graph.invoke(_both("in"), config, durability=durability) + + live = graph.invoke( + Command(resume="yes", goto=[Send("q", _both("unused"))]), + config, + durability=durability, + ) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == live["log"], ( + f"'p' ran in the fan-out the resume replaced, but the reload reads " + f"{state.values['log']} against the live {live['log']}" + ) + + +@NEEDS_CONTEXTVARS +async def test_aresume_that_replaces_the_pending_sends_drops_the_finished_task( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + graph = _build_send_fan_out(async_checkpointer) + config = _thread("t") + await graph.ainvoke(_both("in"), config, durability=durability) + + live = await graph.ainvoke( + Command(resume="yes", goto=[Send("q", _both("unused"))]), + config, + durability=durability, + ) + + state = await graph.aget_state(config) + assert state.values["log"] == state.values["plain"] == live["log"], ( + f"'p' ran in the fan-out the resume replaced, but the reload reads " + f"{state.values['log']} against the live {live['log']}" + ) + + +def test_replay_interrupted_in_its_first_step_still_seals_the_fork( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + graph = _build_send_fan_out(sync_checkpointer) + config = _thread("t") + graph.invoke(_both("in"), config, durability=durability) + + graph.invoke(None, graph.get_state(config).config, durability=durability) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == ["in", "p"], ( + f"the replay reran p, so the fork must not also replay the first p, " + f"but it reads {state.values['log']}" + ) + + +@NEEDS_CONTEXTVARS +async def test_areplay_interrupted_in_its_first_step_still_seals_the_fork( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + graph = _build_send_fan_out(async_checkpointer) + config = _thread("t") + await graph.ainvoke(_both("in"), config, durability=durability) + + await graph.ainvoke( + None, (await graph.aget_state(config)).config, durability=durability + ) + + state = await graph.aget_state(config) + assert state.values["log"] == state.values["plain"] == ["in", "p"], ( + f"the replay reran p, so the fork must not also replay the first p, " + f"but it reads {state.values['log']}" + ) + + +def test_resume_addressed_at_an_interrupted_head_reruns_its_tasks_once( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + config = _thread("t") + graph = _build_parallel_interrupt(sync_checkpointer) + graph.invoke(_both("in-1"), config, durability=durability) + + graph.invoke( + Command(resume="yes"), graph.get_state(config).config, durability=durability + ) + + state = graph.get_state(config) + assert state.next == () + assert state.values["log"] == state.values["plain"] == ["in-1", "p"] + + +def test_update_state_with_the_head_checkpoint_id_keeps_a_deferred_node( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + graph = _build_deferred_after_interrupt(sync_checkpointer) + config = _thread("t") + graph.invoke(_both("in"), config) + + graph.update_state(graph.get_state(config).config, _both("u"), as_node="c") + graph.invoke(None, config) + + state = graph.get_state(config) + assert state.next == (), f"deferred node never ran, still pending: {state.next}" + assert state.values["log"] == state.values["plain"] == ["in", "a", "u", "b"] + + +async def test_aupdate_state_with_the_head_checkpoint_id_keeps_a_deferred_node( + async_checkpointer: BaseCheckpointSaver, +) -> None: + graph = _build_deferred_after_interrupt(async_checkpointer) + config = _thread("t") + await graph.ainvoke(_both("in"), config) + + await graph.aupdate_state( + (await graph.aget_state(config)).config, _both("u"), as_node="c" + ) + await graph.ainvoke(None, config) + + state = await graph.aget_state(config) + assert state.next == (), f"deferred node never ran, still pending: {state.next}" + assert state.values["log"] == state.values["plain"] == ["in", "a", "u", "b"] + + +def test_turns_addressed_at_the_head_store_no_snapshot( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + config = _thread("t") + graph = _build(sync_checkpointer, "turn") + graph.invoke(_both("in-1"), config) + for turn in range(2, 5): + graph.invoke(_both(f"in-{turn}"), graph.get_state(config).config) + + assert not _snapshotted_checkpoints(sync_checkpointer, config) + assert ( + graph.get_state(config).values["log"] == graph.get_state(config).values["plain"] + ) + + +def test_resume_that_reruns_an_error_handler_drops_its_stored_writes( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + handled = threading.Event() + attempts = {"c": 0} + + def fails(state: _State) -> dict: + raise RuntimeError("f always fails") + + def handler(state: _State, error: NodeError) -> dict: + handled.set() + return _both("h") + + def flaky(state: _State) -> dict: + attempts["c"] += 1 + if attempts["c"] == 1: + assert handled.wait(5) + raise ValueError("c fails once") + return _both("c") + + graph = ( + StateGraph(_State) + .add_node("f", fails, error_handler=handler) + .add_node("c", flaky) + .add_edge(START, "f") + .add_edge(START, "c") + .compile(checkpointer=sync_checkpointer) + ) + config = _thread("t") + with pytest.raises(ValueError): + graph.invoke(_both("in"), config, durability=durability) + + live = graph.invoke(None, config, durability=durability) + + state = graph.get_state(config) + assert state.values["log"] == state.values["plain"] == live["log"] diff --git a/libs/langgraph/tests/test_delta_channel_supersteps_bound.py b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py index a2c3a2776..098fe78f4 100644 --- a/libs/langgraph/tests/test_delta_channel_supersteps_bound.py +++ b/libs/langgraph/tests/test_delta_channel_supersteps_bound.py @@ -94,6 +94,28 @@ async def test_forced_snapshot_single_run() -> None: assert "seed-a" in state.values["a"] +async def test_supersteps_bound_skips_a_channel_never_written() -> None: + with patch( + "langgraph.pregel._checkpoint.DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT", + 3, + ): + saver = InMemorySaver() + graph = _build_two_channel_graph(saver, n_loops=4) + config = {"configurable": {"thread_id": "never-written"}} + + graph.invoke({"a": ["seed-a"]}, config) + + minted = [ + t.config["configurable"]["checkpoint_id"] + for t in saver.list(config) + if "b" in t.checkpoint["channel_versions"] + ] + assert not minted, ( + f"b was never written, but {len(minted)} checkpoints minted it a version" + ) + assert graph.get_state(config).values["b"] == [] + + async def test_forced_snapshot_accumulates_across_runs() -> None: """Supersteps counter for an unwritten channel persists across separate invoke() calls. After enough runs, the channel is force-snapshotted.""" @@ -139,13 +161,17 @@ async def test_predicate_fires_on_supersteps_overflow() -> None: channels = {"x": ch_instance} counters: dict[str, tuple[int, int]] = {"x": (0, 5000)} - result = delta_channels_to_snapshot(channels, counters) + result = delta_channels_to_snapshot(channels, counters, {"x": 1}) assert "x" in result counters_below: dict[str, tuple[int, int]] = {"x": (0, 4999)} - result2 = delta_channels_to_snapshot(channels, counters_below) + result2 = delta_channels_to_snapshot(channels, counters_below, {"x": 1}) assert "x" not in result2 + assert not delta_channels_to_snapshot(channels, counters, {}), ( + "a channel with no version was never written, so it has nothing to snapshot" + ) + async def test_counter_reset_after_supersteps_snapshot() -> None: """After the supersteps bound triggers a snapshot, the counters for diff --git a/libs/langgraph/tests/test_delta_channel_update_state.py b/libs/langgraph/tests/test_delta_channel_update_state.py index da6db40f6..62475deeb 100644 --- a/libs/langgraph/tests/test_delta_channel_update_state.py +++ b/libs/langgraph/tests/test_delta_channel_update_state.py @@ -113,6 +113,37 @@ def test_fresh_update_state_head_snapshots_delta_channel() -> None: assert "counters_since_delta_snapshot" not in head.metadata +def test_fresh_update_state_stores_nothing_for_a_delta_channel_it_did_not_write() -> ( + None +): + saver = InMemorySaver() + State = TypedDict( # type: ignore[call-overload] # noqa: UP013 + "State", + { + "messages": Annotated[list, DeltaChannel(_messages_delta_reducer)], + "notes": Annotated[list, DeltaChannel(_messages_delta_reducer)], + }, + ) + graph = ( + StateGraph(State) + .add_node("model", lambda state: {}) + .add_edge(START, "model") + .compile(checkpointer=saver) + ) + config = {"configurable": {"thread_id": "fresh-unwritten"}} + + graph.update_state( + config, + {"messages": [HumanMessage(content="hello", id="m1")]}, + as_node="model", + ) + + head = saver.get_tuple(config) + assert head is not None + assert "notes" not in head.checkpoint["channel_versions"] + assert graph.get_state(config).values["notes"] == [] + + # --------------------------------------------------------------------------- # Non-fresh thread: update_state after invoke # --------------------------------------------------------------------------- @@ -338,3 +369,29 @@ def test_state_history_chain_after_fresh_update_state_delta_channel() -> None: assert update_snapshot.metadata["step"] == 0 assert update_snapshot.parent_config is None assert [m.content for m in update_snapshot.values["messages"]] == ["hello"] + + +def test_update_state_that_snapshots_keeps_a_deferred_node_pending() -> None: + channel = DeltaChannel(_messages_delta_reducer, snapshot_frequency=1) + + class State(TypedDict): + messages: Annotated[list, channel] + + builder = StateGraph(State) + builder.add_node("a", lambda state: {"messages": [HumanMessage("a", id="a")]}) + builder.add_node( + "b", lambda state: {"messages": [HumanMessage("b", id="b")]}, defer=True + ) + builder.add_node("c", lambda state: {}) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("a", "c") + graph = builder.compile(checkpointer=InMemorySaver(), interrupt_after=["a"]) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"messages": [HumanMessage("s", id="s")]}, config) + + graph.update_state(config, {"messages": [HumanMessage("u", id="u")]}, as_node="c") + final = graph.invoke(None, config) + + assert [m.content for m in final["messages"]] == ["s", "a", "u", "b"] + assert graph.get_state(config).next == () diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2e6661e04..bd76c5ea7 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9410,6 +9410,192 @@ def test_fork_does_not_apply_pending_writes( assert result == {"value": 121} +def _extend(state: list, writes: list[list]) -> list: + return [*state, *(v for write in writes for v in write)] + + +class _IntVersionSaver(InMemorySaver): + """Integer versions tie exactly where `InMemorySaver`'s break at random.""" + + get_next_version = BaseCheckpointSaver.get_next_version + + +def _build_chain_after_a_delta_channel() -> Pregel: + return Pregel( + nodes={ + "a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"), + "b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("x"), + "c": NodeBuilder().subscribe_only("x").do(lambda _: "c").write_to("out"), + }, + channels={ + "inp": LastValue(str), + "d": DeltaChannel(_extend, snapshot_frequency=1), + "x": LastValue(str), + "out": LastValue(str), + }, + input_channels=["inp"], + output_channels=["out"], + checkpointer=_IntVersionSaver(), + ) + + +def test_update_state_after_an_exit_snapshot_infers_the_last_writer() -> None: + graph = _build_chain_after_a_delta_channel() + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"inp": "go"}, config, durability="exit") + + graph.update_state(config, "u") + + values = graph.get_state(config).values + assert values["out"] == "u", ( + f"the update should apply as c, the last node to write, but state is {values}" + ) + + +async def test_aupdate_state_after_an_exit_snapshot_infers_the_last_writer() -> None: + graph = _build_chain_after_a_delta_channel() + config = {"configurable": {"thread_id": "t"}} + await graph.ainvoke({"inp": "go"}, config, durability="exit") + + await graph.aupdate_state(config, "u") + + values = (await graph.aget_state(config)).values + assert values["out"] == "u", ( + f"the update should apply as c, the last node to write, but state is {values}" + ) + + +def test_update_state_after_a_fork_seal_infers_the_subscriber_that_ran() -> None: + graph = Pregel( + nodes={ + "a": NodeBuilder() + .subscribe_only("inp") + .do(lambda _: ["a"]) + .write_to("d", go="go"), + "c": NodeBuilder().subscribe_only("go").do(lambda _: ["c"]).write_to("d"), + "b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("out"), + }, + channels={ + "inp": LastValue(str), + "go": EphemeralValue(str), + "d": DeltaChannel(_extend), + "out": LastValue(str), + }, + input_channels=["inp"], + output_channels=["out"], + checkpointer=_IntVersionSaver(), + ) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"inp": "go"}, config) + base = next(s for s in graph.get_state_history(config) if s.metadata["step"] == 0) + graph.invoke(None, graph.update_state(base.config, "u", as_node="b")) + + graph.update_state(config, "w") + + values = graph.get_state(config).values + assert values["out"] == "w", ( + f"the update should apply as b, the last node to run, but state is {values}" + ) + + +def test_update_state_after_a_fork_seal_infers_the_node_whose_read_it_advanced() -> ( + None +): + graph = Pregel( + nodes={ + "a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"), + "b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("out"), + }, + channels={ + "inp": LastValue(str), + "d": DeltaChannel(_extend), + "out": LastValue(str), + }, + input_channels=["inp"], + output_channels=["out"], + checkpointer=_IntVersionSaver(), + ) + config = {"configurable": {"thread_id": "t"}} + graph.invoke({"inp": "1"}, config) + graph.invoke({"inp": "2"}, config) + base = next(s for s in graph.get_state_history(config) if s.next == ("a",)) + fork = graph.update_state(base.config, "u", as_node="b") + + graph.update_state(fork, "w") + + values = graph.get_state(config).values + assert values["out"] == "w", ( + f"the update should apply as b, the last node to read, but state is {values}" + ) + + +def test_update_state_after_a_snapshotting_update_infers_the_last_writer() -> None: + graph = _build_chain_after_a_delta_channel() + config = {"configurable": {"thread_id": "t"}} + graph.update_state(config, ["u"], as_node="a") + graph.invoke(None, config, interrupt_after=["b"]) + + graph.update_state(config, "u") + + values = graph.get_state(config).values + assert values.get("x") == "u", ( + f"the update should apply as b, the last node to write, but state is {values}" + ) + + +async def test_aupdate_state_after_a_snapshotting_update_infers_the_last_writer() -> ( + None +): + graph = _build_chain_after_a_delta_channel() + config = {"configurable": {"thread_id": "t"}} + await graph.aupdate_state(config, ["u"], as_node="a") + await graph.ainvoke(None, config, interrupt_after=["b"]) + + await graph.aupdate_state(config, "u") + + values = (await graph.aget_state(config)).values + assert values.get("x") == "u", ( + f"the update should apply as b, the last node to write, but state is {values}" + ) + + +class _UpdatesOnlyState(TypedDict): + x: Annotated[list[str], operator.add] + + +def _graph_never_run() -> Any: + return ( + StateGraph(_UpdatesOnlyState) + .add_node("a", lambda _: {"x": ["a"]}) + .add_node("b", lambda _: {"x": ["b"]}) + .add_edge(START, "a") + .add_edge("a", "b") + .compile(checkpointer=InMemorySaver()) + ) + + +def test_update_state_on_a_thread_seeded_by_updates_applies_as_input() -> None: + graph = _graph_never_run() + config = {"configurable": {"thread_id": "t"}} + graph.update_state(config, {"x": ["u1"]}) + + graph.update_state(config, {"x": ["u2"]}) + + state = graph.get_state(config) + assert (state.values["x"], state.next) == (["u1", "u2"], ("a",)) + + +async def test_aupdate_state_on_a_thread_seeded_by_updates_applies_as_input() -> None: + graph = _graph_never_run() + config = {"configurable": {"thread_id": "t"}} + await graph.aupdate_state(config, {"x": ["u1"]}) + + await graph.aupdate_state(config, {"x": ["u2"]}) + + state = await graph.aget_state(config) + assert (state.values["x"], state.next) == (["u1", "u2"], ("a",)) + + async def test_delta_channel_end_to_end_inmemory() -> None: """Full graph run: DeltaChannel accumulates correctly across multiple turns."""