Compare commits

...
Author SHA1 Message Date
Sydney RunkleandClaude Opus 4.7 18cbe46baf feat(prebuilt): hydrate ToolNode state from channels via CONFIG_KEY_READ
When ToolNode receives a bare `[tool_call]` list via the Send API (the
new dispatch shape that create_agent uses after langchain-ai/langchain
drops the ToolCallWithContext wrapper), hydrate ToolRuntime.state from
the current channel values instead of requiring the dispatcher to
inline the full agent state dict in the Send payload.

Implementation stays entirely in tool_node:

- Pregel installs CONFIG_KEY_READ as
  `functools.partial(local_read, scratchpad, channels, managed, task)`.
  Introspect the partial's positional args to learn channel + managed
  names, then read them all via `ChannelRead.do_read` with an explicit
  list. No changes to the pregel read machinery.
- Gracefully falls back to {} when invoked outside a Pregel context
  (e.g. direct ToolNode(...).invoke(...) from test harnesses).
- Legacy ToolCallWithContext path is preserved for external dispatchers.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 07:30:51 -04:00
Sydney RunkleandClaude Opus 4.7 e7af9869bb refactor(add_messages): clearer step-by-step structure
Restructure add_messages so fast vs. slow path, REMOVE_ALL handling,
and format application each live in a single numbered section with a
short lead-in comment. No behavior changes; the previous commits'
optimizations are preserved.

- fold the two path branches into one `if pure_append else slow_path`
  so format handling happens at a single exit instead of being
  duplicated between fast and slow paths
- drop the now-redundant `left_seq` local; index `left` directly after
  the coerce step (with a single `cast(list, left)` for the type
  checker)
- use `set.isdisjoint` for the overlap check

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 07:30:51 -04:00
Sydney RunkleandClaude Opus 4.7 e5bfae0f9d test(add_messages): cover fast-path guards and format handling
Nine new tests pin the behavioral boundaries introduced by the
optimization: chunk / dict / tuple / missing-id left inputs must fall
through to full conversion, duplicate right ids and None right ids
must still be handled correctly, format="langchain-openai" and invalid
format must round-trip through the fast path, and the fast path must
return a fresh list rather than aliasing left.

Also picks up a ruff-format reflow in test_time_travel.py.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 07:30:51 -04:00
Sydney RunkleandClaude Sonnet 4.6 3ef54c1ec8 fix(add_messages): guard fast path against None IDs and intra-right duplicates
Two bugs in the fast-path optimisation:

1. The left-side type guard checked isinstance(BaseMessage) but not
   id is not None. Messages without IDs (e.g. HumanMessage(content="hi"))
   would skip ID assignment and return None IDs.

2. The pure-append short-circuit only checked for overlaps between
   right and left, not duplicates within right itself. A right list
   containing two messages with the same ID would bypass the slow-path
   deduplication and return both.

Fixes:
- Add left_seq[0].id is not None to the type-guard condition.
- Replace the any() overlap check with a set-intersection check that
  also verifies len(right_id_set) == len(right_msgs) (no intra-right
  duplicates) before taking the fast return.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-24 07:30:51 -04:00
Sydney RunkleandClaude Sonnet 4.6 2a974d1b73 perf(add_messages): skip left-side conversion and fast-path pure appends
Two optimizations for the hot path in add_messages, which is called on
every write to a messages channel:

1. Skip conversion of left: when left is already list[BaseMessage] with
   IDs assigned (true for every call after the first), skip
   convert_to_messages + message_chunk_to_message + the ID-None loop.
   These are O(n) no-ops on already-resolved messages that allocate two
   intermediate lists.

2. Pure-append short-circuit: when right contains no RemoveMessage and
   no ID overlaps with left, return left + right directly. Replaces the
   O(n) copy + dict build + filter with a single set-membership check.

Benchmarks (median of 2000 iterations, pure-append scenario):
  10-msg thread:   2.9x faster
  100-msg thread:  6.6x faster
  1000-msg thread: 7.3x faster
  200-step simulation (2 msgs/step): 3.4x faster end-to-end

Also adds tests/test_add_messages_benchmark.py with correctness tests
for all scenarios (append, update, remove) and a runnable benchmark.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-24 07:30:51 -04:00
Sydney Runkle afec98f369 internal for now 2026-04-24 07:29:52 -04:00
Sydney RunkleandClaude Opus 4.7 f25d1935ef fix(postgres): handle missing checkpoint_id in _get_channel_writes_history; update test signatures
Two fixes exposed by running the postgres test suite against a local
postgres instance:

1. `PostgresSaver._get_channel_writes_history` /
   `AsyncPostgresSaver._aget_channel_writes_history` required
   `checkpoint_id` in the passed config, raising `KeyError` when called
   with just `thread_id` (e.g. `graph.aget_state({"thread_id": "..."})`).
   Now resolves to the latest checkpoint via `get_tuple`/`aget_tuple`
   when the id is missing.

2. `test_get_checkpoint_no_channel_values` (sync + async) monkeypatched
   `_load_checkpoint_tuple` with the old `(value, cur)` signature. Method
   now takes `(value)` only since delta reconstruction moved out of the
   tuple-load path — updated both tests.

Local postgres (`brew install pgvector postgresql@16`, running on port
5441) now exercises all 40 non-vector postgres tests green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:47:22 -04:00
Sydney RunkleandClaude Opus 4.7 3a7ed5b454 refactor(delta-channel): honest data model, private experimental API
Restructure DeltaChannel reconstruction so the hydration path matches
pregel's storage axes (blobs + writes) without leaking internal DTOs
into the public checkpoint contract.

Key changes:

* Deleted `DeltaChannelWrites` dataclass and `SEED_UNSET` sentinel.
  Reconstruction data no longer flows through `Checkpoint.channel_values`
  as a wrapped DTO — that field now carries a value or `DELTA_SENTINEL`,
  never a reconstruction shape.
* Added private `_ChannelWritesHistory(seed: Any, writes: list[PendingWrite])`
  NamedTuple as the return type for the new storage-level query.
* Added private, experimental `_get_channel_writes_history` /
  `_aget_channel_writes_history` on `BaseCheckpointSaver` — reference
  impl via `get_tuple` + `parent_config` walk, overridden on
  `InMemorySaver` / `PostgresSaver` / `AsyncPostgresSaver` for perf.
  Fixes a latent migration bug in the base fallback (now inspects
  ancestor `channel_values` for pre-delta seed).
* `DeltaChannel.from_checkpoint(seed)` simplified to two cases
  (sentinel/MISSING → empty, else → seed). New `replay_writes` method
  folds `list[PendingWrite]` through the reducer.
* Delta hydration consolidated inside `channels_from_checkpoint` via
  optional `saver` + `config` kwargs (+ async mirror
  `achannels_from_checkpoint`). All six pregel call sites updated.
  `get_tuple` no longer patches `channel_values` — removed
  `_resolve_delta_channels` (memory) and per-tuple reconstruction from
  `_load_checkpoint_tuple` (postgres sync + async).
* Hydration short-circuits on the target's own blob: if
  `channel_values[k]` is a real value (pre-migration tip, `update_state`
  result), use it directly. Only walks ancestors when the target holds
  sentinel or is missing. Fixes a correctness bug where migration-tip
  and `update_state` values would be lost.
* New test_delta_channel_migration.py: 10 scenarios covering
  BinaryOperatorAggregate → DeltaChannel migration (basic + async,
  time-travel, fork, `update_state`, tip-of-pre-migration, base-saver
  fallback parity, cross-thread isolation).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:15:02 -04:00
Sydney Runkle 31ef0e942a refactor(delta-channel): drop snapshot_every and saver Overwrite terminator
snapshot_every was a knob for bounding reconstruction cost on deep threads.
Benchmarks (notes/add_messages_replay_problem.md + scratch work on
sr/add-messages-replay-bench) showed the add_messages fast-path
(optimize/add-messages-fast-path) closes the quadratic replay cost for
threads under ~1000 turns, where the crossover to snapshots makes sense.
For deeper threads we'll ship a first-class compaction primitive instead.

Removals:

* DeltaChannel: snapshot_every ctor param, _writes_since_snapshot counter,
  should_snapshot() / snapshot_write() methods, counter threading through
  _apply_write / update / from_checkpoint / copy.
* Pregel loop: post-checkpoint snapshot-injection block and
  SNAPSHOT_TASK_ID import + constant.
* Checkpoint base: _overwrite_types() helper and the ancestor-walk
  short-circuit on user-emitted Overwrite in sync + async
  get_channel_writes.
* InMemory + Postgres savers: same walk-terminator shortcut. The
  pre-delta blob terminator (seed-from-ancestor-blob) stays — it's
  required for migration correctness, not a snapshot optimization.
* Tests for all of the above.

Preserved:

* Channel-level Overwrite semantics in DeltaChannel / BinOpAggregate:
  Overwrite still resets the value at reducer level; same-super-step
  dedup and InvalidUpdateError on multiple Overwrites still enforced.
* Pre-delta migration seeding.
2026-04-23 09:54:12 -04:00
Sydney Runkle d120f127ca refactor(delta-channel): plain SELECT WHERE replaces recursive CTE
The recursive CTE was bottlenecked by a JSON-expression join
(`bl.version = checkpoint->'channel_versions'->>bl.channel`) that the
planner could not index, producing an O(ancestors x blobs) nested-loop.
At depth 1000 it ran ~275 ms and removed ~2M filter rows; the recursion
itself was 2.4 ms.

Switch to three plain indexed SELECTs per delta channel
(checkpoints, checkpoint_writes, checkpoint_blobs); a pure helper on
BasePostgresSaver walks the parent chain and assembles
DeltaChannelWrites. Sync (__init__.py) and async (aio.py) each own
their three-roundtrip I/O wrappers.

Bench numbers (notes/delta_channel_query_bench.md): 3x at depth 50,
15x at depth 200, ~100x at depth 1000. Plain over-fetches sibling rows
when the thread branches but still wins at every realistic depth on
both local and remote postgres.

