WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-11 20:27:54 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk

Browse Source 
Four fixes from an independent review of the reconstruction pipeline, plus
a structural cleanup:

1. Ancestor walk excludes the target checkpoint itself (matches pregel:
   writes stored under checkpoint_id=T are pending for the NEXT step and
   applied separately via apply_writes). Memory saver previously included
   them, diverging from Postgres and causing pending writes to be folded
   into the reconstructed snapshot — visible via get_state during
   interrupts and time-travel into a non-leaf checkpoint.

2. Pre-delta blob terminator. When the walk hits an ancestor whose blob
   for the channel is a real value (not DELTA_SENTINEL), bind that blob
   as DeltaChannelWrites.seed and stop. Without this, threads migrated
   from pre-delta storage would replay ancestor writes to the root
   forever AND lose any value that lived only in the old blob
   (e.g. from update_state). Per-ancestor, the blob is checked BEFORE
   its writes — a pre-delta blob subsumes writes at the same checkpoint,
   so including them would double-count.

3. Base-fallback get_channel_writes follows parent_checkpoint_id instead
   of list(before=...). The previous form returned every tuple with
   id<target, including sibling branches on forked threads.

4. seed replaces the Overwrite-wrapping hack for pre-delta values.
   DeltaChannelWrites(writes, seed=SEED_UNSET) makes the saver's
   reconstruction terminator semantically explicit; drops the lazy
   _make_overwrite import dance. User-emitted Overwrite still reset the
   chain via _apply_write as before.

Postgres: recursive CTE enumerates on-path ancestors and joins once
against checkpoint_writes and once against checkpoint_blobs for every
delta channel in the get_tuple — one roundtrip instead of the previous
3 queries × N channels.

Tests added:
- Pre-delta blob seeding (seed binding, no double-counting of ancestor
  writes at the terminator, pending-at-target excluded).
- Root checkpoint returns empty writes.
- Seed-based from_checkpoint replay (three scenarios: with writes,
  seed-only, seed=None distinct from SEED_UNSET).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-30 14:49:05 -04:00
