diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index de35d0a14..686fca289 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -15,6 +15,7 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, DeltaChannelWrites, + _overwrite_types, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -402,7 +403,7 @@ class AsyncPostgresSaver(BasePostgresSaver): cur: Any, ) -> list[Any]: """Async version of _get_channel_writes_cur — see sync version for rationale.""" - from langgraph.types import Overwrite # type: ignore[import-untyped] + overwrite_types = _overwrite_types() await cur.execute( "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " @@ -435,7 +436,7 @@ class AsyncPostgresSaver(BasePostgresSaver): for type_tag, blob in writes_by_cp.get(cid, []): val = self.serde.loads_typed((type_tag, blob)) collected.append(val) - if isinstance(val, Overwrite): + if isinstance(val, overwrite_types): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 976b36194..5c06d5648 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -14,6 +14,7 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, ChannelVersions, DeltaChannelWrites, + _overwrite_types, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS @@ -235,7 +236,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): 1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread. 2. Walk ancestry in Python, then fetch writes with a plain ANY() filter. """ - from langgraph.types import Overwrite # type: ignore[import-untyped] + overwrite_types = _overwrite_types() cur.execute( "SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints " @@ -268,7 +269,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): for type_tag, blob in writes_by_cp.get(cid, []): val = self.serde.loads_typed((type_tag, blob)) collected.append(val) - if isinstance(val, Overwrite): + if isinstance(val, overwrite_types): collected.reverse() return collected collected.reverse() diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 2ee1452af..a007f0e71 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -55,6 +55,26 @@ def _overwrite_types() -> tuple[type, ...]: return (Overwrite,) +def _split_list_config( + config: RunnableConfig, +) -> tuple[RunnableConfig, RunnableConfig | None]: + """Split a `get_channel_writes` config into `(list_config, before_config)`. + + Most savers collapse `list(config)` to a single row when `config` carries a + `checkpoint_id`. For ancestor traversal we need the opposite: every tuple + strictly earlier than the target. Drop `checkpoint_id` from `list_config` + and pass the original as `before=` (which filters `checkpoint_id < target`). + """ + configurable = config.get("configurable", {}) or {} + target_id = configurable.get("checkpoint_id") + if target_id is None: + return config, None + list_config: RunnableConfig = { + "configurable": {k: v for k, v in configurable.items() if k != "checkpoint_id"} + } + return list_config, config + + logger = logging.getLogger(__name__) @@ -503,14 +523,12 @@ class BaseCheckpointSaver(Generic[V]): if getattr(_DELTA_RECONSTRUCTION, "active", False): return [] overwrite_types = _overwrite_types() + list_config, before_config = _split_list_config(config) _DELTA_RECONSTRUCTION.active = True try: collected: list[Any] = [] # newest first - target_id = config["configurable"].get("checkpoint_id") - for tup in self.list(config): # newest → oldest - if tup.config["configurable"].get("checkpoint_id") == target_id: - continue + 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 @@ -531,17 +549,15 @@ class BaseCheckpointSaver(Generic[V]): self, config: RunnableConfig, channel: str ) -> List[Any]: # noqa: UP006 """Async version of get_channel_writes.""" - overwrite_types = _overwrite_types() - if getattr(_DELTA_RECONSTRUCTION, "active", False): return [] + overwrite_types = _overwrite_types() + list_config, before_config = _split_list_config(config) + _DELTA_RECONSTRUCTION.active = True try: collected: list[Any] = [] - target_id = config["configurable"].get("checkpoint_id") - async for tup in self.alist(config): - if tup.config["configurable"].get("checkpoint_id") == target_id: - continue + 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): diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 6890eeff8..2bc18ccd4 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -390,3 +390,125 @@ class TestInMemorySaverDeltaChannel: } result = saver.get_channel_writes(config, channel) assert result == [{"content": "hi"}, {"content": "bye"}] + + +class TestBaseFallbackGetChannelWrites: + """Exercises the `BaseCheckpointSaver.get_channel_writes` default + implementation — the path third-party savers inherit when they don't + override `get_channel_writes` themselves. + + Regression guard for a bug where the fallback passed the caller's config + (with `checkpoint_id`) straight to `self.list()`, which most savers + collapse to a single row — causing the fallback to return `[]`. + """ + + def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]: + """Build an InMemorySaver with a 3-checkpoint chain and per-step writes + for a `messages` channel. + + Returns `(saver, thread_id, namespace)`. The saver subclass deletes the + InMemorySaver override so the base class fallback is exercised. + """ + + class _ThirdPartyStyleSaver(InMemorySaver): + get_channel_writes = ( + InMemorySaver.__mro__[1].get_channel_writes # type: ignore[attr-defined] + ) + aget_channel_writes = ( + InMemorySaver.__mro__[1].aget_channel_writes # type: ignore[attr-defined] + ) + + saver = _ThirdPartyStyleSaver() + serde = JsonPlusSerializer() + thread_id, ns, channel = "t1", "", "messages" + + cp0 = empty_checkpoint() + cp0["id"] = "00000000000000000000000000000001.0000000000000000" + cp1 = empty_checkpoint() + cp1["id"] = "00000000000000000000000000000002.0000000000000000" + cp2 = empty_checkpoint() + cp2["id"] = "00000000000000000000000000000003.0000000000000000" + saver.storage[thread_id][ns] = { + cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None), + cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]), + cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]), + } + # Writes under cp0 produced cp1's state; writes under cp1 produced cp2's. + saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = ( + "task1", + channel, + serde.dumps_typed({"content": "first"}), + "", + ) + saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = ( + "task2", + channel, + serde.dumps_typed({"content": "second"}), + "", + ) + return saver, thread_id, ns + + def test_fallback_returns_ancestor_writes_oldest_first(self) -> None: + saver, thread_id, ns = self._build_saver_with_chain() + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + result = saver.get_channel_writes(config, "messages") + + assert result == [{"content": "first"}, {"content": "second"}] + + async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None: + saver, thread_id, ns = self._build_saver_with_chain() + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + result = await saver.aget_channel_writes(config, "messages") + + assert result == [{"content": "first"}, {"content": "second"}] + + def test_fallback_stops_at_first_overwrite(self) -> None: + """An `Overwrite` dominates older history: scan newest→oldest stops at + the first one (so `snapshot_every` / user Overwrites bound replay cost). + """ + langgraph_types = pytest.importorskip( + "langgraph.types", reason="langgraph core not installed" + ) + Overwrite = langgraph_types.Overwrite + + saver, thread_id, ns = self._build_saver_with_chain() + serde = JsonPlusSerializer() + cp1_id = "00000000000000000000000000000002.0000000000000000" + # Replace cp1's write with an Overwrite — cp0's write must be dropped. + saver.writes[(thread_id, ns, cp1_id)][("task2", 0)] = ( + "task2", + "messages", + serde.dumps_typed(Overwrite([{"content": "reset"}])), + "", + ) + + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + result = saver.get_channel_writes(config, "messages") + + assert len(result) == 1 + assert isinstance(result[0], Overwrite) + assert result[0].value == [{"content": "reset"}] diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index eb93e7079..6ab6d0914 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -95,6 +95,24 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): new._writes_since_snapshot = self._writes_since_snapshot return new + def _apply_write(self, value: Any, write: Any, counter: int) -> tuple[Any, int]: + """Apply one write to `value`; return (new_value, new_counter). + + An `Overwrite` resets the counter to 0; any other write increments it. + Centralizes the Overwrite/reducer branching used by both `update` (live + super-step) and `from_checkpoint` (ancestor replay). + """ + is_overwrite, overwrite_value = _get_overwrite(write) + if is_overwrite: + new_value = ( + _copy.copy(overwrite_value) + if overwrite_value is not None + else _empty(self.typ) + ) + return new_value, 0 + base = _empty(self.typ) if value is MISSING else value + return self.operator(base, write), counter + 1 + def from_checkpoint(self, checkpoint: Any) -> Self: new: DeltaChannel[Value] = DeltaChannel( self.operator, snapshot_every=self.snapshot_every @@ -106,22 +124,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): new._writes_since_snapshot = 0 elif isinstance(checkpoint, DeltaChannelWrites): # Saver reconstructed per-step writes; replay through the operator. - # Count writes since last Overwrite so the snapshot cadence stays - # accurate across reloads. + # Counter tracks writes since the last Overwrite so snapshot cadence + # stays accurate across reloads. value: Any = _empty(new.typ) counter = 0 for write in checkpoint.writes: - is_overwrite, overwrite_value = _get_overwrite(write) - if is_overwrite: - value = ( - _copy.copy(overwrite_value) - if overwrite_value is not None - else _empty(new.typ) - ) - counter = 0 - else: - value = new.operator(value, write) - counter += 1 + value, counter = new._apply_write(value, write, counter) new.value = value new._writes_since_snapshot = counter else: @@ -135,9 +143,8 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): if not values: return False seen_overwrite = False - applied = 0 for value in values: - is_overwrite, overwrite_value = _get_overwrite(value) + is_overwrite, _ = _get_overwrite(value) if is_overwrite: if seen_overwrite: from langgraph.errors import ( @@ -151,19 +158,13 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE, ) raise InvalidUpdateError(msg) - self.value = ( - _copy.copy(overwrite_value) - if overwrite_value is not None - else _empty(self.typ) - ) seen_overwrite = True - self._writes_since_snapshot = 0 - elif not seen_overwrite: - base = _empty(self.typ) if self.value is MISSING else self.value - self.value = self.operator(base, value) - applied += 1 - if not seen_overwrite: - self._writes_since_snapshot += applied + elif seen_overwrite: + # Post-Overwrite writes within the same super-step are dropped. + continue + self.value, self._writes_since_snapshot = self._apply_write( + self.value, value, self._writes_since_snapshot + ) return True def get(self) -> Any: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 46beb59c3..c440e77f9 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1049,14 +1049,13 @@ class Pregel( step = saved.metadata.get("step", -1) + 1 stop = step + 2 - checkpoint = saved.checkpoint channels, managed = channels_from_checkpoint( self.channels, - checkpoint, + saved.checkpoint, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( - checkpoint, + saved.checkpoint, saved.pending_writes or [], self.nodes, channels, @@ -1169,14 +1168,13 @@ class Pregel( step = saved.metadata.get("step", -1) + 1 stop = step + 2 - checkpoint = saved.checkpoint channels, managed = channels_from_checkpoint( self.channels, - checkpoint, + saved.checkpoint, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( - checkpoint, + saved.checkpoint, saved.pending_writes or [], self.nodes, channels, @@ -1522,8 +1520,9 @@ class Pregel( saved = checkpointer.get_tuple(config) if saved is not None: self._migrate_checkpoint(saved.checkpoint) - base_checkpoint = saved.checkpoint if saved else empty_checkpoint() - checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint + checkpoint = ( + copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() + ) checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} ) @@ -1967,8 +1966,9 @@ class Pregel( saved = await checkpointer.aget_tuple(config) if saved is not None: self._migrate_checkpoint(saved.checkpoint) - base_checkpoint = saved.checkpoint if saved else empty_checkpoint() - checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint + checkpoint = ( + copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() + ) checkpoint_previous_versions = ( saved.checkpoint["channel_versions"].copy() if saved else {} ) diff --git a/pr_delta_channel.md b/pr_delta_channel.md deleted file mode 100644 index 785f7898d..000000000 --- a/pr_delta_channel.md +++ /dev/null @@ -1,133 +0,0 @@ -# feat(channels): DeltaChannel — O(N) incremental checkpoint storage - -## The problem - -LangGraph checkpoints store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, that means each checkpoint blob contains the entire conversation history up to that point. - -Storage cost grows **O(N²)** in the number of turns: - -| Step | Checkpoint blob | -|------|----------------| -| 1 | [msg_1] | -| 2 | [msg_1, msg_2] | -| N | [msg_1, ..., msg_N] | - -At 100K tokens of conversation data, a single thread accumulates ~250 MB; with large messages or file attachments costs scale even faster. - -## The fix: `DeltaChannel` - -`DeltaChannel` is an opt-in wrapper around any binary reducer that stores only a **sentinel marker** in `checkpoint_blobs` rather than the full accumulated value. The actual per-step writes stay in `checkpoint_writes` (which every checkpointer already writes unconditionally). At read time the saver walks the ancestor chain, collects all writes for the channel, and replays them through the reducer. - -Storage scales **O(N)** — the sentinel blob is effectively zero bytes, and the writes table already exists. - -```python -from langgraph.channels.delta import DeltaChannel -from langgraph.graph.message import add_messages - -class State(TypedDict): - # Before: O(N²) storage - messages: Annotated[list[AnyMessage], add_messages] - - # After: O(N) storage - messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)] -``` - -## Benchmarks - -Simulated with realistic paragraph-length messages (~100 tokens each, ~400 chars). Each turn = one human + one AI message (~200 tokens total). - -### Storage (InMemorySaver) - -| turns | ctx | add_msgs | delta | savings | -|------:|----:|---------:|------:|--------:| -| 10 | ~2K tok | 108.6 KB | 4.0 KB | 27x | -| 25 | ~5K tok | 649.0 KB | 10.1 KB | 64x | -| 50 | ~10K tok | 2.6 MB | 20.2 KB | 126x | -| 100 | ~20K tok | 10.2 MB | 40.5 KB | 251x | -| 500 | ~100K tok | 252.6 MB | 202.8 KB | 1245x | - -Savings grow with N because `add_messages` is O(N²) while `DeltaChannel` is O(N). The sentinel blob itself is essentially zero bytes. - -### Read latency (avg of 5 `get_state` calls = cost per `invoke`) - -| turns | ctx | add_msgs | delta | -|------:|----:|---------:|------:| -| 10 | ~2K tok | 0.1ms | 0.2ms | -| 25 | ~5K tok | 0.2ms | 0.5ms | -| 50 | ~10K tok | 0.4ms | 1.5ms | -| 100 | ~20K tok | 0.7ms | 4.8ms | -| 500 | ~100K tok | 5.8ms | 114.9ms | - -**This cost is paid once per `invoke`/`stream` call, not per node.** Within a single invocation, all channels are loaded into memory once at the start and shared across every node — there is no per-node reconstruction. The 114.9ms at 500 turns is what you pay each time a user sends a new message, not on each step of the graph. - -## How it works - -**Write:** `DeltaChannel.checkpoint()` always emits `DeltaChannelSentinel()` — a tiny marker (zero payload bytes) stored in `checkpoint_blobs`. Per-step writes flow into `checkpoint_writes` as they normally do for every channel. - -**Read:** The saver detects `DeltaChannelSentinel` values in `channel_values` and replaces them by calling `get_channel_writes` / `aget_channel_writes`, which walks the ancestor checkpoint chain and collects all writes for that channel (oldest→newest). `DeltaChannel.from_checkpoint()` replays those writes through the operator to reconstruct the full value. - -**Saver implementations:** -- `InMemorySaver` — direct dict traversal of `self.storage` and `self.writes`, no I/O -- `PostgresSaver` (sync + async) — two queries: one cheap ID walk across the thread, one `ANY()` fetch of writes; no recursive CTE -- All other savers — `BaseCheckpointSaver.get_channel_writes` fallback via `list()`, with a re-entrancy guard to prevent infinite recursion - -## Changes - -**`libs/checkpoint`** -- `base/__init__.py` — add `DeltaChannelSentinel` marker dataclass; add `get_channel_writes` / `aget_channel_writes` to `BaseCheckpointSaver` with a `list()`-based fallback and re-entrancy guard - -**`libs/checkpoint/memory`** -- `memory/__init__.py` — `get_channel_writes` via direct dict traversal; `_resolve_delta_channels` helper called in `get_tuple` / `aget_tuple` to replace sentinels with reconstructed write lists - -**`libs/langgraph`** -- `channels/delta.py` — `DeltaChannel` implementation: `checkpoint()` always emits sentinel, `from_checkpoint()` replays writes list -- `channels/__init__.py` — export `DeltaChannel` -- `graph/state.py` — recognize `DeltaChannel` as a valid channel annotation -- `pregel/_checkpoint.py` / `pregel/_loop.py` — wire `after_checkpoint` hook; call it after each checkpointing step so `DeltaChannel` can advance internal state - -**`libs/checkpoint-postgres`** -- `postgres/base.py` — `_get_channel_writes_cur` two-query ancestor walk (sync); `_resolve_delta_channels` called after `_load_blobs` -- `postgres/aio.py` — `_aget_channel_writes_cur` (async counterpart) - -## Open questions - -**Should we add a compile-time capability check?** - -Currently misconfiguring `DeltaChannel` with an unsupported saver only errors at runtime on first reload. A protocol-based check at `compile()` time would give an early warning without requiring a manual boolean flag. - -**`snapshot_every` for bounded reconstruction cost?** - -Both per-invoke read latency and total write wall time grow O(N) per invoke / O(N²) total as the conversation lengthens. A `snapshot_every` parameter — periodically store a full snapshot in `checkpoint_blobs` to cap chain depth — would bound reconstruction cost and is a natural follow-up once the core design is stable. - -## Backwards compatibility - -| Scenario | Behaviour | -|----------|-----------| -| Existing graph using `add_messages` | Unaffected — no code or schema changes | -| `DeltaChannel` loading an old full-list checkpoint blob | Handled via backwards-compat path in `from_checkpoint` | -| `DeltaChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported | -| Time-travel to a past checkpoint | Ancestor walk uses the version at that checkpoint — correct by construction | -| `Overwrite` value | Resets the effective chain; reconstruction starts from that step | - -## Test plan - -- [x] `DeltaChannel` unit tests: `update` → `checkpoint` lifecycle, `from_checkpoint` chain replay, backwards-compat with plain list, `Overwrite` resets chain -- [x] `InMemorySaver` `get_channel_writes`: assembles write list from dict storage -- [x] Serde round-trip for `DeltaChannelSentinel` -- [x] End-to-end graph tests: multi-turn conversations accumulate correctly, time-travel reconstructs correct partial history -- [x] `PostgresSaver` two-query chain reconstruction (sync + async) -- [x] `BaseCheckpointSaver` fallback path via `list()` with re-entrancy guard -- [x] Storage benchmark: `DeltaChannel` uses strictly less storage than `add_messages` at all measured turn counts - ---- - -## Changes from previous base branch - -The previous version stored `DeltaValue` objects (containing the per-step writes) directly in `checkpoint_blobs` and used a `DeltaChainValue` to represent the assembled chain. Reconstruction required a dedicated `get_delta_chain` / `aget_delta_chain` protocol and a recursive CTE in Postgres. - -This version pivots to a simpler design: -- **Sentinel in blobs, writes in `checkpoint_writes`** — `checkpoint_blobs` stores only a zero-byte `DeltaChannelSentinel` marker. The actual per-step data already lives in `checkpoint_writes` (written unconditionally by every checkpointer), so blob storage is essentially free. This is why storage savings jump to 1245x at 500 turns. -- **No custom serde type for the delta payload** — `DeltaValue` / `DeltaChainValue` and the `"delta"` serde type tag are gone. Writes are deserialized with the same serde path they were originally written with. -- **Postgres: two queries instead of a recursive CTE** — fetch all `(checkpoint_id, parent_checkpoint_id)` pairs for the thread, walk the ancestor chain in Python, then fetch writes with a plain `ANY()` filter. -- **Universal fallback on `BaseCheckpointSaver`** — the base class now provides `get_channel_writes` via `list()`, so any third-party saver works without modification. -- **`snapshot_every` removed** — deferred as a follow-up; the simpler design is easier to reason about and delivers larger storage savings.