WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-19 06:23:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(langgraph): don't replay an abandoned branch into a DeltaChannel fork (#8548)

Browse Source 
Fixes #8443
Fixes #9089

A checkpoint keeps the pending writes that produced its child, and
nothing records which child consumed them. When a new branch starts from
a checkpoint that already has pending writes (going back in time, or new
input on an interrupted head), the `DeltaChannel` ancestor walk replays
those writes into the new branch too. The live run is correct; only a
reload is wrong:

```
fork base:     ['in-1', 'first-out']
fork returns:  ['in-1', 'first-out', 'in-3', 'third-out']
reload gives:  ['in-1', 'first-out', 'in-2', 'in-3', 'third-out']
                                     ^^^^^^ from the branch the fork replaced
```

Plain channels store their full value and are unaffected, so the tests
use one as the oracle.

## Fix

The first checkpoint of a new branch snapshots the delta channels its
base has pending writes for, so the walk stops inside the branch. Only
the base's own writes are branch-specific; everything above it is shared
history. A base with no pending writes has nothing to leak, so an
ordinary turn that addresses the head (as clients commonly do) stores
nothing. `bulk_update_state` takes the set from its first superstep
only: a `__copy__` is stored under the base's parent, so nothing after
it walks the base's writes. A resume that is not replaying reuses the
head's pending writes instead of rerunning their tasks, so it seals only
the loaded writes that don't go back to their task: a finished task
whose `Send` a `Command(goto=...)` replaced, or an error handler that
runs again. A plain resume stores nothing. A resume addressed by
`checkpoint_id` reruns them, so it still seals.

`put` only stores a blob for a channel whose version moved since the
last stored checkpoint, so a snapshot of one that didn't move needs a
version bump, and scheduling reads versions. `create_checkpoint`
therefore advances every `versions_seen` entry that had seen the old
version, including the interrupt tracker. Without the advance, the bump
re-fires `interrupt_before` on resume and reruns the channel's
subscribers. For each entry it advances, `SNAPSHOT_BUMPS` keeps the
version the node really read, so `update_state`'s `as_node` inference
reads `versions_seen` as if the bump never happened. A never-written
channel gets a version only for the seal; the cadence and a fresh
thread's first `update_state` skip it.

`update_state` no longer records its narrower `updated_channels` when it
snapshots; it skipped a deferred node listed in `next` on resume
(#9089).

The same seal fixes two `update_state` calls on one checkpoint (editing
the same message twice): both store their writes there under the same
task id, the saver keeps the first, and the second branch read back the
first one's edit.

Two things this touches were also wrong on `main`: a resumed error
handler that runs again left its stored writes on the head (an exit
reload read them twice), and `aupdate_state` on a thread seeded only by
updates raised "Ambiguous update" where `update_state` applied the
update as the input. `update_state` and `aupdate_state` now share one
`as_node` inference.

Exit durability has a separate replay bug on `main` when a resumed
checkpoint already holds writes (duplicated or reordered replay),
unrelated to forks. It's fixed in #9114; the resume test here marks exit
durability as a strict expected failure until then.

`tests/memory_assert.py` now compares against the checkpoint as read
back: a delta channel a step didn't write is refilled on read, which the
old comparison reported as a mutation.

Cost: 300 turns addressing the head store no snapshots, as on `main`. A
resume that reruns finished tasks seals every time. After a parallel
task finished, 30 turns of resuming with the head's `checkpoint_id`
(what Studio sends) stored 30 snapshots, 191 KB, against 12 KB of delta
writes, and a subgraph resume with a finished sibling does the same,
since a subgraph loop always counts as replaying. That seal is what
keeps a rerun task's new write from being replayed as its old one:
without it, a subgraph task that returns something different on the
rerun reads back its first result. The reruns happen on `main` too, and
stopping them would remove this cost.

276 of 464 cases in `test_delta_channel_fork.py` fail on `main` and pass
here (memory, sqlite and postgres, all durabilities). #9089's own case
is in `test_delta_channel_update_state.py`, the cadence case in
`test_delta_channel_supersteps_bound.py`, and the `as_node` cases in
`test_pregel.py`.

## Limits

- Threads forked before this change keep their state: the ownership was
never recorded, so there is nothing to recover.
- With exit durability, a fork at a finished turn stores its writes on
the shared base, so the original branch then replays them too (`['h1',
'ai', 'h2-edited', 'ai', 'h2', 'ai']`). Same on `main`.
- `Command(update=..., goto=...)` sent to an old checkpoint stores the
update there, so the original branch replays it too. The fork itself is
correct now; the original branch is the same as on `main`.
- #8551 (the mirror case: `update_state`'s own writes leaking into the
abandoned branch) is fixed in #9165, stacked on this PR. It builds on
this snapshot, but keys off whether the addressed checkpoint is the
thread's latest rather than on pending writes, which a finished turn
that a later run continued from doesn't have.

Thanks to @AnnaSuSu for the report, the reproduction and the snapshot
approach, and to @UditDewan for the implementation in #8476. Both are
co-authors.

---------

Co-authored-by: AnnaSuSu <64579968+AnnaSuSu@users.noreply.github.com>
Co-authored-by: UditDewan <194863456+UditDewan@users.noreply.github.com>
This commit is contained in:
Elior Nataf Lackritz
2026-10-03 08:55:59 -04:00
committed by GitHub
co-authored by AnnaSuSu UditDewan
parent 7dc9195e41
commit 9a0394d88b
9 changed files with 1309 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/_internal/_constants.py
+4
View File
@@ -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,
libs/langgraph/langgraph/pregel/_checkpoint.py
+127 -30
View File
@@ -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.
libs/langgraph/langgraph/pregel/_loop.py
+68 -17
View File
@@ -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
)
libs/langgraph/langgraph/pregel/main.py
+103 -62
View File
@@ -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(
libs/langgraph/tests/memory_assert.py
+5 -3
View File
@@ -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):
libs/langgraph/tests/test_delta_channel_fork.py
+731
View File
@@ -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"]
libs/langgraph/tests/test_delta_channel_supersteps_bound.py
+28 -2
View File
@@ -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
libs/langgraph/tests/test_delta_channel_update_state.py
+57
View File
@@ -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 == ()
libs/langgraph/tests/test_pregel.py
+186
View File
@@ -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."""