co-authored by Claude Opus 4.7
parent 5f1e946b1a
commit ffacba950a
8 changed files with 684 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+75 -45
View File
@@ -4,10 +4,9 @@ import copy
import logging
import threading
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import ( # noqa: UP035
from typing import (
Any,
Generic,
List,
Literal,
NamedTuple,
TypedDict,
@@ -30,6 +29,9 @@ from langgraph.checkpoint.serde.types import (
SCHEDULED,
ChannelProtocol,
)
from langgraph.checkpoint.serde.types import (
SEED_UNSET as SEED_UNSET,
)
from langgraph.checkpoint.serde.types import (
DeltaChannelWrites as DeltaChannelWrites,
)
@@ -504,71 +506,99 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006
"""Collect writes for `channel` across this checkpoint's ancestry, oldest→newest.
def get_channel_writes(
self, config: RunnableConfig, channel: str
) -> DeltaChannelWrites:
"""Reconstruct a `DeltaChannel`'s write history at this checkpoint.
Scans newest→oldest and stops at the first `Overwrite` (either from
`snapshot_every` or user code), so reconstruction cost is bounded.
Default implementation walks the full thread history via `list()`; savers
can override with a more efficient query (InMemorySaver and PostgresSaver
do this).
Returns a `DeltaChannelWrites` carrying:
`List` is used instead of `list` to avoid mypy confusing it with the
saver's own `list` method.
* `writes` — per-step deltas from ancestors, oldest→newest, ready
to be replayed through the reducer in `DeltaChannel.from_checkpoint`.
* `seed` — when the ancestor walk hits a pre-delta blob (a value
stored before `DeltaChannel` was enabled for this field), replay
starts from that snapshot instead of the channel's empty value.
Default `SEED_UNSET` means no seed.
Walks the **parent chain** (not `list(before=...)`): for a thread with
forks, only on-path ancestors contribute. Scans newest→oldest and
stops at the first `Overwrite`, so reconstruction cost is bounded.
Writes stored at the target `checkpoint_id` itself are pending writes
for the next step and are excluded — pregel applies them separately
via `apply_writes`.
The base implementation uses `get_tuple` and `pending_writes`; it
never sees blobs, so it never sets `seed`. Savers that can read the
blob table directly (`InMemorySaver`, `PostgresSaver`) override this
method to set `seed` when appropriate, which both shortens the walk
and recovers state from pre-delta threads after migration.
"""
# Guard against re-entrant calls: when list() triggers reconstruction
# which calls list() again, the inner call returns tuples with
# DELTA_SENTINEL in channel_values (which get_channel_writes ignores —
# it only reads pending_writes). This breaks the recursion safely.
# Guard against re-entrant calls: when get_tuple() triggers
# reconstruction which calls get_tuple() again, the inner call
# returns tuples with DELTA_SENTINEL in channel_values (which this
# method ignores — it only reads pending_writes).
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
return DeltaChannelWrites(writes=[])
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = [] # newest first
for tup in self.list(list_config, before=before_config): # newest → oldest
if not tup.pending_writes:
continue
# Within a superstep, pending_writes are oldest→newest; reverse
# to scan newest-first.
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return collected
target_tuple = self.get_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = self.get_tuple(cursor_config)
if tup is None:
break
if tup.pending_writes:
# Within a superstep, pending_writes are oldest→newest;
# reverse to scan newest-first.
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return DeltaChannelWrites(writes=collected)
cursor_config = tup.parent_config
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
finally:
_DELTA_RECONSTRUCTION.active = False
async def aget_channel_writes(
self, config: RunnableConfig, channel: str
) -> List[Any]: # noqa: UP006
"""Async version of get_channel_writes."""
) -> DeltaChannelWrites:
"""Async version of `get_channel_writes`. See docstring there."""
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
return DeltaChannelWrites(writes=[])
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = []
async for tup in self.alist(list_config, before=before_config):
if not tup.pending_writes:
continue
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return collected
target_tuple = await self.aget_tuple(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = await self.aget_tuple(cursor_config)
if tup is None:
break
if tup.pending_writes:
for _, ch, value in reversed(tup.pending_writes):
if ch != channel:
continue
collected.append(value)
if isinstance(value, overwrite_types):
collected.reverse()
return DeltaChannelWrites(writes=collected)
cursor_config = tup.parent_config
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
finally:
_DELTA_RECONSTRUCTION.active = False
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+45 -16
View File
@@ -146,22 +146,26 @@ class InMemorySaver(
channel_values: dict[str, Any],
) -> None:
"""Replace DELTA_SENTINEL entries with DeltaChannelWrites so
DeltaChannel.from_checkpoint can distinguish reconstructed writes from
a pre-DeltaChannel accumulated list."""
`DeltaChannel.from_checkpoint` can distinguish reconstructed writes
from a pre-DeltaChannel accumulated value."""
for channel, value in channel_values.items():
if value is DELTA_SENTINEL:
channel_values[channel] = DeltaChannelWrites(
self.get_channel_writes(config, channel)
)
channel_values[channel] = self.get_channel_writes(config, channel)
def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]:
def get_channel_writes(
self, config: RunnableConfig, channel: str
) -> DeltaChannelWrites:
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
# Walk the parent chain newest→oldest collecting checkpoint IDs.
# Walk the parent chain newest→oldest. Skip the target itself —
# writes stored AT `checkpoint_id` are pending for the next step
# (pregel applies them via `apply_writes`; they aren't part of the
# snapshot value AT `checkpoint_id`).
chain: list[str] = []
current: str | None = checkpoint_id
target_entry = ns_storage.get(checkpoint_id)
current: str | None = target_entry[2] if target_entry is not None else None
while current is not None:
entry = ns_storage.get(current)
if entry is None:
@@ -171,14 +175,39 @@ class InMemorySaver(
current = parent
overwrite_types = _overwrite_types()
# Scan writes newest→oldest. Stop at the first `Overwrite` — it
# dominates all older history. Either from `snapshot_every` or from
# user code: the bound applies the same way.
# Scan newest→oldest. Two terminators stop the walk:
# 1. a user-emitted `Overwrite` in writes — replaces prior history;
# 2. a pre-delta blob on an ancestor — bind it as `seed`.
# Without (2), a thread migrated from pre-delta storage would replay
# ancestor writes all the way to the root AND miss any value that
# lived only in the old blob (e.g. from `update_state`).
#
# At each ancestor, check the blob BEFORE processing its pending
# writes: a pre-delta blob represents the state AT that ancestor,
# which already subsumes any writes stored under it. Processing
# those writes first would fold them into the reconstructed value
# twice (once via the blob, once via replay).
collected: list[Any] = [] # newest first
for cp_id in chain: # newest → oldest
entry = ns_storage.get(cp_id)
if entry is not None:
ckpt = self.serde.loads_typed(entry[0])
ver = ckpt.get("channel_versions", {}).get(channel)
if ver is not None:
blob_entry = self.blobs.get(
(thread_id, checkpoint_ns, channel, ver)
)
if blob_entry is not None and blob_entry[0] != "empty":
blob_value = self.serde.loads_typed(blob_entry)
if blob_value is not DELTA_SENTINEL:
# Pre-delta snapshot terminator. Skip this
# ancestor's writes — the blob subsumes them.
collected.reverse()
return DeltaChannelWrites(writes=collected, seed=blob_value)
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
# within a superstep, sorted by (task_id, idx) = oldest → newest,
# so reverse to get newest-first scan.
# Within a superstep, sorted by (task_id, idx) = oldest → newest;
# reverse for newest-first scan.
for (_task_id, _idx), (_, ch, serialized, _) in sorted(
step_writes.items(), reverse=True
):
@@ -188,13 +217,13 @@ class InMemorySaver(
collected.append(val)
if isinstance(val, overwrite_types):
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
collected.reverse()
return collected
return DeltaChannelWrites(writes=collected)
async def aget_channel_writes(
self, config: RunnableConfig, channel: str
) -> list[Any]:
) -> DeltaChannelWrites:
return self.get_channel_writes(config, channel)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
libs/checkpoint/langgraph/checkpoint/serde/types.py
+23
View File
@@ -34,14 +34,37 @@ class _DeltaSentinel:
DELTA_SENTINEL = _DeltaSentinel()
class _SeedUnset:
"""Marker used as the default for `DeltaChannelWrites.seed`.
Distinct from `None`, which is a legitimate pre-delta value
(e.g. an `Optional` field whose accumulated value really was `None`).
"""
__slots__ = ()
def __repr__(self) -> str:
return "SEED_UNSET"
SEED_UNSET = _SeedUnset()
@dataclasses.dataclass
class DeltaChannelWrites:
"""In-memory wrapper around per-step writes reconstructed by a saver.
Consumed by `DeltaChannel.from_checkpoint`. Never serialized — if this
reaches the wire, something upstream forgot to unwrap it.
`seed` is the value from which chain replay should begin. When the saver
encounters a pre-delta blob during the ancestor walk, it uses that blob
as the seed and stops walking further back (the older chain is
represented by the seed). `SEED_UNSET` means "no seed — replay from the
channel's empty value".
"""
writes: list[Any]
seed: Any = SEED_UNSET
Value = TypeVar("Value", covariant=True)
libs/checkpoint/tests/test_memory.py
+165 -20
View File
@@ -6,8 +6,11 @@ from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
SEED_UNSET,
Checkpoint,
CheckpointMetadata,
DeltaChannelWrites,
create_checkpoint,
empty_checkpoint,
)
@@ -208,8 +211,6 @@ class TestMemorySaver:
async def test_memory_saver() -> None:
from langgraph.checkpoint.memory import InMemorySaver
memory_saver = InMemorySaver()
assert isinstance(memory_saver, InMemorySaver)
@@ -325,11 +326,6 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
class TestInMemorySaverDeltaChannel:
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
empty_checkpoint,
)
saver = InMemorySaver()
serde = JsonPlusSerializer()
@@ -349,10 +345,10 @@ class TestInMemorySaverDeltaChannel:
assert channel in result
assert result[channel] is DELTA_SENTINEL
def test_get_channel_writes_collects_writes(self) -> None:
"""get_channel_writes collects per-step writes oldest→newest."""
from langgraph.checkpoint.base import empty_checkpoint
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
"""get_channel_writes collects ancestor writes oldest→newest, and
excludes writes stored at the target checkpoint itself (those are
pending writes for the next step, applied separately by pregel)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
@@ -366,18 +362,20 @@ class TestInMemorySaverDeltaChannel:
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
}
# cp1 has a write for channel
# Writes stored at cp1 produced the cp1 snapshot; part of history.
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "hi"}),
"",
)
# cp2 has a write for channel
# Writes stored at cp2 are pending — they will produce cp3 when the
# step that loaded cp2 completes. They MUST NOT appear in the
# reconstructed snapshot value at cp2.
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed({"content": "bye"}),
serde.dumps_typed({"content": "pending"}),
"",
)
@@ -389,7 +387,37 @@ class TestInMemorySaverDeltaChannel:
}
}
result = saver.get_channel_writes(config, channel)
assert result == [{"content": "hi"}, {"content": "bye"}]
assert result == DeltaChannelWrites(writes=[{"content": "hi"}])
assert result.seed is SEED_UNSET
def test_get_channel_writes_at_root_returns_empty(self) -> None:
"""Reconstructing the root checkpoint's state: no ancestors → []."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
}
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "pending"}),
"",
)
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": "cp1",
}
}
assert saver.get_channel_writes(config, channel) == DeltaChannelWrites(
writes=[]
)
class TestBaseFallbackGetChannelWrites:
@@ -461,7 +489,10 @@ class TestBaseFallbackGetChannelWrites:
result = saver.get_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
assert result == DeltaChannelWrites(
writes=[{"content": "first"}, {"content": "second"}]
)
assert result.seed is SEED_UNSET
async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None:
saver, thread_id, ns = self._build_saver_with_chain()
@@ -476,7 +507,10 @@ class TestBaseFallbackGetChannelWrites:
result = await saver.aget_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
assert result == DeltaChannelWrites(
writes=[{"content": "first"}, {"content": "second"}]
)
assert result.seed is SEED_UNSET
def test_fallback_stops_at_first_overwrite(self) -> None:
"""An `Overwrite` dominates older history: scan newest→oldest stops at
@@ -509,6 +543,117 @@ class TestBaseFallbackGetChannelWrites:
result = saver.get_channel_writes(config, "messages")
assert len(result) == 1
assert isinstance(result[0], Overwrite)
assert result[0].value == [{"content": "reset"}]
assert len(result.writes) == 1
assert isinstance(result.writes[0], Overwrite)
assert result.writes[0].value == [{"content": "reset"}]
assert result.seed is SEED_UNSET
class TestPreDeltaBlobTerminator:
"""Verify the pre-delta blob terminator: when the ancestor walk hits a
checkpoint whose blob for the channel is a real value (not
DELTA_SENTINEL), reconstruction seeds from it and stops. This guards
* back-compat: a thread written by pre-delta code, then extended under
delta — reconstruction must return the correct value without walking
past the last pre-delta ancestor;
* perf: without the terminator, every reconstruct-after-migration would
walk all the way to the thread root.
"""
def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]:
"""Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta,
write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield
seed=[A] + writes=[B, C].
Returns `(saver, thread_id, ns, channel, cp3_id)`.
"""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0"
v2 = "00000000000000000000000000000002.0"
v3 = "00000000000000000000000000000003.0"
# Pre-delta: cp1 stored a real blob for the channel.
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
cp2 = empty_checkpoint()
cp2["id"] = "cp2"
cp2["channel_versions"][channel] = v2
cp3 = empty_checkpoint()
cp3["id"] = "cp3"
cp3["channel_versions"][channel] = v3
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
}
# Write under cp1 would be from the pre-delta era and MUST be ignored
# (the blob already captures it). We add one and assert it is not
# folded into the reconstructed result.
saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = (
"task0",
channel,
serde.dumps_typed("PRE-DELTA-WRITE"),
"",
)
saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed("B"),
"",
)
saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = (
"task3",
channel,
serde.dumps_typed("PENDING-AT-TARGET"),
"",
)
return saver, thread_id, ns, channel, "cp3"
def test_seed_from_pre_delta_ancestor_blob(self) -> None:
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver.get_channel_writes(config, channel)
# Seed came from the pre-delta blob at cp1.
assert result.seed == ["A"]
# Delta-era writes from cp2 replay through the reducer on top of seed.
# cp3 is the target — its own write is pending for the NEXT step and
# must be excluded.
assert result.writes == ["B"]
def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None:
"""Writes stored at the pre-delta ancestor itself must not be replayed
(the blob subsumes them)."""
saver, thread_id, ns, channel, target = self._build_mixed_thread()
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target,
}
}
result = saver.get_channel_writes(config, channel)
# The pre-delta write under cp1 must not appear (the blob subsumes it).
assert "PRE-DELTA-WRITE" not in result.writes
# And the pending write at the target is never folded in.
assert "PENDING-AT-TARGET" not in result.writes