Multi-channel coalescing dropped — reconstruction is per-channel now.
Same shape as InMemorySaver. Can come back as a SQL-level
optimization later if needed.
2026-04-23 08:21:00 -04:00
Sydney Runkle 9e330c96dc contextvar 2026-04-23 07:42:51 -04:00
Sydney RunkleandClaude Opus 4.7 cd8fad5905 fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk
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>
2026-04-22 22:09:08 -04:00
Sydney Runkle acc7eda8c5 optimizations i sure hope 2026-04-22 20:47:51 -04:00
Sydney Runkle 5b7fdf5655 eh 2026-04-22 18:41:06 -04:00
Sydney RunkleandClaude Opus 4.7 4cad68f767 fix(delta-channel): unwrap NotRequired[X] for dict/set reducers
Annotated[NotRequired[dict[...]], DeltaChannel(reducer)] (the shape used
by deepagents' filesystem middleware) fell through type inference to
`list`, so the first operator call blew up with
"'list' object is not a mapping". `_is_field_channel` now unwraps a
parameterized Required[X]/NotRequired[X] before stripping extras, which
lets dict/set/mapping outer types reach the abc normalization block.

Also type-annotates the `new` locals in DeltaChannel.copy() and
from_checkpoint() so mypy can infer them through the abstract return
type.

Adds tests covering: dict Overwrite in update and in writes replay,
snapshot_write with a dict reducer, dict backwards-compat checkpoints,
NotRequired type inference, and a filesystem-shaped end-to-end graph.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-22 17:58:36 -04:00
Sydney Runkle 9d8c0be068 arbitrary 2026-04-22 17:54:12 -04:00
Sydney Runkle 51154be4ab lint 2026-04-22 17:47:53 -04:00
Sydney Runkle 2e7edb2b60 lint 2026-04-22 17:30:22 -04:00
Sydney Runkle 325cb42f19 lint and snapshot every 2026-04-22 17:06:53 -04:00
Sydney Runkle b9fad696ec cleanup 2026-04-22 16:45:32 -04:00
Sydney Runkle 96760e6267 chore(delta-channel): add PR description 2026-04-22 14:35:48 -04:00
Sydney RunkleandClaude Sonnet 4.6 9342ae215a chore(delta-channel): remove snapshot_every — simpler design, better storage savings
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:12:57 -04:00
Sydney Runkle ee5b3c2639 refactor(postgres): replace recursive CTE with two-query ancestor walk for DeltaChannel
Instead of a recursive SQL CTE, collect the ancestor checkpoint ID chain in
Python by fetching all (checkpoint_id, parent_checkpoint_id) for the thread
in one query, then fetch writes with a plain WHERE checkpoint_id = ANY(...).

Simpler, avoids recursive query planner overhead, and uses well-indexed lookups.
2026-04-22 14:03:37 -04:00
Sydney Runkle d7c3616620 feat(delta-channel): store sentinel in blobs, reconstruct from checkpoint_writes
DeltaChannel.checkpoint() now returns a zero-byte DeltaChannelSentinel
instead of duplicating delta data in checkpoint_blobs. Reconstruction
walks the parent checkpoint chain via checkpoint_writes (which already
holds per-step writes) and replays them through the operator.

In-memory benchmark (100 turns, ~20K tokens):
  storage: 10.2 MB → 40.5 KB (251x reduction)
  read:    0.6ms → 7.9ms (reconstruction cost, amortized by storage savings)

InMemorySaver and PostgresSaver override get_channel_writes() with
efficient implementations (Python dict walk and recursive CTE respectively).
The base class fallback uses self.list() with a thread-local recursion guard.
2026-04-22 14:03:37 -04:00
Sydney Runkle 06e302bbff refactor(delta-channel): infer typ from Annotated outer type instead of constructor arg
Remove the positional `typ` parameter from `DeltaChannel.__init__`. The type is
now injected automatically from the `Annotated` outer type in `_is_field_channel`
(matching how `BinaryOperatorAggregate` receives its type). `copy()` and
`from_checkpoint()` propagate `self.typ` explicitly. Test helpers updated to
use `_get_channel` with the proper `Annotated` path.
2026-04-22 14:03:37 -04:00
Sydney Runkle 10da2326a9 chore(delta-channel): remove supports_delta_channels flag
Rely on the runtime raise in DeltaChannel.from_checkpoint() instead of
a compile-time boolean flag. Savers that assemble DeltaChainValue inside
_load_blobs work transparently; savers that don't will pass through a raw
DeltaValue and hit a clear ValueError on first reload.

Removes: BaseCheckpointSaver.supports_delta_channels, the attribute on
InMemorySaver / PostgresSaver / AsyncPostgresSaver, the compile-time
UserWarning in StateGraph.compile(), and the associated test.
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 84f00c51eb chore: remove docs/ from PR
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
ccurmeandSydney Runkle 3cab106bdf fix(prebuilt): handle injected NotRequired keys (#7392)
Resolves https://github.com/langchain-ai/langchain/issues/35585

This would previously raise KeyError:
```python
from typing import Annotated

from langchain_core.tools import tool
from langchain.agents import create_agent
from typing_extensions import NotRequired
from langgraph.prebuilt import InjectedState
from langchain.agents import AgentState


class CustomAgentState(AgentState):
    city: NotRequired[str]


@tool
def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> str:
    """Get weather for a given city."""
    if city is None:
        city = "Boston"
    return f"It's always sunny in {city}!"


agent = create_agent(
    model="claude-sonnet-4-6",
    tools=[get_weather],
    system_prompt="You are a helpful assistant",
    state_schema=CustomAgentState,
)

input_message = {
    "role": "user",
    "content": "What's the weather?",
}

result = agent.invoke({"messages": [input_message]})
for m in result["messages"]:
    m.pretty_print()
```

---------

Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle 77bb349309 lint 2026-04-22 14:03:37 -04:00
Sydney Runkle 1d8364a749 latest 2026-04-22 14:03:37 -04:00
Sydney Runkle 06e97b0fd2 fix(delta-channel): support non-list reducers (dict) and fix MISSING handling
Use typ() instead of [] throughout DeltaChannel so reducers over dict
(and other non-list types) work correctly. fromCheckpoint(MISSING) now
leaves value as typ() from __init__ instead of overwriting with MISSING.
copy() uses value.copy() to handle dicts. update() initialises base from
typ() when value is MISSING. Add four tests covering the deepagents-style
dict-merge / file-deletion reducer pattern.
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 e256a31d00 chore(delta): rename _steps_since_rehydrate → _steps_since_snapshot; add audit tests
- Rename `_steps_since_rehydrate` → `_steps_since_snapshot` in DeltaChannel
  for clarity (counts steps since the last snapshot, not since rehydration)
- Pre-seed cycle-detection `visited` set with current checkpoint ID in both
  sync and async `_assemble_delta_channels` to prevent self-referential chains
- Add 4 new unit tests:
  - `test_delta_channel_snapshot_every_emits_plain_list`: verifies counter
    semantics and snapshot/delta transitions
  - `test_delta_channel_snapshot_every_end_to_end`: graph-level smoke test
  - `test_delta_channel_assembly_fast_path_returns_delta_value`: exercises
    chain traversal via get_channel_blob returning DeltaValue then plain list
  - `test_delta_channel_assembly_broken_chain_logs_warning`: partial chain
    when get_tuple returns None

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 04b3ae7cd0 chore: rename serde type tag "diff" → "delta" for DeltaValue
Consistent with channel/type naming (DeltaChannel, DeltaValue).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ec52520389 chore: apply format/lint fixes across checkpoint, checkpoint-postgres, prebuilt
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 82fea763c5 fix: register DeltaValue in SAFE_MSGPACK_TYPES; rename _is_diff_delta; cross-saver benchmark
- Add DeltaValue to SAFE_MSGPACK_TYPES so SQLite and other msgpack-based
  savers don't emit "Deserializing unregistered type" warnings.
- Rename _is_diff_delta → _is_delta_value (leftover from DiffChannel rename).
- Parametrize benchmark by checkpointer: runs InMemory (fast-path) and
  SQLite (get_tuple fallback) in the same table, sharing the _run_turns helper.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 6cfcad18f3 fix(delta-channel): fix chain assembly and get_state paths
- Fix InMemorySaver.get_channel_blob: use correct storage[thread_id][ns]
  nesting and deserialize the checkpoint before extracting channel_versions.
- Pass checkpoint_id to after_checkpoint() in channels_from_checkpoint so
  DeltaChannel seeds _last_checkpoint_id correctly on load; without this
  every turn broke the chain at its boundary.
- Wire _assemble_delta_channels into _prepare_state_snapshot and
  _aprepare_state_snapshot (get_state / get_state_history paths) and into
  perform_superstep / aperform_superstep (update_state paths) — previously
  only the loop __enter__ path did assembly.
- Fix test_get_channel_blob to use the correct storage structure.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle db6b9ec995 test(channels): replace unsupported-saver raise test with fallback assembly test 2026-04-22 14:03:37 -04:00
Sydney Runkle 55fdc7aec6 feat(postgres): remove _load_diff_chains; add get_channel_blob / aget_channel_blob 2026-04-22 14:03:37 -04:00
Sydney Runkle 9969fb9737 feat(memory): implement get_channel_blob; remove diff handling from _load_blobs 2026-04-22 14:03:37 -04:00
Sydney Runkle 40981fdac7 feat(pregel): wire DeltaChannel assembly into loop; pass checkpoint_id to after_checkpoint 2026-04-22 14:03:37 -04:00
Sydney Runkle 5d1b3c4190 feat(pregel): add _assemble_delta_channels helpers for universal DeltaChannel support 2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 fa83d64eff feat(channels): DeltaChannel tracks checkpoint_id; emits prev_checkpoint_id in DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 4608af9615 feat(serde): diff type encodes prev_checkpoint_id; loads_typed returns DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle 9beda5d3fb docs(checkpoint): expand aget_channel_blob docstring for parity 2026-04-22 14:03:37 -04:00
Sydney Runkle 9a6d7e08fb chore: add .worktrees/ to .gitignore 2026-04-22 14:03:37 -04:00
Sydney Runkle bbeb2759ba feat(checkpoint): DeltaValue uses prev_checkpoint_id; add get_channel_blob stubs 2026-04-22 14:03:37 -04:00
Sydney Runkle b799b95138 chore: rename DiffChannel/DiffDelta/DiffChainValue to Delta* across libs
Renames the diff-channel types to DeltaChannel, DeltaValue, and DeltaChainValue
for consistency with the settled naming convention.
2026-04-22 14:03:37 -04:00
Sydney Runkle ed9711fd33 more tests 2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ca883fe4eb chore: format/lint fixes for rehydrate_every benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 4b9e4d25ca feat(channels): add rehydrate_every to DiffChannel for bounded chain traversal
Periodic full-snapshot checkpoints cap chain depth, trading a small
amount of extra storage for bounded reconstruction time.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ec8fd85ea2 test(channels): add DiffChannel vs BinaryOperatorAggregate storage/time benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 ebd98f2e27 chore: format and lint fixes for DiffChannel implementation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 318fee9fc6 fix(checkpoint/postgres): pass cursor to avoid deadlock in diff chain traversal
Fixes a critical deadlock that occurs when _load_diff_chains calls self._cursor()
from within _load_blobs while the outer _load_checkpoint_tuple already holds
self._cursor(). On bare (non-pool) connections, the threading.Lock is not
reentrant, causing a deadlock.

Solution: Pass the cursor as a parameter to _load_diff_chains and _load_blobs
instead of acquiring a new cursor within those methods. Updated _load_checkpoint_tuple
to acquire a cursor once at the top level and pass it through the call chain.

Changes:
- Updated _load_blobs signature to accept optional cur parameter
- Updated _load_diff_chains signature (base and implementations) to accept optional cur parameter
- Modified _load_checkpoint_tuple in PostgresSaver to acquire cursor and pass it
- Modified _load_checkpoint_tuple_async to acquire cursor only when diff_payloads exist
- Removed nested self._cursor() calls in _load_diff_chains and _load_diff_chains_async

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e37299af87 feat(checkpoint/postgres): diff chain reconstruction in async saver
Add `_load_diff_chains_async` to `AsyncPostgresSaver` and override
`_load_checkpoint_tuple` to inline blob-parsing and diff-chain
resolution via async point-lookup traversal, mirroring the sync
`PostgresSaver._load_diff_chains` implementation. Add integration test
`test_diff_channel_chain_reconstruction` that skips gracefully when
`langgraph` core is not installed.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 65d6ab2609 feat(checkpoint/postgres): diff chain reconstruction in _load_blobs (sync)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 888a308814 test(pregel): strengthen DiffChannel time-travel and reply assertions
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 c7086ed7e2 feat(pregel): call after_checkpoint hook when loading and saving channels
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e645c2a085 fix(checkpoint/memory): warn on broken diff chain, guard against cycles
- Add logger.warning when a mid-chain blob is missing (fixes silent truncation bug)
- Add cycle guard to prevent infinite loops on corrupt blob stores
- Fix type annotation on diff_channels from dict[str, Any] to dict[str, str]

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e52b7b2d54 feat(checkpoint/memory): chain-traverse diff blobs in _load_blobs
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney Runkle 6581fdd0a5 fix(channels): align DiffChannel.is_available with BinaryOperatorAggregate 2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 fe2bc286fc feat(channels): implement DiffChannel for incremental checkpoint storage
Adds DiffChannel, a new channel type that stores only per-step write
deltas in checkpoints and reconstructs the full list by replaying the
chain through the operator at load time.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney Runkle c345d337bb feat(channels): add no-op after_checkpoint hook to BaseChannel 2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 d0af83b746 fix(checkpoint/serde): use lazy isinstance check for DiffDelta
Replace duck-typing check with lazy import inside _is_diff_delta helper
function to avoid module-level circular dependency while using proper
isinstance semantics.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.5 eabf926a4f feat(checkpoint/serde): serialize DiffDelta as 'diff' type tag
Add serde support for DiffDelta by implementing dump/load for the "diff" type tag.
This allows the checkpoint system to efficiently store delta objects by serializing
them as msgpack-encoded dicts with {"d": delta, "p": prev_version} structure.

The implementation uses runtime type checking to avoid circular imports and
leverages the existing msgpack ext hooks for proper deserialization of complex
types like LangChain messages.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 6852e0478e feat(checkpoint): add DiffDelta and DiffChainValue protocol types
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 68b2f7bfac docs: add DiffChannel implementation plan
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 95e0fe060c docs: add DiffChannel incremental checkpoint storage design spec
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
27 changed files with 3551 additions and 111 deletions
+1
View File
@@ -100,3 +100,4 @@ dmypy.json
.turbo
.editorconfig
.scratch
.worktrees/
@@ -8,11 +8,13 @@ from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
_ChannelWritesHistory,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -23,7 +25,12 @@ from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
Conn = _internal.Conn # For backward compatibility
@@ -430,6 +437,42 @@ class PostgresSaver(BasePostgresSaver):
with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`) each filtered by `(thread_id, checkpoint_ns)` and
the per-table key. Plain SELECTs let the planner pick straight index
scans; rationale + benchmark in `notes/delta_channel_query_bench.md`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
# Caller didn't specify a target — resolve to the latest
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
# returns the newest; its config carries the resolved id.
target = self.get_tuple(config)
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
with self._cursor() as cur:
cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
parents_rows = cur.fetchall()
cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
writes_rows = cur.fetchall()
cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
blobs_rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
)
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
Convert a database row into a CheckpointTuple object.
@@ -442,6 +485,7 @@ class PostgresSaver(BasePostgresSaver):
including its configuration, metadata, parent checkpoint (if any),
and pending writes.
"""
channel_values = self._load_blobs(value["channel_values"])
return CheckpointTuple(
{
"configurable": {
@@ -454,7 +498,7 @@ class PostgresSaver(BasePostgresSaver):
**value["checkpoint"],
"channel_values": {
**(value["checkpoint"].get("channel_values") or {}),
**self._load_blobs(value["channel_values"]),
**channel_values,
},
},
value["metadata"],
@@ -8,11 +8,13 @@ from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
_ChannelWritesHistory,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -23,7 +25,12 @@ from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_BLOBS_SQL,
SELECT_DELTA_PARENTS_SQL,
SELECT_DELTA_WRITES_SQL,
BasePostgresSaver,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
Conn = _ainternal.Conn # For backward compatibility
@@ -391,6 +398,45 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
`checkpoint_blobs`); rows assembled by the shared pure helper on
`BasePostgresSaver`. Rationale + benchmark in
`notes/delta_channel_query_bench.md`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
target = await self.aget_tuple(config)
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
async with self._cursor() as cur:
await cur.execute(
SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns)
)
parents_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel)
)
writes_rows = await cur.fetchall()
await cur.execute(
SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel)
)
blobs_rows = await cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
parents_rows=parents_rows,
writes_rows=writes_rows,
blobs_rows=blobs_rows,
)
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
Convert a database row into a CheckpointTuple object.
@@ -403,11 +449,18 @@ class AsyncPostgresSaver(BasePostgresSaver):
including its configuration, metadata, parent checkpoint (if any),
and pending writes.
"""
thread_id = value["thread_id"]
checkpoint_ns = value["checkpoint_ns"]
blob_values = value["channel_values"]
channel_values: dict[str, Any] = {}
if blob_values:
channel_values = self._load_blobs(blob_values)
return CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
}
},
@@ -415,15 +468,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
**value["checkpoint"],
"channel_values": {
**(value["checkpoint"].get("channel_values") or {}),
**self._load_blobs(value["channel_values"]),
**channel_values,
},
},
value["metadata"],
(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["parent_checkpoint_id"],
}
}
@@ -8,9 +8,12 @@ from typing import Any, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
PendingWrite,
_ChannelWritesHistory,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.types import TASKS
@@ -152,6 +155,30 @@ INSERT_CHECKPOINT_WRITES_SQL = """
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
# DeltaChannel reconstruction: three plain indexed SELECTs per channel.
# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE
# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are
# 3x-100x faster in the realistic depth range and the Python walk is O(n).
SELECT_DELTA_PARENTS_SQL = """
SELECT checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
"""
SELECT_DELTA_WRITES_SQL = """
SELECT checkpoint_id, type, blob, task_id, idx
FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
SELECT_DELTA_BLOBS_SQL = """
SELECT version, type, blob
FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
@@ -185,16 +212,91 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
)
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
self,
blob_values: Any,
) -> dict[str, Any]:
if not blob_values:
return {}
return {
k.decode(): self.serde.loads_typed((t.decode(), v))
for k, t, v in blob_values
if t.decode() != "empty"
result: dict[str, Any] = {}
for k, t, v in blob_values:
type_tag = t.decode()
if type_tag != "empty":
result[k.decode()] = self.serde.loads_typed((type_tag, v))
return result
def _build_delta_channel_writes_history(
self,
*,
channel: str,
target_id: str,
parents_rows: Sequence[Any],
writes_rows: Sequence[Any],
blobs_rows: Sequence[Any],
) -> _ChannelWritesHistory:
"""Reconstruct one delta channel's history from rows of the three SELECTs.
Pure data transform shared by sync (`PostgresSaver`) and async
(`AsyncPostgresSaver`); both paths run the queries themselves and
feed the rows here.
Walk is newest → oldest from the target's parent. A non-sentinel
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
walk and is returned as the seed so replay starts from it.
Writes stored at `target_id` itself are pending writes for the next
step and are excluded — the walk begins at the target's parent.
"""
parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {}
for r in parents_rows:
cid = r["checkpoint_id"]
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
ancestors: list[str] = []
cid = parent_of.get(target_id)
while cid is not None:
ancestors.append(cid)
cid = parent_of.get(cid)
if not ancestors:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
ancestor_set = set(ancestors)
# Group writes by ancestor cid; sort within (task_id DESC, idx DESC)
# to match the prior CTE ordering — newest write first per ancestor.
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
for r in writes_rows:
cid = r["checkpoint_id"]
if cid not in ancestor_set:
continue
writes_by_cid.setdefault(cid, []).append(
(r["type"], r["blob"], r["task_id"], r["idx"])
)
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
blob_by_ver: dict[str, tuple[str, bytes]] = {
r["version"]: (r["type"], r["blob"]) for r in blobs_rows
}
collected: list[PendingWrite] = [] # newest first; reversed at the end
for cid in ancestors:
# Pre-delta blob terminator: subsumes any writes at this ancestor.
ver = ver_of.get(cid)
if ver is not None:
seed_blob = blob_by_ver.get(ver)
if seed_blob is not None and seed_blob[0] != "empty":
blob_value = self.serde.loads_typed(seed_blob)
if blob_value is not DELTA_SENTINEL:
collected.reverse()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, channel, val))
collected.reverse() # oldest → newest
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
def _dump_blobs(
self,
thread_id: str,
+46 -2
View File
@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
load_checkpoint_tuple = saver._load_checkpoint_tuple
def patched_load_checkpoint_tuple(value):
async def patched_load_checkpoint_tuple(value):
value["checkpoint"].pop("channel_values", None)
return load_checkpoint_tuple(value)
return await load_checkpoint_tuple(value)
monkeypatch.setattr(
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
@@ -371,3 +371,47 @@ async def test_get_checkpoint_no_channel_values(
checkpoint = await saver.aget_tuple(config)
assert checkpoint.checkpoint["channel_values"] == {}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
pytest.importorskip(
"langgraph.channels._delta", reason="langgraph core not installed"
)
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
async with _saver(saver_name) as saver:
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "diff-channel-test-1"}}
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
await graph.ainvoke(
{"messages": [HumanMessage(content="there", id="h2")]}, config
)
state = await graph.aget_state(config)
msgs = state.values["messages"]
assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
assert msgs[0].content == "hi"
assert msgs[1].content == "reply-1"
assert msgs[2].content == "there"
assert msgs[3].content == "reply-3"
@@ -1,9 +1,10 @@
from __future__ import annotations
import contextvars
import copy
import logging
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import ( # noqa: UP035
from typing import (
Any,
Generic,
Literal,
@@ -18,6 +19,9 @@ from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.serde.types import (
DELTA_SENTINEL as DELTA_SENTINEL,
)
from langgraph.checkpoint.serde.types import (
ERROR,
INTERRUPT,
@@ -28,6 +32,16 @@ from langgraph.checkpoint.serde.types import (
V = TypeVar("V", int, float, str)
PendingWrite = tuple[str, str, Any]
# Task-local guard: ContextVar is copied per asyncio Task, so concurrent
# requests on the same event-loop thread do not share this flag. A plain
# `threading.local()` would leak across tasks and let one in-flight
# reconstruction silently short-circuit another.
_DELTA_RECONSTRUCTION: contextvars.ContextVar[bool] = contextvars.ContextVar(
"_DELTA_RECONSTRUCTION", default=False
)
logger = logging.getLogger(__name__)
@@ -119,6 +133,30 @@ class CheckpointTuple(NamedTuple):
pending_writes: list[PendingWrite] | None = None
class _ChannelWritesHistory(NamedTuple):
"""Result of `BaseCheckpointSaver._get_channel_writes_history`.
Storage-level view of what one channel wrote across the ancestor chain
of a target checkpoint:
* `seed` — the nearest ancestor's stored blob value for this channel,
or `DELTA_SENTINEL` if the walk reached the root without finding a
stored value. A non-sentinel seed typically indicates a pre-delta
snapshot preserved across a channel-type migration (e.g.
`BinaryOperatorAggregate` storage extended under `DeltaChannel`).
* `writes` — on-path deltas oldest→newest, one `PendingWrite` per
step that wrote to this channel. Writes stored at the target
checkpoint itself are pending for the next super-step and are
excluded.
Experimental: method surface may change; the NamedTuple shape is the
contract.
"""
seed: Any
writes: list[PendingWrite]
class BaseCheckpointSaver(Generic[V]):
"""Base class for creating a graph checkpointer.
@@ -457,6 +495,104 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""**Experimental.** Query one channel's writes along the parent chain.
Storage-level query, not channel semantics: returns `(seed, writes)`
reflecting what storage knows about a single channel across the
ancestor chain of the target checkpoint identified by `config`.
* `writes` — on-path deltas oldest→newest as `PendingWrite` tuples.
Writes stored at the target `checkpoint_id` itself are pending
for the next super-step and are excluded.
* `seed` — the nearest ancestor's stored blob value for this
channel; `DELTA_SENTINEL` if the walk reached the root without
finding a stored value. A non-sentinel seed typically indicates
a pre-delta snapshot preserved across a channel-type migration.
Walks the **parent chain** (not `list(before=...)`): for forked
threads, only on-path ancestors contribute.
Reference implementation walks `get_tuple` + `parent_config`,
inspecting each ancestor's `channel_values[channel]` for the seed
terminator. Savers with direct storage access (`InMemorySaver`,
`PostgresSaver`) override for performance; the return contract is
fixed here.
Underscore-prefixed because the method surface is experimental.
"""
# Guard against re-entrant calls: when get_tuple() triggers
# reconstruction which calls get_tuple() again, the inner call
# short-circuits here.
if _DELTA_RECONSTRUCTION.get():
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
token = _DELTA_RECONSTRUCTION.set(True)
try:
collected: list[PendingWrite] = [] # newest first; reversed at the end
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
# Pre-delta seed terminator: if the ancestor has a stored
# (non-sentinel) value for this channel, that snapshot
# subsumes any earlier writes on the chain. Stop here.
ancestor_value = tup.checkpoint["channel_values"].get(channel)
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
collected.reverse()
return _ChannelWritesHistory(seed=ancestor_value, writes=collected)
if tup.pending_writes:
# Within a superstep, pending_writes are oldest→newest;
# reverse to scan newest-first.
for write in reversed(tup.pending_writes):
if write[1] != channel:
continue
collected.append(write)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
finally:
_DELTA_RECONSTRUCTION.reset(token)
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""Async version of `_get_channel_writes_history`. See docstring there."""
if _DELTA_RECONSTRUCTION.get():
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
token = _DELTA_RECONSTRUCTION.set(True)
try:
collected: list[PendingWrite] = []
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
ancestor_value = tup.checkpoint["channel_values"].get(channel)
if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL:
collected.reverse()
return _ChannelWritesHistory(seed=ancestor_value, writes=collected)
if tup.pending_writes:
for write in reversed(tup.pending_writes):
if write[1] != channel:
continue
collected.append(write)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
finally:
_DELTA_RECONSTRUCTION.reset(token)
def get_next_version(self, current: V | None, channel: None) -> V:
"""Generate the next version ID for a channel.
@@ -9,18 +9,21 @@ from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
from types import TracebackType
from typing import Any
from typing import Any, cast
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
PendingWrite,
SerializerProtocol,
_ChannelWritesHistory,
get_checkpoint_id,
get_checkpoint_metadata,
)
@@ -121,16 +124,91 @@ class InMemorySaver(
return self.stack.__exit__(__exc_type, __exc_value, __traceback)
def _load_blobs(
self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
self,
thread_id: str,
checkpoint_ns: str,
versions: ChannelVersions,
) -> dict[str, Any]:
channel_values: dict[str, Any] = {}
for k, v in versions.items():
kk = (thread_id, checkpoint_ns, k, v)
if kk in self.blobs:
vv = self.blobs[kk]
if vv[0] != "empty":
channel_values[k] = self.serde.loads_typed(vv)
return channel_values
result: dict[str, Any] = {}
for k, ver in versions.items():
kk = (thread_id, checkpoint_ns, k, ver)
if kk not in self.blobs:
continue
vv = self.blobs[kk]
if vv[0] == "empty":
continue
result[k] = self.serde.loads_typed(vv)
return result
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
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. 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] = []
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:
break
chain.append(current)
_, _, parent = entry
current = parent
# Scan newest→oldest. A pre-delta blob on an ancestor terminates the
# walk and is bound as `seed`; without this, 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[PendingWrite] = [] # 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 _ChannelWritesHistory(
seed=blob_value, writes=collected
)
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
# Within a superstep, sorted by (task_id, idx) = oldest → newest;
# reverse for newest-first scan.
for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
step_writes.items(), reverse=True
):
if ch != channel:
continue
val = self.serde.loads_typed(serialized)
collected.append((tid, ch, val))
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
return self._get_channel_writes_history(config, channel)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Get a checkpoint tuple from the in-memory storage.
@@ -153,13 +231,16 @@ class InMemorySaver(
checkpoint, metadata, parent_checkpoint_id = saved
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
channel_values = self._load_blobs(
thread_id,
checkpoint_ns,
checkpoint_["channel_versions"],
)
return CheckpointTuple(
config=config,
checkpoint={
**checkpoint_,
"channel_values": self._load_blobs(
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
),
"channel_values": channel_values,
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
@@ -183,19 +264,26 @@ class InMemorySaver(
checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
checkpoint_ = self.serde.loads_typed(checkpoint)
return CheckpointTuple(
config={
resolved_config = cast(
RunnableConfig,
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
)
channel_values = self._load_blobs(
thread_id,
checkpoint_ns,
checkpoint_["channel_versions"],
)
return CheckpointTuple(
config=resolved_config,
checkpoint={
**checkpoint_,
"channel_values": self._load_blobs(
thread_id, checkpoint_ns, checkpoint_["channel_versions"]
),
"channel_values": channel_values,
},
metadata=self.serde.loads_typed(metadata),
pending_writes=[
@@ -290,21 +378,27 @@ class InMemorySaver(
checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
yield CheckpointTuple(
config={
list_config = cast(
RunnableConfig,
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
)
channel_values = self._load_blobs(
thread_id,
checkpoint_ns,
checkpoint_["channel_versions"],
)
yield CheckpointTuple(
config=list_config,
checkpoint={
**checkpoint_,
"channel_values": self._load_blobs(
thread_id,
checkpoint_ns,
checkpoint_["channel_versions"],
),
"channel_values": channel_values,
},
metadata=metadata,
parent_config=(
@@ -33,14 +33,13 @@ from langchain_core.load.load import Reviver
from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
from langgraph.checkpoint.serde.types import SendProtocol
from langgraph.checkpoint.serde.types import DELTA_SENTINEL, SendProtocol
from langgraph.store.base import Item
if TYPE_CHECKING:
from langgraph.checkpoint.serde._msgpack import (
AllowedMsgpackModules,
)
from langgraph.checkpoint.serde.types import SendProtocol
LC_REVIVER = Reviver()
EMPTY_BYTES = b""
@@ -252,6 +251,8 @@ class JsonPlusSerializer(SerializerProtocol):
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
if obj is None:
return "null", EMPTY_BYTES
elif obj is DELTA_SENTINEL:
return "delta", EMPTY_BYTES
elif isinstance(obj, bytes):
return "bytes", obj
elif isinstance(obj, bytearray):
@@ -278,6 +279,8 @@ class JsonPlusSerializer(SerializerProtocol):
return ormsgpack.unpackb(
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
elif type_ == "delta":
return DELTA_SENTINEL
elif self.pickle_fallback and type_ == "pickle":
return pickle.loads(data_)
else:
@@ -14,6 +14,25 @@ INTERRUPT = "__interrupt__"
RESUME = "__resume__"
TASKS = "__pregel_tasks"
class _DeltaSentinel:
"""Singleton marker stored (as zero bytes) in checkpoint_blobs for a
DeltaChannel field. The actual per-step writes live in checkpoint_writes
and are replayed through the reducer at load time.
Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same
module-level instance.
"""
__slots__ = ()
def __repr__(self) -> str:
return "DELTA_SENTINEL"
DELTA_SENTINEL = _DeltaSentinel()
Value = TypeVar("Value", covariant=True)
Update = TypeVar("Update", contravariant=True)
C = TypeVar("C")
+13
View File
@@ -997,3 +997,16 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
# No blocking should occur - inner is serialized as dict, not ext
assert "blocked" not in caplog.text.lower()
assert result == obj
def test_delta_sentinel_serde_round_trip() -> None:
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
# Zero-byte "delta" tag — no allowlist change needed.
assert type_tag == "delta"
assert blob == b""
loaded = serde.loads_typed((type_tag, blob))
assert loaded is DELTA_SENTINEL
+345 -2
View File
@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
@@ -208,8 +209,6 @@ class TestMemorySaver:
async def test_memory_saver() -> None:
from langgraph.checkpoint.memory import InMemorySaver
memory_saver = InMemorySaver()
assert isinstance(memory_saver, InMemorySaver)
@@ -320,3 +319,347 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
assert direct is not None
expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict()
assert direct.checkpoint["channel_values"]["foo"] == expected
class TestInMemorySaverDeltaChannel:
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0000000000000000"
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
}
result = saver._load_blobs(thread_id, ns, {channel: v1})
assert channel in result
assert result[channel] is DELTA_SENTINEL
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
"""_get_channel_writes_history 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()
thread_id, ns, channel = "t1", "", "messages"
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp2 = empty_checkpoint()
cp2["id"] = "cp2"
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
}
# 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"}),
"",
)
# 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": "pending"}),
"",
)
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": "cp2",
}
}
result = saver._get_channel_writes_history(config, channel)
assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes]
assert values == [{"content": "hi"}]
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",
}
}
result = saver._get_channel_writes_history(config, channel)
assert result.seed is DELTA_SENTINEL
assert result.writes == []
class TestBaseFallbackGetChannelWrites:
"""Exercises the `BaseCheckpointSaver._get_channel_writes_history` default
implementation — the path third-party savers inherit when they don't
override `_get_channel_writes_history` 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_history = (
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
)
_aget_channel_writes_history = (
InMemorySaver.__mro__[1]._aget_channel_writes_history # 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_history(config, "messages")
assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes]
assert values == [{"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_history(config, "messages")
assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes]
assert values == [{"content": "first"}, {"content": "second"}]
async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None:
"""Regression: the re-entrancy guard must be task-local, not thread-local.
Two concurrent `_aget_channel_writes_history` calls on the same
event-loop thread must each see their full reconstructed writes. A
`threading.local()` guard would let whichever task set it first
short-circuit the other to `writes=[]`.
"""
import asyncio
saver, thread_id, ns = self._build_saver_with_chain()
# Force the two tasks to interleave across the `set(True)` boundary:
# each `aget_tuple` yields control, so if the guard were thread-local
# the second task would observe `active=True` set by the first.
orig_aget_tuple = saver.aget_tuple
async def slow_aget_tuple(config: RunnableConfig) -> Any:
await asyncio.sleep(0)
return await orig_aget_tuple(config)
saver.aget_tuple = slow_aget_tuple # type: ignore[method-assign]
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
results = await asyncio.gather(
saver._aget_channel_writes_history(config, "messages"),
saver._aget_channel_writes_history(config, "messages"),
)
expected_values = [{"content": "first"}, {"content": "second"}]
for result in results:
assert result.seed is DELTA_SENTINEL
values = [v for _, _, v in result.writes]
assert values == expected_values
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_history(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.
values = [v for _, _, v in result.writes]
assert values == ["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_history(config, channel)
values = [v for _, _, v in result.writes]
# The pre-delta write under cp1 must not appear (the blob subsumes it).
assert "PRE-DELTA-WRITE" not in values
# And the pending write at the target is never folded in.
assert "PENDING-AT-TARGET" not in values
+177
View File
@@ -0,0 +1,177 @@
from __future__ import annotations
import copy as _copy
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
from typing_extensions import Self
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel, Value
from langgraph.channels.binop import _get_overwrite
from langgraph.errors import EmptyChannelError
__all__ = ("DeltaChannel",)
def _empty(typ: Any) -> Any:
try:
return typ()
except Exception:
return []
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
"""Experimental — private API, subject to change or removal without notice.
Imported from the underscored module `langgraph.channels._delta` on purpose;
not re-exported from `langgraph.channels`. Intended for internal use only
while we validate the design on real workloads.
A channel that stores only a sentinel in checkpoints; per-step writes are
stored in checkpoint_writes and replayed through the operator at load time.
Use with append-style reducers (e.g. `add_messages`) on long-running threads
to eliminate O(N²) blob growth — storage is O(N) using the writes table that
every checkpointer already maintains.
Reconstruction replays every ancestor write through the operator, so
per-get cost scales with thread depth. Compaction for deep threads is
a follow-up — today, use this on threads of a few hundred turns.
Usage::
from langgraph.channels._delta import DeltaChannel
class State(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
"""
__slots__ = (
"value",
"operator",
)
def __init__(
self,
operator: Callable[[Any, Any], Any],
) -> None:
super().__init__(list)
self.operator = operator
self.value: Any = []
def __eq__(self, other: object) -> bool:
if not isinstance(other, DeltaChannel):
return False
if (
self.operator.__name__ != "<lambda>"
and other.operator.__name__ != "<lambda>"
):
return self.operator is other.operator
return True
@property
def ValueType(self) -> Any:
return self.typ
@property
def UpdateType(self) -> Any:
return self.typ
def copy(self) -> Self:
new: DeltaChannel[Value] = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
return new
def _apply_write(self, value: Any, write: Any) -> Any:
"""Apply one write to `value` and return the new value.
An `Overwrite` replaces the value; any other write is folded through
the operator. 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:
return (
_copy.copy(overwrite_value)
if overwrite_value is not None
else _empty(self.typ)
)
base = _empty(self.typ) if value is MISSING else value
return self.operator(base, write)
def from_checkpoint(self, checkpoint: Any) -> Self:
"""Initialize from a seed value.
Pregel's hydration path calls this with the `seed` returned by
`saver.get_channel_history`:
* `MISSING` / `DELTA_SENTINEL` → channel starts empty. The walk
either reached the root (fresh delta thread) or found nothing
to seed from.
* any other value → use as the base value. Typically a pre-delta
blob preserved across a channel-type migration; `replay_writes`
folds subsequent deltas on top.
"""
new: DeltaChannel[Value] = DeltaChannel(self.operator)
new.typ = self.typ
new.key = self.key
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
new.value = _empty(new.typ)
else:
new.value = checkpoint
return new
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
"""Fold a sequence of `PendingWrite` tuples into the current value.
Called after `from_checkpoint` during pregel hydration to replay
per-step deltas from on-path ancestors through the reducer. Writes
are oldest→newest. `Overwrite` values inside the stream reset the
reducer state at that point, same as during a live super-step.
The `task_id` and `channel` fields of each `PendingWrite` are
ignored — `_get_channel_writes_history` has already filtered to
this channel.
"""
for _, _, value in writes:
self.value = self._apply_write(self.value, value)
def update(self, values: Sequence[Any]) -> bool:
if not values:
return False
seen_overwrite = False
for value in values:
is_overwrite, _ = _get_overwrite(value)
if is_overwrite:
if seen_overwrite:
from langgraph.errors import (
ErrorCode,
InvalidUpdateError,
create_error_message,
)
msg = create_error_message(
message="Can receive only one Overwrite value per super-step.",
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
)
raise InvalidUpdateError(msg)
seen_overwrite = True
elif seen_overwrite:
# Post-Overwrite writes within the same super-step are dropped.
continue
self.value = self._apply_write(self.value, value)
return True
def get(self) -> Any:
if self.value is MISSING:
raise EmptyChannelError()
return self.value
def is_available(self) -> bool:
return self.value is not MISSING
def checkpoint(self) -> Any:
return DELTA_SENTINEL
+74 -40
View File
@@ -184,63 +184,97 @@ def add_messages(
```
"""
remove_all_idx = None
# coerce to list
# 1. Coerce scalars to lists.
if not isinstance(left, list):
left = [left] # type: ignore[assignment]
if not isinstance(right, list):
right = [right] # type: ignore[assignment]
# coerce to message
left = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(left)
]
right = [
left = cast(list, left)
# 2. Normalize `left`. After the first call, `left` is the previous return
# value of `add_messages` — a list of fully-resolved BaseMessages with
# IDs — and needs no work. Fresh user input (dicts, tuples, message
# chunks, BaseMessages without IDs) falls through to full conversion.
left_msgs: list[BaseMessage]
if (
left
and isinstance(left[0], BaseMessage)
and not isinstance(left[0], BaseMessageChunk)
and left[0].id is not None
):
left_msgs = left
else:
left_msgs = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(left)
]
for m in left_msgs:
if m.id is None:
m.id = str(uuid.uuid4())
# 3. Normalize `right` — always fresh input. Assign missing IDs and detect
# any RemoveMessage sentinels in a single pass.
right_msgs: list[BaseMessage] = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(right)
]
# assign missing ids
for m in left:
remove_all_idx: int | None = None
has_remove = False
for idx, m in enumerate(right_msgs):
if m.id is None:
m.id = str(uuid.uuid4())
for idx, m in enumerate(right):
if m.id is None:
m.id = str(uuid.uuid4())
if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
remove_all_idx = idx
if isinstance(m, RemoveMessage):
has_remove = True
if m.id == REMOVE_ALL_MESSAGES:
remove_all_idx = idx
# 4. REMOVE_ALL_MESSAGES: discard everything up to and including the sentinel.
if remove_all_idx is not None:
return right[remove_all_idx + 1 :]
return right_msgs[remove_all_idx + 1 :]
# merge
merged = left.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
# 5. Decide fast vs. slow path. The fast path (pure append) is only valid
# when `right` has no removals, no intra-right duplicate IDs, and no IDs
# that overlap with `left` — any of those would force the indexed merge
# below to update or dedup.
pure_append = False
if not has_remove:
left_ids = {m.id for m in left_msgs}
right_ids = {m.id for m in right_msgs}
pure_append = len(right_ids) == len(right_msgs) and right_ids.isdisjoint(
left_ids
)
if pure_append:
merged = left_msgs + right_msgs
else:
# 6. Slow path: build id→index map over `left`, then replay `right`.
# In-place replacement for matching IDs, append for new IDs, and a
# deferred removal pass so RemoveMessages can target either side.
merged = left_msgs.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right_msgs:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
merged = [m for m in merged if m.id not in ids_to_remove]
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
merged = [m for m in merged if m.id not in ids_to_remove]
# 7. Apply optional output format.
if format == "langchain-openai":
merged = _format_messages(merged)
elif format:
return _format_messages(merged)
if format:
msg = f"Unrecognized {format=}. Expected one of 'langchain-openai', None."
raise ValueError(msg)
else:
pass
return merged
+34 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import collections.abc
import inspect
import logging
import typing
@@ -47,7 +48,8 @@ from langgraph._internal._pydantic import create_model
from langgraph._internal._runnable import coerce_to_runnable
from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.binop import BinaryOperatorAggregate, _strip_extras
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue, LastValueAfterFinish
from langgraph.channels.named_barrier_value import (
@@ -1082,6 +1084,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
CompiledStateGraph: The compiled `StateGraph`.
"""
checkpointer = ensure_valid_checkpointer(checkpointer)
serde_allowlist: set[tuple[str, ...]] | None = None
if _serde.STRICT_MSGPACK_ENABLED:
schema_types: list[type[Any]] = [
@@ -1667,6 +1670,36 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
# Search through all annotated medata to find channel annotations
for item in meta:
if isinstance(item, BaseChannel):
if isinstance(item, DeltaChannel) and hasattr(typ, "__origin__"):
origin = typ.__origin__
# Unwrap parameterized Required[X]/NotRequired[X] to X
# (e.g. Annotated[NotRequired[dict[...]], ...]).
if hasattr(origin, "__origin__") and origin.__origin__ in (
Required,
NotRequired,
):
origin = origin.__args__[0]
outer = _strip_extras(origin)
if outer in (
collections.abc.Sequence,
collections.abc.MutableSequence,
):
outer = list
elif outer in (
collections.abc.Mapping,
collections.abc.MutableMapping,
):
outer = dict
elif outer in (
collections.abc.Set,
collections.abc.MutableSet,
):
outer = set
item.typ = outer
try:
item.value = outer()
except Exception:
item.value = []
return item
elif isclass(item) and issubclass(item, BaseChannel):
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
+79 -9
View File
@@ -3,11 +3,13 @@ from __future__ import annotations
from collections.abc import Mapping
from datetime import datetime, timezone
from langgraph.checkpoint.base import Checkpoint
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel
from langgraph.channels._delta import DeltaChannel
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 4
@@ -58,8 +60,22 @@ def create_checkpoint(
def channels_from_checkpoint(
specs: Mapping[str, BaseChannel | ManagedValueSpec],
checkpoint: Checkpoint,
*,
saver: BaseCheckpointSaver | None = None,
config: RunnableConfig | None = None,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Get channels from a checkpoint."""
"""Hydrate channels from a checkpoint.
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
is sufficient the stored value IS the reconstructed state.
`DeltaChannel` is the exception: its stored value is a sentinel; the
full state is spread across `checkpoint_writes` along the ancestor
chain. When `saver` and `config` are provided, this function fetches
that history via `saver._get_channel_writes_history` and folds it
through the channel's reducer. Without them (static contexts — graph
drawing, unit tests), delta channels fall back to empty.
"""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -67,13 +83,67 @@ def channels_from_checkpoint(
channel_specs[k] = v
else:
managed_specs[k] = v
return (
{
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
for k, v in channel_specs.items()
},
managed_specs,
)
channels: dict[str, BaseChannel] = {}
for k, spec in channel_specs.items():
ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING)
if (
isinstance(spec, DeltaChannel)
and saver is not None
and config is not None
and (stored is MISSING or stored is DELTA_SENTINEL)
):
# Target's own blob is empty/sentinel — walk ancestors for
# seed + writes. Skipping this when `stored` is a real value
# preserves state written via `update_state` or sitting at the
# tip of a pre-migration thread: the saver's ancestor walk
# intentionally excludes the target's own blob, so without
# this short-circuit we'd lose it.
history = saver._get_channel_writes_history(config, k)
delta_ch = spec.from_checkpoint(history.seed)
delta_ch.replay_writes(history.writes)
ch = delta_ch
else:
ch = spec.from_checkpoint(stored)
channels[k] = ch
return channels, managed_specs
async def achannels_from_checkpoint(
specs: Mapping[str, BaseChannel | ManagedValueSpec],
checkpoint: Checkpoint,
*,
saver: BaseCheckpointSaver | None = None,
config: RunnableConfig | None = None,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Async version of `channels_from_checkpoint`. See docstring there."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
if isinstance(v, BaseChannel):
channel_specs[k] = v
else:
managed_specs[k] = v
channels: dict[str, BaseChannel] = {}
for k, spec in channel_specs.items():
ch: BaseChannel
stored = checkpoint["channel_values"].get(k, MISSING)
if (
isinstance(spec, DeltaChannel)
and saver is not None
and config is not None
and (stored is MISSING or stored is DELTA_SENTINEL)
):
history = await saver._aget_channel_writes_history(config, k)
delta_ch = spec.from_checkpoint(history.seed)
delta_ch.replay_writes(history.writes)
ch = delta_ch
else:
ch = spec.from_checkpoint(stored)
channels[k] = ch
return channels, managed_specs
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
+10 -3
View File
@@ -92,6 +92,7 @@ from langgraph.pregel._algo import (
task_path_str,
)
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -1273,7 +1274,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
)
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
@@ -1476,8 +1480,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
)
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
self.channels, self.managed = await achannels_from_checkpoint(
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
+21 -2
View File
@@ -122,6 +122,7 @@ from langgraph.pregel._algo import (
)
from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -1052,6 +1053,10 @@ class Pregel(
channels, managed = channels_from_checkpoint(
self.channels,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
@@ -1168,9 +1173,13 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
channels, managed = channels_from_checkpoint(
channels, managed = await achannels_from_checkpoint(
self.channels,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
@@ -1541,6 +1550,11 @@ class Pregel(
channels, managed = channels_from_checkpoint(
self.channels,
checkpoint,
saver=self.checkpointer
if saved is not None
and isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config if saved is not None else None,
)
values, as_node = updates[0][:2]
@@ -1984,9 +1998,14 @@ class Pregel(
)
if saved:
checkpoint_config = patch_configurable(config, saved.config[CONF])
channels, managed = channels_from_checkpoint(
channels, managed = await achannels_from_checkpoint(
self.channels,
checkpoint,
saver=self.checkpointer
if saved is not None
and isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config if saved is not None else None,
)
values, as_node = updates[0][:2]
# no values, just clear all tasks
@@ -0,0 +1,290 @@
"""Benchmark: add_messages fast-path optimizations.
Both implementations are inlined so the benchmark is self-contained and
immune to import-cache or installed-vs-local confusion.
Run directly:
python tests/test_add_messages_benchmark.py
Or via pytest (correctness only, numbers printed to stdout):
pytest tests/test_add_messages_benchmark.py -s -v
"""
import statistics
import time
import tracemalloc
import uuid
from typing import cast
from langchain_core.messages import (
AIMessage,
BaseMessage,
BaseMessageChunk,
HumanMessage,
RemoveMessage,
convert_to_messages,
message_chunk_to_message,
)
from langgraph.graph.message import REMOVE_ALL_MESSAGES
# ── original implementation (pre-optimisation) ────────────────────────────────
def _add_messages_original(left, right):
remove_all_idx = None
if not isinstance(left, list):
left = [left]
if not isinstance(right, list):
right = [right]
left = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(left)
]
right = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(right)
]
for m in left:
if m.id is None:
m.id = str(uuid.uuid4())
for idx, m in enumerate(right):
if m.id is None:
m.id = str(uuid.uuid4())
if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
remove_all_idx = idx
if remove_all_idx is not None:
return right[remove_all_idx + 1 :]
merged = left.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
return [m for m in merged if m.id not in ids_to_remove]
# ── optimised implementation ──────────────────────────────────────────────────
def _add_messages_optimized(left, right):
if not isinstance(left, list):
left = [left]
if not isinstance(right, list):
right = [right]
# Optimisation 1: skip conversion + ID assignment on left when it already
# contains fully-resolved BaseMessage objects (the common case after the
# first call, since add_messages always returns list[BaseMessage] with IDs).
if (
left
and isinstance(left[0], BaseMessage)
and not isinstance(left[0], BaseMessageChunk)
):
left = cast(list[BaseMessage], left)
else:
left = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(left)
]
for m in left:
if m.id is None:
m.id = str(uuid.uuid4())
# always normalise right — it's fresh external input
right = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(right)
]
remove_all_idx = None
has_remove = False
for idx, m in enumerate(right):
if m.id is None:
m.id = str(uuid.uuid4())
if isinstance(m, RemoveMessage):
has_remove = True
if m.id == REMOVE_ALL_MESSAGES:
remove_all_idx = idx
if remove_all_idx is not None:
return right[remove_all_idx + 1 :]
# Optimisation 2: pure-append fast path — no removals and no ID overlaps.
# Builds one set over left instead of copying left + building a full dict.
if not has_remove:
left_ids = {m.id for m in left}
if not any(m.id in left_ids for m in right):
return left + right
# slow path: updates or removals present — full indexed merge
merged = left.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
return [m for m in merged if m.id not in ids_to_remove]
# ── helpers ───────────────────────────────────────────────────────────────────
def _make_messages(n: int) -> list[BaseMessage]:
return [
(HumanMessage if i % 2 == 0 else AIMessage)(
content=f"message {i}", id=str(uuid.uuid4())
)
for i in range(n)
]
def _bench_time(fn, left, right, *, iters: int = 2_000) -> float:
"""Return median latency in microseconds."""
for _ in range(100):
fn(list(left), list(right))
times = []
for _ in range(iters):
left_copy, right_copy = list(left), list(right)
t0 = time.perf_counter()
fn(left_copy, right_copy)
times.append(time.perf_counter() - t0)
return statistics.median(times) * 1e6
def _bench_memory(fn, left, right) -> int:
"""Return peak memory allocated during a single call (bytes)."""
# one warm-up so any lazy init is excluded
fn(list(left), list(right))
left_copy, right_copy = list(left), list(right)
tracemalloc.start()
tracemalloc.clear_traces()
fn(left_copy, right_copy)
_, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()
return peak
# ── scenarios ─────────────────────────────────────────────────────────────────
SCENARIOS = [
("pure append 1 → 1 msg", 1, 1, "append"),
("pure append 10 → 1 msg", 10, 1, "append"),
("pure append 100 → 1 msg", 100, 1, "append"),
("pure append 1000 → 1 msg", 1000, 1, "append"),
("pure append 1000 → 5 msgs", 1000, 5, "append"),
("update existing 100 → 1 msg", 100, 1, "update"),
("remove message 100 → 1 msg", 100, 1, "remove"),
]
def _make_inputs(n_left, n_right, mode):
left = _make_messages(n_left)
right = _make_messages(n_right)
if mode == "update":
right[0] = AIMessage(content="updated", id=left[0].id)
elif mode == "remove":
right = [RemoveMessage(id=left[0].id)]
return left, right
# ── main output ───────────────────────────────────────────────────────────────
COL = 36
def run_benchmarks() -> None:
print()
print("=" * 88)
print("add_messages benchmark — time (µs, median of 2 000 iterations)")
print("=" * 88)
print(f"{'Scenario':<{COL}} {'Original':>10} {'Optimized':>11} {'Speedup':>8}")
print("-" * 88)
for label, n_left, n_right, mode in SCENARIOS:
left, right = _make_inputs(n_left, n_right, mode)
t_orig = _bench_time(_add_messages_original, left, right)
t_opt = _bench_time(_add_messages_optimized, left, right)
print(f"{label:<{COL}} {t_orig:>10.2f} {t_opt:>11.2f} {t_orig / t_opt:>7.2f}x")
print()
print("=" * 88)
print("add_messages benchmark — peak memory allocated per call (bytes)")
print("=" * 88)
print(f"{'Scenario':<{COL}} {'Original':>10} {'Optimized':>11} {'Reduction':>10}")
print("-" * 88)
for label, n_left, n_right, mode in SCENARIOS:
left, right = _make_inputs(n_left, n_right, mode)
m_orig = _bench_memory(_add_messages_original, left, right)
m_opt = _bench_memory(_add_messages_optimized, left, right)
reduction = (1 - m_opt / m_orig) * 100 if m_orig else 0.0
print(f"{label:<{COL}} {m_orig:>10,} {m_opt:>11,} {reduction:>9.1f}%")
print()
print("=" * 88)
print("Simulated long thread — 200 steps × 2 msgs appended per step")
print("=" * 88)
for name, fn in [
("original", _add_messages_original),
("optimized", _add_messages_optimized),
]:
state: list = []
t0 = time.perf_counter()
for step in range(200):
new_msgs = [
HumanMessage(content=f"step {step} human", id=str(uuid.uuid4())),
AIMessage(content=f"step {step} ai", id=str(uuid.uuid4())),
]
state = fn(state, new_msgs)
elapsed = (time.perf_counter() - t0) * 1_000
print(f" {name:<12} {elapsed:.2f} ms ({len(state)} messages)")
print()
# ── pytest entry-points ───────────────────────────────────────────────────────
def test_add_messages_correctness():
"""Optimised implementation must match original output for every scenario."""
for label, n_left, n_right, mode in SCENARIOS:
left, right = _make_inputs(n_left, n_right, mode)
expected = _add_messages_original(list(left), list(right))
actual = _add_messages_optimized(list(left), list(right))
assert len(actual) == len(expected), f"[{label}] length mismatch"
for a, b in zip(actual, expected):
assert type(a) is type(b), f"[{label}] type mismatch"
assert a.id == b.id, f"[{label}] id mismatch"
assert a.content == b.content, f"[{label}] content mismatch"
def test_add_messages_benchmark(capsys):
run_benchmarks()
out = capsys.readouterr().out
assert "Speedup" in out
assert "Optimized" in out
if __name__ == "__main__":
run_benchmarks()
+504
View File
@@ -2,13 +2,17 @@ import operator
from collections.abc import Sequence
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph._internal._typing import MISSING
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels._delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.graph.message import add_messages
pytestmark = pytest.mark.anyio
@@ -117,3 +121,503 @@ def test_untracked_value() -> None:
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
with pytest.raises(EmptyChannelError):
new_channel.get()
def test_delta_channel_basic_two_steps() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
# Step 1: one message added
ch.update([HumanMessage(content="hi", id="h1")])
d1 = ch.checkpoint()
assert d1 is DELTA_SENTINEL
# Step 2: another message
ch.update([AIMessage(content="hello", id="a1")])
d2 = ch.checkpoint()
assert d2 is DELTA_SENTINEL
# Full accumulated value is preserved in memory
assert len(ch.get()) == 2
assert ch.get()[0].content == "hi"
assert ch.get()[1].content == "hello"
def test_delta_channel_from_checkpoint_writes_list() -> None:
"""replay_writes on a fresh channel replays through the operator."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "messages", HumanMessage(content="hi", id="h1")),
("t1", "messages", AIMessage(content="hello", id="a1")),
("t2", "messages", HumanMessage(content="bye", id="h2")),
]
)
msgs = ch.get()
assert len(msgs) == 3
assert msgs[0].content == "hi"
assert msgs[1].content == "hello"
assert msgs[2].content == "bye"
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
from langchain_core.messages import HumanMessage
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
# Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat
spec = DeltaChannel(add_messages)
old_value = [HumanMessage(content="old", id="h1")]
ch = spec.from_checkpoint(old_value)
assert ch.get() == old_value
def test_delta_channel_overwrite() -> None:
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
from langgraph.types import Overwrite
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
ch.update([HumanMessage(content="old", id="h1")])
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
d = ch.checkpoint()
assert d is DELTA_SENTINEL
# After overwrite, value is reset to only the new message
assert len(ch.get()) == 1
assert ch.get()[0].content == "new"
def test_delta_channel_remove_message_and_replay() -> None:
"""RemoveMessage must round-trip correctly when writes are replayed."""
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(MISSING)
# Step 1: add two messages
ch.update([HumanMessage(content="hi", id="h1")])
ch.update([AIMessage(content="hello", id="a1")])
assert ch.get() == [
HumanMessage(content="hi", id="h1"),
AIMessage(content="hello", id="a1"),
]
# Step 2: remove the AI message
ch.update([RemoveMessage(id="a1")])
assert ch.get() == [HumanMessage(content="hi", id="h1")]
# Replay the writes list from scratch — must reproduce the post-remove state
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "messages", HumanMessage(content="hi", id="h1")),
("t1", "messages", AIMessage(content="hello", id="a1")),
("t2", "messages", RemoveMessage(id="a1")),
]
)
assert ch2.get() == [HumanMessage(content="hi", id="h1")]
def test_delta_channel_update_by_id_and_replay() -> None:
"""Updating a message by ID must round-trip correctly through writes replay."""
from langchain_core.messages import HumanMessage
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(MISSING)
# Step 1: add a message
ch.update([HumanMessage(content="original", id="h1")])
# Step 2: update the same message by ID
ch.update([HumanMessage(content="updated", id="h1")])
assert ch.get() == [HumanMessage(content="updated", id="h1")]
# Replay writes — must produce the updated message, not the original
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "messages", HumanMessage(content="original", id="h1")),
("t1", "messages", HumanMessage(content="updated", id="h1")),
]
)
assert len(ch2.get()) == 1
assert ch2.get()[0].content == "updated"
def test_delta_channel_checkpoint_returns_sentinel() -> None:
"""checkpoint() always returns DELTA_SENTINEL regardless of state."""
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
assert ch.checkpoint() is DELTA_SENTINEL
from langchain_core.messages import HumanMessage
ch.update([HumanMessage(content="hi", id="h1")])
assert ch.checkpoint() is DELTA_SENTINEL
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
n = {"v": 0}
def respond(state: State) -> dict:
n["v"] += 1
return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
saver = InMemorySaver()
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "t1"}}
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
# get_tuple returns raw storage shape — channel_values stores DELTA_SENTINEL
# for delta channels; the reconstructed writes flow separately via
# saver._get_channel_writes_history.
saved = saver.get_tuple(config)
assert saved is not None
assert "messages" in saved.checkpoint["channel_values"]
assert saved.checkpoint["channel_values"]["messages"] is DELTA_SENTINEL
state = graph.get_state(config)
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
def _delta_channel_with_type(operator, typ):
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
from typing import Annotated
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.state import _get_channel
return _get_channel("_test", Annotated[typ, DeltaChannel(operator)])
def test_delta_channel_dict_reducer_fresh_channel() -> None:
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
# Should be available (not raise EmptyChannelError) and start empty
assert ch.is_available()
assert ch.get() == {}
def test_delta_channel_dict_reducer_basic_updates() -> None:
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
from langgraph.checkpoint.base import DELTA_SENTINEL
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
ch.update([{"a": 1}])
d1 = ch.checkpoint()
assert d1 is DELTA_SENTINEL
ch.update([{"b": 2}])
d2 = ch.checkpoint()
assert d2 is DELTA_SENTINEL
assert ch.get() == {"a": 1, "b": 2}
def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
"""replay_writes on a fresh channel replays through a dict merge reducer."""
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "files", {"a": 1}),
("t1", "files", {"b": 2}),
("t2", "files", {"c": 3}),
]
)
assert ch.get() == {"a": 1, "b": 2, "c": 3}
def test_delta_channel_dict_reducer_with_deletions() -> None:
"""Dict reducer that treats None values as deletions works end-to-end (deepagents pattern)."""
def merge_files(left: dict | None, right: dict) -> dict:
if left is None:
return {k: v for k, v in right.items() if v is not None}
result = {**left}
for k, v in right.items():
if v is None:
result.pop(k, None)
else:
result[k] = v
return result
ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING)
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
# Delete file1, add file3
ch.update([{"file1.py": None, "file3.py": "content3"}])
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
# Confirm writes reconstruction produces the same result
spec = _delta_channel_with_type(merge_files, dict)
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "files", {"file1.py": "content1", "file2.py": "content2"}),
("t1", "files", {"file1.py": None, "file3.py": "content3"}),
]
)
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
"""Overwrite(dict) in update() must preserve dict shape, not coerce to list."""
from langgraph.types import Overwrite
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
ch.update([{"a": 1}])
ch.update([Overwrite({"b": 2, "c": 3})])
assert ch.get() == {"b": 2, "c": 3}
def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
"""Overwrite(dict) embedded in replayed writes must reconstruct as dict."""
from langgraph.types import Overwrite
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
ch = spec.from_checkpoint(DELTA_SENTINEL)
ch.replay_writes(
[
("t0", "files", {"a": 1}),
("t1", "files", Overwrite({"x": 10, "y": 20})),
("t2", "files", {"z": 30}),
]
)
assert ch.get() == {"x": 10, "y": 20, "z": 30}
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`.
This is the shape the deepagents filesystem middleware uses for its
`files` field; without unwrapping NotRequired we'd fall through to `list`
and blow up on the first dict operator call.
"""
from typing import Annotated
from typing_extensions import NotRequired
from langgraph.channels._delta import DeltaChannel
from langgraph.graph.state import _get_channel
def merge_dicts(left: dict | None, right: dict) -> dict:
if left is None:
return dict(right)
return {**left, **right}
annotation = Annotated[
NotRequired[dict[str, int]],
DeltaChannel(merge_dicts),
]
ch = _get_channel("files", annotation).from_checkpoint(MISSING)
assert ch.get() == {}
ch.update([{"a": 1}])
ch.update([{"b": 2}])
assert ch.get() == {"a": 1, "b": 2}
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel.
Mirrors the deepagents filesystem pattern: `files: Annotated[dict, reducer]`
where the reducer merges dicts and treats None values as deletions.
"""
from typing import Annotated
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
def merge_files(left: dict | None, right: dict) -> dict:
if left is None:
return {k: v for k, v in right.items() if v is not None}
result = {**left}
for k, v in right.items():
if v is None:
result.pop(k, None)
else:
result[k] = v
return result
class State(TypedDict):
files: Annotated[dict[str, str], DeltaChannel(merge_files)]
turn = {"v": 0}
def write_file(state: State) -> dict:
turn["v"] += 1
n = turn["v"]
return {"files": {f"/doc_{n}.txt": f"content for turn {n}"}}
builder = StateGraph(State)
builder.add_node("write_file", write_file)
builder.add_edge(START, "write_file")
saver = InMemorySaver()
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "fs"}}
for _ in range(3):
graph.invoke({"files": {}}, config)
# Checkpoint stores only the sentinel — per-step writes live in checkpoint_writes.
saved = saver.get_tuple(config)
assert saved is not None
cv = saved.checkpoint["channel_values"]["files"]
assert cv is DELTA_SENTINEL
state = graph.get_state(config)
assert state.values["files"] == {
"/doc_1.txt": "content for turn 1",
"/doc_2.txt": "content for turn 2",
"/doc_3.txt": "content for turn 3",
}
# Deletion path must round-trip through writes replay.
def delete_file(state: State) -> dict:
return {"files": {"/doc_1.txt": None}}
builder2 = StateGraph(State)
builder2.add_node("write_file", write_file)
builder2.add_node("delete_file", delete_file)
builder2.add_edge(START, "write_file")
builder2.add_edge("write_file", "delete_file")
turn["v"] = 0
saver2 = InMemorySaver()
graph2 = builder2.compile(checkpointer=saver2)
config2 = {"configurable": {"thread_id": "fs2"}}
graph2.invoke({"files": {}}, config2)
state2 = graph2.get_state(config2)
assert state2.values["files"] == {}
def test_delta_channel_dict_reducer_backwards_compat() -> None:
"""A pre-DeltaChannel dict checkpoint must load as a dict, not be listified."""
def merge_dicts(left: dict, right: dict) -> dict:
return {**left, **right}
spec = _delta_channel_with_type(merge_dicts, dict)
old_value = {"a": 1, "b": 2}
ch = spec.from_checkpoint(old_value)
assert ch.get() == {"a": 1, "b": 2}
# ---------------------------------------------------------------------------
# seed / pre-delta migration
# ---------------------------------------------------------------------------
def test_delta_channel_from_checkpoint_honors_seed() -> None:
"""A non-sentinel value to from_checkpoint is used as the pre-delta seed.
Guards the pre-delta migration path: when the saver's ancestor walk hits
a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs
the post-migration state correctly rather than replaying from empty.
"""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="pre-delta", id="p1")]
ch = spec.from_checkpoint(seed)
ch.replay_writes(
[
("t0", "messages", AIMessage(content="delta-1", id="d1")),
("t1", "messages", HumanMessage(content="delta-2", id="d2")),
]
)
msgs = ch.get()
assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"]
def test_delta_channel_from_checkpoint_seed_without_writes() -> None:
"""Reconstruction at a pre-delta ancestor with no newer deltas returns
just the seed the saver's terminator fired immediately."""
spec = DeltaChannel(add_messages)
seed = [HumanMessage(content="only-snap", id="s1")]
ch = spec.from_checkpoint(seed)
ch.replay_writes([])
assert ch.get() == seed
def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> None:
"""`seed=None` must start replay from None, not from an empty channel.
The DELTA_SENTINEL / MISSING sentinels mean 'no seed'; passing `None`
explicitly should feed None to the reducer as the left operand.
"""
def replace(left, right):
return right
spec = DeltaChannel(replace)
ch = spec.from_checkpoint(None)
ch.replay_writes([("t0", "x", "after")])
# Reducer replaces; seed=None → first write produces "after".
assert ch.get() == "after"
@@ -0,0 +1,398 @@
"""Benchmark: DeltaChannel vs BinaryOperatorAggregate storage and time.
Run directly: python tests/test_delta_channel_benchmark.py
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
Simulates realistic multi-turn conversations with paragraph-length messages
(~100 tokens each) scaling up to 1M-token-equivalent histories.
Token estimates: 1 token 4 chars; each turn 200 tokens (human + AI).
A 1M-token conversation 5,000 turns of realistic messages.
DeltaChannel stores only a zero-byte sentinel in checkpoint_blobs; the actual
write data lives in checkpoint_writes (already stored there). Reconstruction
walks the parent chain and replays writes through the operator O(N) total
storage vs O() for plain add_messages.
"""
from __future__ import annotations
import sys
import time
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from typing_extensions import TypedDict
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import END, StateGraph
from langgraph.graph.message import add_messages
try:
from langgraph.checkpoint.sqlite import SqliteSaver
_SQLITE_AVAILABLE = True
except ImportError:
_SQLITE_AVAILABLE = False
try:
from langgraph.checkpoint.postgres import PostgresSaver
_POSTGRES_AVAILABLE = True
_POSTGRES_URI = (
"postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
)
except ImportError:
_POSTGRES_AVAILABLE = False
# ---------------------------------------------------------------------------
# Realistic message payload (~100 tokens / ~400 chars each)
# ---------------------------------------------------------------------------
_HUMAN_TEMPLATE = (
"I need help understanding the implications of {topic} on our system architecture. "
"Specifically, I'm concerned about how this interacts with our existing {concern} "
"and whether we need to refactor the {component} layer before proceeding. "
"We've had prior incidents in this area and want to be deliberate. "
"What should we prioritize first, and are there known failure modes we should design around from the start?"
)
_AI_TEMPLATE = (
"Great question about {topic}. The key insight here is that {concern} introduces "
"a subtle ordering dependency that most teams overlook until they hit it in production. "
"For your {component} layer specifically, I'd recommend starting with a careful audit "
"of the interface boundaries before making any structural changes. This will give you "
"a clear picture of the blast radius and let you sequence the migration safely."
)
_TOPICS = [
"distributed tracing",
"eventual consistency",
"schema migration",
"backpressure handling",
"idempotency guarantees",
"cache invalidation",
"connection pooling",
"rate limiting",
"circuit breaking",
"observability pipelines",
]
_CONCERNS = [
"concurrency model",
"retry semantics",
"state management",
"error propagation",
"latency budget",
]
_COMPONENTS = [
"persistence",
"routing",
"ingestion",
"aggregation",
"serialization",
]
def _human_content(i: int) -> str:
return _HUMAN_TEMPLATE.format(
topic=_TOPICS[i % len(_TOPICS)],
concern=_CONCERNS[i % len(_CONCERNS)],
component=_COMPONENTS[i % len(_COMPONENTS)],
)
def _ai_content(i: int) -> str:
return _AI_TEMPLATE.format(
topic=_TOPICS[i % len(_TOPICS)],
concern=_CONCERNS[i % len(_CONCERNS)],
component=_COMPONENTS[i % len(_COMPONENTS)],
)
# ---------------------------------------------------------------------------
# State definitions
# ---------------------------------------------------------------------------
class BinaryState(TypedDict):
messages: Annotated[list, add_messages]
class DeltaState(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
# ---------------------------------------------------------------------------
# Graph factory
# ---------------------------------------------------------------------------
def _make_graph(state_cls: type, checkpointer: Any = None) -> Any:
def human_node(state: Any) -> dict:
return {}
def ai_node(state: Any) -> dict:
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]}
g = StateGraph(state_cls)
g.add_node("human", human_node)
g.add_node("ai", ai_node)
g.add_edge("human", "ai")
g.add_edge("ai", END)
g.set_entry_point("human")
return g.compile(checkpointer=checkpointer or MemorySaver())
# ---------------------------------------------------------------------------
# Measurement helpers
# ---------------------------------------------------------------------------
def _total_blob_bytes(saver: MemorySaver) -> int:
total = 0
for (_, _, _, _), (type_tag, blob) in saver.blobs.items():
if blob is not None:
total += len(blob)
return total
def _run_turns(
n_turns: int,
state_cls: type,
checkpointer: Any = None,
) -> tuple[float, float, int]:
"""Run n_turns conversation turns.
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
blob_bytes is -1 for savers without in-memory blob stores (e.g. SQLite).
Read latency is measured as the time to invoke the graph with no new
messages after the full history is built this forces state rehydration.
"""
graph = _make_graph(state_cls, checkpointer)
config = {"configurable": {"thread_id": "bench"}}
t0 = time.perf_counter()
for i in range(n_turns):
graph.invoke(
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
config,
)
write_elapsed = time.perf_counter() - t0
# Measure read/rehydration: get_state forces the channel to rebuild
t1 = time.perf_counter()
for _ in range(5):
graph.get_state(config)
read_elapsed = (time.perf_counter() - t1) / 5
if isinstance(graph.checkpointer, MemorySaver):
blob_bytes = _total_blob_bytes(graph.checkpointer)
else:
blob_bytes = -1
return write_elapsed, read_elapsed, blob_bytes
def _fmt_bytes(n: int) -> str:
if n >= 1_000_000:
return f"{n / 1_000_000:.1f} MB"
if n >= 1_000:
return f"{n / 1_000:.1f} KB"
return f"{n} B"
def _approx_tokens(n_turns: int) -> str:
# ~100 tokens human + ~100 tokens AI per turn
tokens = n_turns * 200
if tokens >= 1_000_000:
return f"~{tokens / 1_000_000:.1f}M tok"
if tokens >= 1_000:
return f"~{tokens / 1_000:.0f}K tok"
return f"~{tokens} tok"
# ---------------------------------------------------------------------------
# Benchmark matrix
# ---------------------------------------------------------------------------
# Turn counts chosen to demonstrate O(N²) vs O(N) storage growth without running too long.
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
TURN_COUNTS = [10, 25, 50, 100, 500]
# Deep-thread counts where add_messages blob storage would exceed 1 GB;
# only DeltaChannel runs here.
DELTA_ONLY_TURN_COUNTS = [1000]
def _checkpointer_factories() -> list[tuple[str, Any]]:
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
return [("InMemory", None)]
def run_benchmark() -> None:
print()
print(
"DeltaChannel vs add_messages (BinaryOperatorAggregate) — checkpoint storage & latency"
)
print("Simulating realistic multi-turn conversations up to ~1M-token histories")
print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)")
print()
checkpointers: list[tuple[str, Any]] = [("InMemory", None)]
if _POSTGRES_AVAILABLE:
try:
import psycopg
psycopg.connect(_POSTGRES_URI).close()
checkpointers.append(("Postgres (plain SELECT)", "postgres"))
except Exception:
pass
for cp_label, cp_hint in checkpointers:
print(f"--- Checkpointer: {cp_label} ---")
_run_benchmark_for_checkpointer(cp_hint)
def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
import contextlib
import tempfile
@contextlib.contextmanager
def _make_saver():
if cp_hint is None:
yield None
elif cp_hint == "postgres":
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
saver.setup()
with saver._cursor() as cur:
cur.execute("DELETE FROM checkpoints WHERE thread_id = 'bench'")
cur.execute(
"DELETE FROM checkpoint_blobs WHERE thread_id = 'bench'"
)
cur.execute(
"DELETE FROM checkpoint_writes WHERE thread_id = 'bench'"
)
yield saver
else:
with tempfile.NamedTemporaryFile(suffix=".db") as f:
with SqliteSaver.from_conn_string(f.name) as saver:
yield saver
rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = []
for turns in TURN_COUNTS:
with _make_saver() as saver:
b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver)
with _make_saver() as saver:
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt))
for turns in DELTA_ONLY_TURN_COUNTS:
with _make_saver() as saver:
d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver)
rows.append((turns, None, d_bytes, None, d_rt, None, d_wt))
# ── Table 1: Storage ─────────────────────────────────────────────────────
W = 64
print("Storage (checkpoint blob bytes)")
print("=" * W)
print(
f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} "
f"{'savings':>8}"
)
print("-" * W)
def _bytes_or_na(v: Any) -> str:
if v is None:
return "n/a"
if v < 0:
return "n/a"
return _fmt_bytes(v)
def _ms_or_na(v: Any) -> str:
return "n/a" if v is None else f"{v * 1000:.1f}ms"
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0:
ratio_str = "n/a"
else:
ratio = b_bytes / d_bytes if d_bytes else float("inf")
ratio_str = f"{ratio:.0f}x"
print(
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} "
f"{ratio_str:>8}"
)
print("=" * W)
print()
# ── Table 2: Read latency ─────────────────────────────────────────────────
print("Read latency (avg of 5 get_state calls)")
print("=" * W)
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
print("-" * W)
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
print(
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
)
print("=" * W)
print()
# ── Table 3: Per-invoke latency (total write_elapsed / turns) ─────────────
print("Per-invoke latency (total graph.invoke time / turns)")
print("=" * W)
print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}")
print("-" * W)
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
def _per(wt: Any) -> str:
if wt is None:
return "n/a"
return f"{(wt / turns) * 1000:.1f}ms"
print(
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{_per(b_wt):>12} {_per(d_wt):>12}"
)
print("=" * W)
print()
print("Legend:")
print(" add_msgs = Annotated[list, add_messages] — O(N²) storage")
print(" delta = Annotated[list, DeltaChannel(add_messages)] — O(N) storage")
print()
# ---------------------------------------------------------------------------
# Pytest entry point
# ---------------------------------------------------------------------------
@pytest.mark.skip(
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
)
def test_delta_channel_benchmark(capsys: Any) -> None:
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
with capsys.disabled():
run_benchmark()
# Correctness assertion: DeltaChannel must use less storage at scale.
for turns in [25, 50]:
_, _, b_bytes = _run_turns(turns, BinaryState)
_, _, d_bytes = _run_turns(turns, DeltaState)
assert d_bytes < b_bytes, (
f"DeltaChannel should use less storage at {turns} turns, "
f"got delta={d_bytes} binary={b_bytes}"
)
# ---------------------------------------------------------------------------
# Script entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
run_benchmark()
sys.exit(0)
@@ -0,0 +1,504 @@
"""Tests for the BinaryOperatorAggregate -> DeltaChannel migration path.
A thread written under `BinaryOperatorAggregate(...)` must keep working
after its annotation is swapped to `DeltaChannel(...)` on the same
checkpointer pre-migration state visible at each *settled* ancestor
checkpoint is preserved, and post-migration writes fold on top through
the reducer.
Mechanism under test: the saver's `_get_channel_writes_history(config,
channel)` walks the parent chain; when it encounters an ancestor whose
`channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it
returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses
it as the base value, and `replay_writes(writes)` folds on-path deltas.
Scenarios covered:
1. **Basic migration (sync + async)**: build pre-migration state with
`BinaryOperatorAggregate`, swap the annotation to `DeltaChannel` on
the same checkpointer, and verify that every settled pre-migration
super-step boundary (`next=('__start__',)`) round-trips exactly
under the delta-channel view.
2. **Time travel into a pre-migration checkpoint** after migration
`graph.get_state(pre_migration_config)` at a settled ancestor
returns the same state as under the binop channel.
3. **Continuing a migrated thread**: driving one more super-step after
migration produces a state that includes the pre-migration settled
prefix plus the new delta write proving `from_checkpoint(seed)` +
`replay_writes` correctly fold post-migration deltas onto the
pre-migration seed.
4. **Base-saver fallback path**: a third-party-style subclass that
removes the optimized `InMemorySaver` override and falls back to
`BaseCheckpointSaver._get_channel_writes_history` must produce the
same result as the optimized path.
5. **Channel-type isolation across threads**: two threads on the same
checkpointer under the delta-channel graph one freshly-started,
one migrated from pre-migration state don't cross-contaminate.
The parent-chain walk is scoped to the thread.
TODO: add postgres variants in the existing `libs/checkpoint-postgres`
test files (different fixture setup; not this file).
"""
from __future__ import annotations
import operator
from typing import Annotated, Any
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import END, START, StateGraph
pytestmark = pytest.mark.anyio
# ---------------------------------------------------------------------------
# Graph factories
#
# A minimal reducer (`operator.add` on lists of str) with a noop node keeps
# state change localized to the HumanMessage-like payload passed through
# `invoke`. That isolates the pre/post-migration parity assertions to
# channel-hydration semantics.
# ---------------------------------------------------------------------------
def _noop(_state: Any) -> dict:
return {}
def _binop_graph(checkpointer: Any) -> Any:
class BinopState(TypedDict):
items: Annotated[list, BinaryOperatorAggregate(list, operator.add)]
return (
StateGraph(BinopState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def _delta_graph(checkpointer: Any) -> Any:
class DeltaState(TypedDict):
items: Annotated[list, DeltaChannel(operator.add)]
return (
StateGraph(DeltaState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def _drive(graph: Any, config: dict, tag: str, n: int) -> None:
for i in range(n):
graph.invoke({"items": [f"{tag}{i}"]}, config)
async def _adrive(graph: Any, config: dict, tag: str, n: int) -> None:
for i in range(n):
await graph.ainvoke({"items": [f"{tag}{i}"]}, config)
def _settled_boundaries(history: list) -> list[tuple[dict, list]]:
"""Return `[(config, items), ...]` for every checkpoint in `history`
whose `next == ('__start__',)` the stable boundaries between invokes.
"""
return [
(s.config, list(s.values.get("items", [])))
for s in history
if s.next == ("__start__",)
]
# ---------------------------------------------------------------------------
# 1. Basic migration (sync + async)
# ---------------------------------------------------------------------------
def test_basic_migration_preserves_pre_migration_state() -> None:
"""Build state under `BinaryOperatorAggregate`, migrate to
`DeltaChannel` on the same checkpointer, and verify that every
settled pre-migration super-step boundary round-trips exactly.
Settled boundaries (`next=('__start__',)`) are the stable hydration
targets for the migration path: writes that produced the NEXT
super-step are kept as `pending_writes` on the ancestor, so walking
from a descendant finds the ancestor's blob as the seed and
reconstructs the correct state.
"""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "basic-sync"}}
# Pre-migration: accumulate items across 3 invokes.
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
assert len(pre_boundaries) >= 2, "expected multiple settled boundaries"
# Migrate: swap the annotation on the same checkpointer.
delta = _delta_graph(checkpointer)
for cfg, items in pre_boundaries:
snap = delta.get_state(cfg)
assert list(snap.values.get("items", [])) == items, (
f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: "
f"expected {items}, got {snap.values.get('items', [])}"
)
async def test_basic_migration_preserves_pre_migration_state_async() -> None:
"""Async variant of the basic migration scenario."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "basic-async"}}
binop = _binop_graph(checkpointer)
await _adrive(binop, config, "u", 3)
pre_history = [s async for s in binop.aget_state_history(config)]
pre_boundaries = _settled_boundaries(pre_history)
assert len(pre_boundaries) >= 2
delta = _delta_graph(checkpointer)
for cfg, items in pre_boundaries:
snap = await delta.aget_state(cfg)
assert list(snap.values.get("items", [])) == items, (
f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}"
)
# ---------------------------------------------------------------------------
# 2. Time travel into a pre-migration checkpoint after migration
# ---------------------------------------------------------------------------
def test_time_travel_into_pre_migration_checkpoint() -> None:
"""After migration, `graph.get_state(pre_migration_config)` at a
settled ancestor returns the state as stored at that point."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "time-travel"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
assert pre_boundaries, "no settled ancestors to time-travel to"
delta = _delta_graph(checkpointer)
# Pick the oldest non-empty boundary — a long distance to walk back.
non_empty = [(cfg, items) for cfg, items in pre_boundaries if items]
assert non_empty, "expected at least one non-empty boundary"
target_cfg, expected_items = non_empty[-1]
snap = delta.get_state(target_cfg)
assert list(snap.values.get("items", [])) == expected_items
# ---------------------------------------------------------------------------
# 3. Continuing a migrated thread: deltas fold onto pre-migration seed
# ---------------------------------------------------------------------------
def test_continuing_migrated_thread_folds_deltas_on_seed() -> None:
"""Resume a pre-migration settled ancestor via `invoke(None, cfg)`
under the delta-channel graph. Since the pre-migration checkpoint
has an existing `pending_writes` entry (the input for the NEXT
super-step), re-running from that ancestor reproduces the same
post-ancestor state as the original binop run.
This proves the seed-terminator + write-replay pipeline works
end-to-end across the migration boundary.
"""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "continue"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 2)
# Pick the oldest settled boundary with non-empty state.
pre_boundaries = _settled_boundaries(list(binop.get_state_history(config)))
target_cfg, seed_items = next(
(cfg, items) for cfg, items in reversed(pre_boundaries) if items
)
assert seed_items, "need a non-empty seed boundary"
# Migrate and resume from the pre-migration ancestor. `invoke(None,
# cfg)` replays the pending writes staged at `cfg` under the new
# channel; the reducer folds those deltas onto the seed.
delta = _delta_graph(checkpointer)
result = delta.invoke(None, target_cfg)
# The resumed state must include the pre-migration seed items in order.
result_items = list(result.get("items", []))
for idx, prefix_item in enumerate(seed_items):
assert result_items[idx] == prefix_item, (
f"pre-migration seed item at {idx} not preserved: "
f"got {result_items[: idx + 1]}, expected {seed_items}"
)
# ---------------------------------------------------------------------------
# 4. Base-saver fallback path
# ---------------------------------------------------------------------------
class _ThirdPartyStyleSaver(InMemorySaver):
"""Simulates a third-party saver that inherits the reference
`_get_channel_writes_history` implementation from
`BaseCheckpointSaver` rather than overriding it.
We rebind the two methods to the base-class versions (via MRO) so
the fallback path is exercised even though the storage layer is
still the in-memory one.
"""
# MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...]
_get_channel_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined]
)
_aget_channel_writes_history = ( # type: ignore[assignment]
InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined]
)
def test_base_saver_fallback_matches_optimized_override() -> None:
"""The reference `BaseCheckpointSaver` implementation must produce
the same migration behavior as the optimized `InMemorySaver`
override. We drive the same migration scenario through both savers
and assert per-snapshot parity in the delta-channel view."""
# Fast path: optimized InMemorySaver override.
fast_saver = InMemorySaver()
fast_config = {"configurable": {"thread_id": "fast"}}
fast_binop = _binop_graph(fast_saver)
_drive(fast_binop, fast_config, "u", 3)
fast_delta = _delta_graph(fast_saver)
fast_history = [
(s.next, list(s.values.get("items", [])))
for s in fast_delta.get_state_history(fast_config)
]
# Slow path: base-class fallback.
slow_saver = _ThirdPartyStyleSaver()
slow_config = {"configurable": {"thread_id": "slow"}}
slow_binop = _binop_graph(slow_saver)
_drive(slow_binop, slow_config, "u", 3)
slow_delta = _delta_graph(slow_saver)
slow_history = [
(s.next, list(s.values.get("items", [])))
for s in slow_delta.get_state_history(slow_config)
]
assert slow_history == fast_history, (
"base-saver fallback should match optimized-override behavior; "
f"fast={fast_history}, slow={slow_history}"
)
# ---------------------------------------------------------------------------
# 5. Thread isolation under mixed-generation storage
# ---------------------------------------------------------------------------
def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None:
"""Two threads sharing a checkpointer — one migrated from
pre-migration state, one freshly-started under DeltaChannel must
maintain independent state. The parent-chain walk in
`_get_channel_writes_history` must be scoped to the target thread.
"""
checkpointer = InMemorySaver()
migrated_cfg = {"configurable": {"thread_id": "migrated"}}
fresh_cfg = {"configurable": {"thread_id": "fresh"}}
# Thread A: pre-migration build-up.
binop = _binop_graph(checkpointer)
_drive(binop, migrated_cfg, "m", 2)
# Thread B: fresh delta-channel run.
delta = _delta_graph(checkpointer)
_drive(delta, fresh_cfg, "f", 2)
# Thread A: migrate and confirm its state is anchored in its own
# thread's pre-migration history (tag 'm'), never mixing in tag 'f'.
migrated_boundaries = _settled_boundaries(
list(delta.get_state_history(migrated_cfg))
)
assert migrated_boundaries, "migrated thread has no settled boundaries"
for _, items in migrated_boundaries:
for it in items:
assert it.startswith("m"), (
f"migrated thread leaked item from other thread: {it}"
)
# Thread B: settled boundaries must only contain 'f' tags.
fresh_boundaries = _settled_boundaries(list(delta.get_state_history(fresh_cfg)))
assert fresh_boundaries, "fresh thread has no settled boundaries"
for _, items in fresh_boundaries:
for it in items:
assert it.startswith("f"), (
f"fresh thread leaked item from migrated thread: {it}"
)
# ---------------------------------------------------------------------------
# 6. Tip-of-pre-migration hydration: the latest checkpoint from a binop-run
# thread has a real accumulated value in its own `channel_values["items"]`.
# When hydrated under the delta-channel graph via `get_state(config)` with no
# `checkpoint_id`, the short-circuit must use that value directly instead of
# walking ancestors (which would skip the tip's own blob).
# ---------------------------------------------------------------------------
def test_tip_of_pre_migration_hydrates_directly() -> None:
"""`graph.get_state(config)` at the latest (pre-migration) checkpoint
returns the full accumulated list stored in that checkpoint's own
`channel_values`. The hydration must not walk ancestors past it."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "tip-sync"}}
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 3)
binop_tip = binop.get_state(config)
expected_items = list(binop_tip.values.get("items", []))
assert expected_items == ["u0", "u1", "u2"], (
f"sanity: pre-migration tip should accumulate all 3 items, got {expected_items}"
)
delta = _delta_graph(checkpointer)
snap = delta.get_state(config)
assert list(snap.values.get("items", [])) == expected_items, (
f"tip hydration mismatch: expected {expected_items}, "
f"got {snap.values.get('items', [])}"
)
async def test_tip_of_pre_migration_hydrates_directly_async() -> None:
"""Async variant of the tip-of-pre-migration hydration scenario."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "tip-async"}}
binop = _binop_graph(checkpointer)
await _adrive(binop, config, "u", 3)
binop_tip = await binop.aget_state(config)
expected_items = list(binop_tip.values.get("items", []))
assert expected_items == ["u0", "u1", "u2"]
delta = _delta_graph(checkpointer)
snap = await delta.aget_state(config)
assert list(snap.values.get("items", [])) == expected_items, (
f"async tip hydration mismatch: expected {expected_items}, "
f"got {snap.values.get('items', [])}"
)
# ---------------------------------------------------------------------------
# 7. `update_state` after migration writes a real value to the new
# checkpoint's `channel_values` (not a sentinel). Hydration must use it
# directly — the ancestor walk would skip this blob and return stale state.
# ---------------------------------------------------------------------------
def test_update_state_after_migration_uses_written_value() -> None:
"""After migrating and running at least one post-migration super-step
(so the thread's tip has a `DELTA_SENTINEL`), `update_state` writes a
concrete value to a new checkpoint's `channel_values`. `get_state`
must reflect that concrete value."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "update-state"}}
# Pre-migration: accumulate a little state.
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 2)
# Migrate and run one more super-step so the tip is a post-migration
# checkpoint with `DELTA_SENTINEL` in its own `channel_values`.
delta = _delta_graph(checkpointer)
delta.invoke({"items": ["post"]}, config)
# `update_state` writes a concrete value into a new checkpoint's blob
# via the reducer against the hydrated prior state.
delta.update_state(config, {"items": ["x", "y"]})
snap = delta.get_state(config)
updated_items = list(snap.values.get("items", []))
# Must include the "x","y" update; without the hydration fix, the
# update_state-written blob would be skipped in favor of an ancestor
# walk, and the update values would disappear.
assert "x" in updated_items and "y" in updated_items, (
f"update_state values missing from snapshot: {updated_items}"
)
# The "x","y" items should be folded onto the prior accumulated state,
# not stand alone. This verifies the update-written blob is used
# directly by `get_state` (no ancestor walk past it).
assert len(updated_items) >= 4, (
f"update_state snapshot should preserve pre-update state, got {updated_items}"
)
assert updated_items[-2:] == ["x", "y"], (
f"update_state deltas should be at the tail, got {updated_items}"
)
# ---------------------------------------------------------------------------
# 8. Fork from an `update_state` checkpoint: a new run branched off the
# update_state-produced checkpoint must see that checkpoint's concrete
# `channel_values` as its base, with new deltas folded on top.
# ---------------------------------------------------------------------------
def test_fork_from_update_state_checkpoint() -> None:
"""Branching a new run from the checkpoint produced by `update_state`
must use that checkpoint's concrete blob as the base. Additional
deltas from the forked run fold onto it through the reducer."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "fork"}}
# Pre-migration build-up, then migrate and add one post-migration step.
binop = _binop_graph(checkpointer)
_drive(binop, config, "u", 2)
delta = _delta_graph(checkpointer)
delta.invoke({"items": ["post"]}, config)
# Apply `update_state` and capture the returned config (references
# the new checkpoint produced by the update).
update_cfg = delta.update_state(config, {"items": ["x", "y"]})
update_snap = delta.get_state(update_cfg)
base_items = list(update_snap.values.get("items", []))
assert "x" in base_items and "y" in base_items, (
f"update_state values missing from snapshot: {base_items}"
)
assert base_items[-2:] == ["x", "y"], (
f"sanity: update_state deltas should be at the tail, got {base_items}"
)
# Fork: invoke from the update_state checkpoint with a new delta.
forked = delta.invoke({"items": ["fork0"]}, update_cfg)
forked_items = list(forked.get("items", []))
# The fork must see the update_state-written blob as its base (not
# walk past it), and the new delta must fold on top of it.
assert forked_items[: len(base_items)] == base_items, (
f"fork lost update_state base: base={base_items}, forked={forked_items}"
)
assert forked_items[-1] == "fork0", f"fork delta not appended: {forked_items}"
+118
View File
@@ -5,6 +5,7 @@ import langchain_core
import pytest
from langchain_core.messages import (
AIMessage,
AIMessageChunk,
AnyMessage,
HumanMessage,
RemoveMessage,
@@ -338,6 +339,123 @@ def test_remove_all_messages():
]
def test_fast_path_preserves_format_openai():
"""Pure-append fast path must still apply the `langchain-openai` formatter."""
left = [HumanMessage(content="prior", id="1")]
right = [
AIMessage(
content=[
{
"type": "tool_use",
"name": "foo",
"input": {"bar": "baz"},
"id": "t1",
}
],
id="2",
)
]
result = add_messages(left, right, format="langchain-openai")
assert isinstance(result[0], HumanMessage)
assert result[0].content == "prior"
assert isinstance(result[1], AIMessage)
# formatter collapses the tool_use content block into `tool_calls`
assert result[1].content == ""
assert len(result[1].tool_calls) == 1
assert result[1].tool_calls[0]["name"] == "foo"
assert result[1].tool_calls[0]["args"] == {"bar": "baz"}
assert result[1].tool_calls[0]["id"] == "t1"
def test_fast_path_rejects_invalid_format():
"""Pure-append fast path must validate the `format` arg like the slow path."""
left = [HumanMessage(content="prior", id="1")]
right = [AIMessage(content="new", id="2")]
with pytest.raises(ValueError, match="Unrecognized format="):
add_messages(left, right, format="bogus") # type: ignore[arg-type]
def test_left_starting_with_chunk_is_normalized():
"""Opt-1 guard: a `BaseMessageChunk` at left[0] must trigger full conversion."""
chunk = AIMessageChunk(content="chunk", id="c1")
result = add_messages([chunk], [HumanMessage(content="h", id="h1")])
assert len(result) == 2
# chunk must be converted to a non-chunk message
assert type(result[0]).__name__ == "AIMessage"
assert result[0].id == "c1"
assert result[1].id == "h1"
def test_left_as_dicts_is_normalized():
"""Opt-1 guard: dicts at left[0] must trigger full conversion."""
left = [{"role": "user", "content": "hi", "id": "d1"}]
right = [AIMessage(content="reply", id="a1")]
result = add_messages(left, right)
assert len(result) == 2
assert isinstance(result[0], HumanMessage)
assert result[0].id == "d1"
assert result[0].content == "hi"
def test_left_as_tuples_is_normalized():
"""Opt-1 guard: tuple-form messages must trigger full conversion."""
left = [("user", "hi")]
right = [AIMessage(content="reply", id="a1")]
result = add_messages(left, right)
assert len(result) == 2
assert isinstance(result[0], HumanMessage)
# id is auto-assigned
assert isinstance(result[0].id, str) and UUID(result[0].id, version=4)
def test_left_first_msg_missing_id_is_normalized():
"""Opt-1 guard: a BaseMessage without an id at left[0] falls to the else branch."""
left = [HumanMessage(content="hi")] # no id
right = [AIMessage(content="reply", id="a1")]
result = add_messages(left, right)
assert len(result) == 2
# left's id must have been auto-assigned
assert isinstance(result[0].id, str) and UUID(result[0].id, version=4)
def test_duplicate_ids_in_right_with_nonempty_left():
"""Opt-2 guard: intra-right duplicate ids must take slow path (dedup kept)."""
left = [HumanMessage(content="prior", id="1")]
right = [
AIMessage(content="first", id="2"),
AIMessage(content="second", id="2"),
]
result = add_messages(left, right)
assert len(result) == 2
assert result[0].id == "1"
assert result[1].id == "2"
assert result[1].content == "second"
def test_right_with_none_ids_pure_append():
"""Fast path still correct when right entries start with id=None (fresh uuids assigned)."""
left = [HumanMessage(content="prior", id="1")]
right = [AIMessage(content="a"), AIMessage(content="b")]
result = add_messages(left, right)
assert len(result) == 3
assert result[0].id == "1"
for m in result[1:]:
assert isinstance(m.id, str) and UUID(m.id, version=4)
# fresh uuids must be distinct
assert result[1].id != result[2].id
def test_fast_path_returns_fresh_list():
"""Fast path must return a new list object (not mutate or alias left)."""
left = [HumanMessage(content="prior", id="1")]
right = [AIMessage(content="new", id="2")]
result = add_messages(left, right)
assert result is not left
# left must be untouched
assert len(left) == 1
assert left[0].id == "1"
def test_push_messages_in_graph():
class MessagesState(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
+187
View File
@@ -9400,3 +9400,190 @@ def test_fork_does_not_apply_pending_writes(
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121}
async def test_delta_channel_end_to_end_inmemory() -> None:
"""Full graph run: DeltaChannel accumulates correctly across multiple turns."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-test-1"}}
# Turn 1
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
# Turn 2
graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config)
# Turn 3
graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# 3 human + 3 AI = 6 total
assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}"
assert msgs[0].content == "hello"
assert msgs[2].content == "world"
assert msgs[4].content == "bye"
assert msgs[1].content == "reply-1"
assert msgs[3].content == "reply-3"
assert msgs[5].content == "reply-5"
async def test_delta_channel_time_travel() -> None:
"""Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
counter = {"n": 0}
def respond(state: State) -> dict:
counter["n"] += 1
return {
"messages": [
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
]
}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
saver = InMemorySaver()
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "diff-time-travel"}}
# Run 2 turns: h1→ai-1, h2→ai-2
graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config)
graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config)
# Find the checkpoint after turn 1 (2 messages: h1 + ai-1)
history = list(graph.get_state_history(config))
after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2)
assert len(after_turn1.values["messages"]) == 2
assert after_turn1.values["messages"][0].content == "h1"
assert after_turn1.values["messages"][1].content == "ai-1"
# Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N)
# NOT 5 messages (turn-2 deltas must not bleed into the resumed run)
result = graph.invoke(
{"messages": [HumanMessage(content="h3", id="h3")]},
after_turn1.config,
)
msgs = result["messages"]
# Should be: h1, ai-1, h3, ai-N — 4 messages total
assert len(msgs) == 4, (
f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}"
)
assert msgs[0].content == "h1"
assert msgs[1].content == "ai-1"
assert msgs[2].content == "h3"
async def test_delta_channel_remove_message_end_to_end() -> None:
"""RemoveMessage inside a DeltaChannel graph must persist and reload correctly."""
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
return {"messages": [AIMessage(content="reply", id="ai-1")]}
def delete_first(state: State) -> dict:
# removes the first message
return {"messages": [RemoveMessage(id=state["messages"][0].id)]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_node("delete_first", delete_first)
builder.add_edge(START, "respond")
builder.add_edge("respond", "delete_first")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-remove-test"}}
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# h1 was removed, only ai-1 should remain
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
assert msgs[0].id == "ai-1"
# A subsequent turn must reconstruct from the checkpoint correctly
graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed
# more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2
# just assert h1 is still gone
assert all(m.id != "h1" for m in msgs), (
"h1 should still be absent after second turn"
)
async def test_delta_channel_update_by_id_end_to_end() -> None:
"""Updating a message by ID via DeltaChannel must persist and reload correctly."""
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels._delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def update_msg(state: State) -> dict:
# re-send h1 with updated content
return {"messages": [HumanMessage(content="updated", id="h1")]}
builder = StateGraph(State)
builder.add_node("update_msg", update_msg)
builder.add_edge(START, "update_msg")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-update-id-test"}}
graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
assert msgs[0].content == "updated"
assert msgs[0].id == "h1"
# Second turn: verify the updated state is the base for further accumulation
graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
ids = [m.id for m in msgs]
assert "h1" in ids # h1 persists (updated, not duplicated)
assert "h2" in ids
assert ids.count("h1") == 1, "h1 must not be duplicated"
+1 -3
View File
@@ -1161,9 +1161,7 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
assert called == ["step_a", "ask_human"]
# Resume with explicit head checkpoint_id in config
head_checkpoint_id = graph.get_state(config).config["configurable"][
"checkpoint_id"
]
head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"]
called.clear()
resume_config = {
"configurable": {
+28 -8
View File
@@ -82,6 +82,7 @@ from langchain_core.tools.base import (
_is_injected_arg_type,
get_all_basemodel_annotations,
)
from langgraph._internal._constants import CONF, CONFIG_KEY_READ
from langgraph._internal._runnable import RunnableCallable
from langgraph.errors import GraphBubbleUp
from langgraph.graph.message import REMOVE_ALL_MESSAGES
@@ -800,7 +801,7 @@ class ToolNode(RunnableCallable):
# Construct ToolRuntime instances at the top level for each tool call
tool_runtimes = []
for call, cfg in zip(tool_calls, config_list, strict=False):
state = self._extract_state(input)
state = self._extract_state(input, cfg)
tool_runtime = ToolRuntime(
state=state,
tool_call_id=call["id"],
@@ -835,7 +836,7 @@ class ToolNode(RunnableCallable):
# Construct ToolRuntime instances at the top level for each tool call
tool_runtimes = []
for call, cfg in zip(tool_calls, config_list, strict=False):
state = self._extract_state(input)
state = self._extract_state(input, cfg)
tool_runtime = ToolRuntime(
state=state,
tool_call_id=call["id"],
@@ -1273,18 +1274,37 @@ class ToolNode(RunnableCallable):
return None
def _extract_state(
self, input: list[AnyMessage] | dict[str, Any] | BaseModel
self,
input: list[AnyMessage] | dict[str, Any] | BaseModel,
config: RunnableConfig,
) -> list[AnyMessage] | dict[str, Any] | BaseModel:
"""Extract state from input, handling ToolCallWithContext if present.
"""Extract state from input.
Args:
input: The input which may be raw state or ToolCallWithContext.
Three input shapes:
Returns:
The actual state to pass to wrap_tool_call wrappers.
- `ToolCallWithContext` dict legacy Send payload carrying an inlined
state snapshot; return `input["state"]`.
- list of `ToolCall` dicts new Send payload with no inlined state;
hydrate state from channels via `CONFIG_KEY_READ`.
- regular graph state (dict/list/BaseModel) return `input` as-is.
"""
if isinstance(input, dict) and input.get("__type") == "tool_call_with_context":
return input["state"]
if (
isinstance(input, list)
and input
and isinstance(input[-1], dict)
and input[-1].get("type") == "tool_call"
):
read = config.get(CONF, {}).get(CONFIG_KEY_READ)
if read is None:
return {}
# Pregel installs CONFIG_KEY_READ as
# `functools.partial(local_read, scratchpad, channels, managed, task)`.
# Match the previous inlined-state contract by reading channels only;
# managed values have their own injection path (`ToolRuntime.context`).
channels = read.args[1]
return cast("dict[str, Any]", read(list(channels), False))
return input
def _inject_tool_args(
+92
View File
@@ -1320,6 +1320,98 @@ async def test_state_extraction_with_tool_call_with_context_async() -> None:
assert "tool_call" not in state_seen[0]
def _config_with_channel_read(
channel_values: dict[str, object],
store: BaseStore | None = None,
) -> RunnableConfig:
"""Build a config that mimics `CONFIG_KEY_READ` as Pregel installs it.
Pregel always installs a `functools.partial(local_read, scratchpad,
channels, managed, task)`, and `ToolNode` introspects that partial to
learn channel names. The stub matches the shape: partial whose second and
third positional args are `channels` and `managed` mappings.
"""
import functools
channels_stub = {k: None for k in channel_values}
managed_stub: dict[str, object] = {}
# Shape matches pregel's real partial:
# functools.partial(local_read, scratchpad, channels, managed, task)
def _read(scratchpad, channels, managed, task, select, fresh): # noqa: ARG001
if isinstance(select, str):
return channel_values[select]
return {k: channel_values[k] for k in select if k in channel_values}
read = functools.partial(_read, None, channels_stub, managed_stub, None)
cfg = _create_config_with_runtime(store)
cfg["configurable"]["__pregel_read"] = read
return cfg
def test_list_form_send_hydrates_state_from_channel_read() -> None:
"""Send('tools', [tool_call]) with no inlined state should hydrate
ToolRuntime.state from CONFIG_KEY_READ (full state read)."""
state_seen = []
def state_inspector_handler(
request: ToolCallRequest,
execute: Callable[[ToolCallRequest], ToolMessage | Command],
) -> ToolMessage | Command:
state_seen.append(request.state)
return execute(request)
channel_values = {
"messages": [AIMessage("from channels")],
"files": {"/a.md": "body"},
}
tool_node = ToolNode([add], wrap_tool_call=state_inspector_handler)
tool_call: ToolCall = {
"name": "add",
"args": {"a": 1, "b": 2},
"id": "call_1",
"type": "tool_call",
}
tool_node.invoke([tool_call], config=_config_with_channel_read(channel_values))
assert len(state_seen) == 1
got = state_seen[0]
assert got == channel_values
assert "messages" in got and "files" in got
async def test_list_form_send_hydrates_state_async() -> None:
state_seen = []
def state_inspector_handler(
request: ToolCallRequest,
execute: Callable[[ToolCallRequest], ToolMessage | Command],
) -> ToolMessage | Command:
state_seen.append(request.state)
return execute(request)
channel_values = {"messages": [AIMessage("from channels")], "files": {}}
tool_node = ToolNode([add], wrap_tool_call=state_inspector_handler)
tool_call: ToolCall = {
"name": "add",
"args": {"a": 1, "b": 2},
"id": "call_1",
"type": "tool_call",
}
await tool_node.ainvoke(
[tool_call], config=_config_with_channel_read(channel_values)
)
assert len(state_seen) == 1
assert state_seen[0] == channel_values
def test_tool_call_request_is_frozen() -> None:
"""Test that ToolCallRequest raises deprecation warnings on direct attribute reassignment."""
tool_call: ToolCall = {"name": "add", "args": {"a": 1, "b": 2}, "id": "call_1"}
+137
View File
@@ -0,0 +1,137 @@
# Delta-channel reconstruction: query strategy benchmark
**Branch:** `delta-channel-writes-based`
**Question (Nuno):** Is the recursive CTE the right query shape for reconstructing a delta channel inside `get_tuple`, or would a plain `SELECT WHERE` be cheaper even though it returns more rows?
**Answer:** Plain `SELECT WHERE` wins at every realistic depth. The recursion isn't the problem — the JSON-expression join inside the CTE is.
## Setup
- Postgres 16 on `localhost:5441` (the `compose-postgres.yml` instance, run directly without docker for this round)
- Single delta channel `messages`, one write per checkpoint, `DELTA_SENTINEL` blob per checkpoint
- Linear chain (`branch=1`) and 5-way branching at every step (`branch=5`) — branching is the case where plain over-fetches sibling rows
- Median of 20 timed runs after 3 warmups, fresh psycopg cursor per strategy
- Bench script: `bench_get_tuple_strategies.py` at repo root
Three strategies compared:
| name | roundtrips | shape |
|------|-----------|-------|
| `cte` | 1 | Current prod: recursive CTE walks ancestors, LEFT JOINs writes + blobs |
| `plain` | 3 | Nuno's suggestion: thread-wide `SELECT WHERE` per table, Python walks parent chain and filters |
| `cte+narrow` | 2 | CTE returns ancestor IDs only, then one `UNION ALL` of writes + blobs filtered by `ANY(ids)` |
## Results (ms per get_tuple, median of 20)
```
depth branch cte plain cte+narrow rows_cte rows_plain plain/cte
10 1 0.14ms 0.26ms 0.24ms 9 30 1.89x
10 5 0.21ms 0.23ms 0.17ms 9 110 1.11x
50 1 0.89ms 0.27ms 0.33ms 49 150 0.31x
50 5 2.35ms 0.66ms 0.51ms 49 550 0.28x
200 1 11.61ms 0.78ms 1.30ms 199 600 0.07x
200 5 34.79ms 2.33ms 3.07ms 199 2200 0.07x
1000 1 274.60ms 2.59ms 13.29ms 999 3000 0.01x
1000 5 856.01ms 10.14ms 15.31ms 999 11000 0.01x
```
Lower is better. `plain/cte < 1` means plain is faster.
### Headline numbers
- depth 50: plain is **3x** faster
- depth 200: plain is **15x** faster
- depth 1000: plain is **~100x** faster
- Branching makes plain over-fetch (3000 rows → 11000 rows at d=1000), but it remains ~85x faster than the CTE
## Why the CTE collapses
`EXPLAIN (ANALYZE, BUFFERS)` of the CTE at depth 1000 (linear). Excerpt with the load-bearing nodes:
```
Sort ... actual time=137.798..137.827 rows=999
CTE ancestors
-> Recursive Union ... actual time=0.005..2.443 rows=999
^^^^^^
recursion is 2.4 ms — fine
-> Nested Loop Left Join ... actual time=2.676..137.529 rows=999
Join Filter: (cw.checkpoint_id = a.cid)
Rows Removed by Join Filter: 998001
^^^^^^^
999 ancestors x ~1000 writes
-> Nested Loop Left Join ... actual time=2.669..85.061 rows=999
Join Filter: (bl.version = ((c.checkpoint -> 'channel_versions'::text) ->> bl.channel))
Rows Removed by Join Filter: 998001
^^^^^^^
same quadratic blow-up on the blob join
```
Two pathological things are happening:
1. **The blob join filter is on a JSON expression**: `bl.version = (c.checkpoint -> 'channel_versions' ->> bl.channel)`. The planner cannot push this into an index lookup, so it materializes `checkpoint_blobs` for the thread and does a nested-loop comparison against every ancestor — a Cartesian product that grows as `O(ancestors × blobs_in_thread)`.
2. **The writes join is similar**: writes for the thread are materialized once, then nested-loop joined against ancestors with a `Join Filter` rather than a hash/merge join over the indexed `checkpoint_id`.
At depth 1000 that's **~2 million rows evaluated, 99.9% of them discarded**. The recursion itself is a rounding error.
For comparison, the plain Q1 (`SELECT … FROM checkpoints WHERE thread_id=? AND checkpoint_ns=?`) at depth 1000:
```
Seq Scan on checkpoints ... actual time=0.012..0.121 rows=1000
Execution Time: 0.140 ms
```
A simple seq scan over 57 buffers. Q2 and Q3 follow the same shape and complete in well under 1 ms each.
## Crossover and remote-DB reasoning
- Pure local Postgres: plain wins from depth ~30 onward; CTE wins by fractions of a ms below that
- Remote Postgres at ~5 ms RTT adds ~10 ms to plain (3 roundtrips vs 1). Crossover shifts to ~depth 30. Above that, the CTE's quadratic SQL cost still dominates the RTT savings.
There is no realistic conversation depth where the CTE wins on a remote DB. At depth 200+ (anything resembling a real multi-turn agent run) plain is faster regardless of network.
## Recommendation
**Switch to plain SELECT WHERE, one delta channel at a time.**
Three indexed queries per delta channel:
```sql
-- Q1: parent chain + per-checkpoint version of this channel
SELECT checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> 'channel_name' AS ver
FROM checkpoints
WHERE thread_id = ? AND checkpoint_ns = ?;
-- Q2: writes for this channel, anywhere in the thread
SELECT checkpoint_id, type, blob, task_id, idx
FROM checkpoint_writes
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
-- Q3: blobs for this channel, anywhere in the thread
SELECT version, type, blob
FROM checkpoint_blobs
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
```
Python then:
- Builds `parent_of: dict[cid, parent_cid]` from Q1
- Walks from target's parent newest → oldest
- Filters Q2 rows by `ancestor_set`, processes oldest → newest, applies overwrite-terminator
- Picks seed blob via the per-ancestor `ver` map, terminates at first non-sentinel blob
All O(n) on n = thread checkpoints, with tight constants (dict lookups). No recursion, no JSON-expression joins, no quadratic plans.
If the 3-roundtrip cost ever shows up on remote-DB benchmarks, fold Q2 + Q3 into one `UNION ALL` to get back to 2 roundtrips. Bench says it isn't worth the SQL complexity right now.
## Bonus: code simplification from single-channel scope
Multi-channel reconstruction in the current `_reconstruct_delta_channels_cur` carries:
- `rows_by_cid` nested dicts, keyed by cid then channel
- `seen_blob: set[(cid, channel)]` and `seen_write: set[(cid, channel, task_id, idx)]` dedup
- `collected: dict[channel, list]`, `done: set[channel]`, `seeds: dict[channel, value]`
- Inner `for ch in channels_list` loops and an early-exit `if len(done) == len(channels_list)`
Single-channel collapses these to a single list, a single bool, and one `Optional[Any]`. Roughly half the Python in that function, plus an obvious shape for splitting pure post-processing into `base.py` so sync and async stop duplicating it.
If multi-channel coalescing turns out to matter later, it can come back as a SQL-level optimization without re-introducing the bookkeeping in Python.