Compare commits

..
Author SHA1 Message Date
Sydney RunkleandClaude Sonnet 4.6 186d045947 perf(sdk): lazy-load langgraph_sdk top-level exports
Converting __init__.py to use __getattr__ defers loading of Auth,
get_client/get_sync_client, Encryption, and EncryptionContext until
first access. This removes ~13ms from langgraph.runtime import time,
since importing BaseUser from langgraph_sdk.auth.types no longer
triggers the full client/auth/encryption module graph.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-21 12:56:32 -04:00
Sydney RunkleandClaude Sonnet 4.6 c459079e52 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-21 11:03:58 -04:00
Sydney RunkleandClaude Sonnet 4.6 ffda8b5472 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-21 10:58:09 -04:00
Sydney RunkleandClaude Sonnet 4.6 374eebcd65 chore: apply format/lint fixes across checkpoint, checkpoint-postgres, prebuilt
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-21 10:50:27 -04:00
Sydney RunkleandClaude Sonnet 4.6 350182ed18 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-21 10:48:20 -04:00
Sydney RunkleandClaude Sonnet 4.6 4c2ce5c8a9 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-21 10:40:00 -04:00
Sydney Runkle bae7486565 test(channels): replace unsupported-saver raise test with fallback assembly test 2026-04-21 10:06:51 -04:00
Sydney Runkle b082584500 feat(postgres): remove _load_diff_chains; add get_channel_blob / aget_channel_blob 2026-04-21 10:06:13 -04:00
Sydney Runkle 503071c2aa feat(memory): implement get_channel_blob; remove diff handling from _load_blobs 2026-04-21 10:05:08 -04:00
Sydney Runkle c9913afef2 feat(pregel): wire DeltaChannel assembly into loop; pass checkpoint_id to after_checkpoint 2026-04-21 10:04:19 -04:00
Sydney Runkle 9fd6374302 feat(pregel): add _assemble_delta_channels helpers for universal DeltaChannel support 2026-04-21 09:37:49 -04:00
Sydney RunkleandClaude Sonnet 4.6 e6c065739f feat(channels): DeltaChannel tracks checkpoint_id; emits prev_checkpoint_id in DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-21 09:34:44 -04:00
Sydney RunkleandClaude Sonnet 4.6 e18f8fff2b feat(serde): diff type encodes prev_checkpoint_id; loads_typed returns DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-21 09:29:58 -04:00
Sydney Runkle 65438610e8 docs(checkpoint): expand aget_channel_blob docstring for parity 2026-04-21 09:28:21 -04:00
Sydney Runkle 4ebebd686a chore: add .worktrees/ to .gitignore 2026-04-21 09:27:47 -04:00
Sydney Runkle fca3f6d919 feat(checkpoint): DeltaValue uses prev_checkpoint_id; add get_channel_blob stubs 2026-04-21 09:27:22 -04:00
Sydney Runkle 599afd7585 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-21 07:56:43 -04:00
Sydney Runkle c0e6062bfb more tests 2026-04-20 12:53:54 -04:00
Sydney RunkleandClaude Sonnet 4.6 056d3143ff chore: format/lint fixes for rehydrate_every benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 16:31:26 -04:00
Sydney RunkleandClaude Sonnet 4.6 1da43d412b 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-17 16:30:17 -04:00
Sydney RunkleandClaude Sonnet 4.6 df56b7cdf6 test(channels): add DiffChannel vs BinaryOperatorAggregate storage/time benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 16:21:28 -04:00
Sydney RunkleandClaude Sonnet 4.6 566a3150b2 chore: format and lint fixes for DiffChannel implementation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 16:20:21 -04:00
Sydney RunkleandClaude Sonnet 4.6 1bb1811fbd 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-17 16:05:58 -04:00
Sydney RunkleandClaude Sonnet 4.6 dd7f21e3ff 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-17 15:56:19 -04:00
Sydney RunkleandClaude Sonnet 4.6 d4e1efa1f6 feat(checkpoint/postgres): diff chain reconstruction in _load_blobs (sync)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 15:37:18 -04:00
Sydney RunkleandClaude Sonnet 4.6 dba1987c9b test(pregel): strengthen DiffChannel time-travel and reply assertions
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 15:01:50 -04:00
Sydney RunkleandClaude Sonnet 4.6 9fb0493ac5 feat(pregel): call after_checkpoint hook when loading and saving channels
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 14:56:42 -04:00
Sydney RunkleandClaude Sonnet 4.6 1cb057e6bc 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-17 14:54:58 -04:00
Sydney RunkleandClaude Sonnet 4.6 d76127fbbf feat(checkpoint/memory): chain-traverse diff blobs in _load_blobs
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 14:52:34 -04:00
Sydney Runkle 94853fb14c fix(channels): align DiffChannel.is_available with BinaryOperatorAggregate 2026-04-17 14:51:03 -04:00
Sydney RunkleandClaude Sonnet 4.6 e6fab22f0c 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-17 14:35:46 -04:00
Sydney Runkle ea644f413d feat(channels): add no-op after_checkpoint hook to BaseChannel 2026-04-17 14:30:04 -04:00
Sydney RunkleandClaude Sonnet 4.6 3f86b1485d 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-17 14:24:41 -04:00
Sydney RunkleandClaude Sonnet 4.5 9e40dee07f 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-17 14:21:40 -04:00
Sydney RunkleandClaude Sonnet 4.6 4d1f4086eb feat(checkpoint): add DiffDelta and DiffChainValue protocol types
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 14:18:09 -04:00
Sydney RunkleandClaude Sonnet 4.6 afcf6c03dd docs: add DiffChannel implementation plan
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 14:14:50 -04:00
Sydney RunkleandClaude Sonnet 4.6 4b303ceb39 docs: add DiffChannel incremental checkpoint storage design spec
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-17 14:00:41 -04:00
50 changed files with 3576 additions and 3428 deletions
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,405 @@
# DiffChannel: Incremental Checkpoint Storage for Append-Style Reducers
**Date:** 2026-04-17
**Status:** Approved for implementation
**Scope:** `libs/checkpoint`, `libs/langgraph`, `libs/checkpoint-postgres`
---
## Motivation
LangGraph checkpoints today store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, this means each checkpoint blob contains the entire conversation history. Storage cost grows O(N²) in the number of turns: step 1 stores 1 message, step 100 stores 100 messages, step 1000 stores 1000 messages. For long-running agentic conversations with high-token messages this is untenable.
The fix is to store only the **delta** (new writes) per step, reconstructing the full accumulated value at load time by replaying the chain. This is an opt-in mechanism — existing graphs are unaffected.
---
## Non-Goals
- **Compaction / materialized snapshots**: deferred. Load cost stays O(N) blob fetches but those fetches are batched into a single query — acceptable for now.
- **SQLite saver support**: SQLite stores all channel values inline in one row (no per-channel blob table). Deferred to a follow-up.
- **Automatic migration** of existing `BinaryOperatorAggregate` channels: users opt in explicitly. Old checkpoints load correctly via the backwards-compatibility path in `from_checkpoint`.
---
## Architecture Overview
```
User state definition
└── Annotated[list[AnyMessage], DiffChannel(add_messages)]
Write path (per superstep)
DiffChannel.update() — apply operator, accumulate writes in _pending
DiffChannel.checkpoint() — return DiffDelta(delta=_pending, prev_version=_base_version)
serde.dumps_typed() — serialize DiffDelta as ("diff", msgpack_bytes)
saver.put() — store blob at (thread_id, ns, "messages", version_N)
DiffChannel.after_checkpoint(version_N) — advance _base_version, clear _pending
Read path (on graph load or time-travel)
saver.get_tuple() — fetch current-version blob per channel
saver._load_blobs() — detect "diff" type → follow chain to reconstruct DiffChainValue
DiffChannel.from_checkpoint(DiffChainValue) — replay deltas with operator → full list
DiffChannel.after_checkpoint(version_N) — set _base_version for next write
```
The pregel layer (`_checkpoint.py`, `_loop.py`) is unchanged except for two small additions to call the new `after_checkpoint` hook. The saver public interface (`BaseCheckpointSaver`) gains no new methods. All chain-following logic lives inside each saver's private `_load_blobs`.
---
## New Protocol Types
**Location:** `libs/checkpoint/langgraph/checkpoint/base/__init__.py`
Two dataclasses form the contract between `DiffChannel` and savers:
```python
@dataclass
class DiffDelta:
"""Returned by DiffChannel.checkpoint(). Written to the blob store."""
delta: list[Any] # raw writes passed to update() this step
prev_version: str | None # version of the previous diff blob; None = chain root
```
```python
@dataclass
class DiffChainValue:
"""Passed to DiffChannel.from_checkpoint(). Assembled by _load_blobs()."""
base: list[Any] | None # starting accumulated value (None = empty start)
deltas: list[list[Any]] # write-sets ordered oldest → newest
```
`DiffDelta` lives in the checkpoint base package (not the channel module) so savers can import it without creating a circular dependency. `DiffChainValue` is there for the same reason.
---
## `BaseChannel.after_checkpoint()` Hook
**Location:** `libs/langgraph/langgraph/channels/base.py`
```python
def after_checkpoint(self, version: Any) -> None:
"""Called after checkpoint() (with the new version) and after from_checkpoint()
(with the current version). No-op by default; DiffChannel overrides."""
pass
```
This is a **non-abstract, no-op default** — fully backwards compatible. All existing channels inherit it silently. It is NOT in the abstract interface.
---
## `DiffChannel[V]`
**Location:** `libs/langgraph/langgraph/channels/diff.py` (new file)
### Internal state
| Attribute | Type | Description |
|---|---|---|
| `value` | `list[V]` | Full accumulated value (the reconstructed list) |
| `operator` | `Callable` | The binary reducer (e.g. `add_messages`) |
| `_pending` | `list[Any]` | Raw writes accumulated since last `after_checkpoint` call |
| `_base_version` | `str \| None` | Version this channel was last checkpointed at (= `prev_version` for next delta) |
| `_overwritten` | `bool` | True if an `Overwrite` was applied since last `after_checkpoint`; makes next blob a chain root |
### `update(values)`
Mirrors `BinaryOperatorAggregate.update()` with two additions:
1. For each non-Overwrite value: apply `self.operator(self.value, value)` as before; **also append the raw incoming value to `self._pending`**.
2. For an `Overwrite(v)` value: set `self.value = v`; set `self._pending = list(v)` (full value becomes the new delta); set `self._overwritten = True`.
The key: `_pending` stores the **incoming writes** (what was passed to `update()`), not the diff of `self.value`. This is important because `add_messages` handles removal and update-by-ID — replaying the writes with `operator` during reconstruction applies that logic correctly.
### `checkpoint()`
```python
def checkpoint(self) -> DiffDelta:
return DiffDelta(
delta=self._pending[:],
prev_version=None if self._overwritten else self._base_version,
)
```
- Normal step: `prev_version = self._base_version` → chain link
- After Overwrite: `prev_version = None` → chain root (reconstruction stops here and uses `delta` as the full base value)
Returns `DiffDelta`, never the raw accumulated list. The serde handles serialization.
### `from_checkpoint(checkpoint)`
```python
def from_checkpoint(self, checkpoint) -> Self:
new = DiffChannel(self.typ, self.operator)
new.key = self.key
if checkpoint is MISSING:
new.value = []
elif isinstance(checkpoint, DiffChainValue):
accumulated = checkpoint.base or []
for step_writes in checkpoint.deltas:
# Mirror update() exactly: apply each write individually so operator
# semantics (e.g. add_messages ID-based removal) are respected.
for write in step_writes:
accumulated = new.operator(accumulated, write)
new.value = accumulated
elif isinstance(checkpoint, DiffDelta):
# Unsupported saver: _load_blobs returned a raw DiffDelta instead of
# assembling a DiffChainValue. Raise rather than silently losing history.
raise ValueError(
"DiffChannel received a raw DiffDelta from the checkpoint saver. "
"Your saver does not support incremental channel storage. "
"Use InMemorySaver or PostgresSaver."
)
else:
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
new.value = checkpoint
new._pending = []
new._base_version = None # set by the subsequent after_checkpoint() call
return new
```
The operator is available on `self` (the channel spec) so reconstruction is correct for any reducer — the saver never needs to know about `add_messages`.
`_pending` stores **individual writes** (each `value` from `update()`'s `values` sequence), so each `step_writes` list in `DiffChainValue.deltas` is replayed write-by-write — identical to the `update()` loop.
### `after_checkpoint(version)`
```python
def after_checkpoint(self, version: Any) -> None:
if version != self._base_version:
self._base_version = version
self._pending = []
self._overwritten = False
```
No-op when `version == self._base_version` (channel wasn't updated this step — blob was not written). Clears `_pending` and advances `_base_version` when the channel was actually checkpointed.
### Opt-in API
```python
from langgraph.channels.diff import DiffChannel
class State(TypedDict):
messages: Annotated[list[AnyMessage], DiffChannel(add_messages)]
```
`StateGraph` already handles `BaseChannel` instances as annotation metadata — `DiffChannel` inherits this without any changes to `StateGraph`.
---
## Serde Extension
**Location:** `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py`
Add one branch to `dumps_typed` (before the `else` msgpack fallback), using the existing module-level `_msgpack_enc` so message ext-types (Pydantic v2, etc.) are handled correctly:
```python
elif isinstance(obj, DiffDelta):
return "diff", _msgpack_enc({"d": obj.delta, "p": obj.prev_version})
```
Add one branch to `loads_typed` so savers can decode diff blobs without importing `ormsgpack` directly:
```python
elif type_ == "diff":
return ormsgpack.unpackb(
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
# returns {"d": [writes...], "p": prev_version_str_or_none}
```
Savers call `serde.loads_typed(("diff", raw_bytes))` to decode a diff blob into `{"d": ..., "p": ...}`, then check `type_tag == "diff"` to trigger chain traversal. The serde layer is the only place that knows about `ormsgpack`.
---
## Saver Changes
### InMemorySaver
**`put()` — `libs/checkpoint/langgraph/checkpoint/memory/__init__.py`**
No change needed. The existing `self.serde.dumps_typed(values[k])` call already handles `DiffDelta` via the new serde branch above, storing it as `("diff", bytes)`.
**`_load_blobs()` — same file**
After checking `vv[0] != "empty"`, add a branch for `"diff"` before calling `serde.loads_typed`:
```python
def _load_blobs(self, thread_id, checkpoint_ns, versions):
channel_values = {}
diff_channels = {} # channel_name -> current_version for diff channels
for k, v in versions.items():
kk = (thread_id, checkpoint_ns, k, v)
if kk not in self.blobs:
continue
type_tag, blob_bytes = self.blobs[kk]
if type_tag == "diff":
diff_channels[k] = v # handle below
elif type_tag != "empty":
channel_values[k] = self.serde.loads_typed((type_tag, blob_bytes))
for k, current_version in diff_channels.items():
# Follow chain: newest → oldest, then reverse
chain_deltas = []
base = None
version = current_version
while version is not None:
kk = (thread_id, checkpoint_ns, k, version)
if kk not in self.blobs:
break
type_tag, blob_bytes = self.blobs[kk]
if type_tag == "diff":
# Use serde so we don't need to import ormsgpack directly
payload = self.serde.loads_typed((type_tag, blob_bytes))
chain_deltas.append(payload["d"])
version = payload["p"] # prev_version; None = root
else:
# Old non-diff blob encountered: treat as base accumulated value
base = self.serde.loads_typed((type_tag, blob_bytes))
break
chain_deltas.reverse()
channel_values[k] = DiffChainValue(base=base, deltas=chain_deltas)
return channel_values
```
Each blob lookup is O(1) on the dict. Total: N dict lookups for a chain of depth N. Memory usage is identical to loading a single full-list blob (same total bytes, split across N entries).
### PostgresSaver
**`_load_blobs()` — `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py`**
The existing `SELECT_SQL` fetches one blob per channel via a JOIN. After running that query, detect any `"diff"` channels in the result and issue one additional range query:
```python
def _load_blobs(self, blob_values):
if not blob_values:
return {}
result = {}
diff_channels = {} # channel_name -> current_version (as str)
for k, t, v in blob_values:
channel = k.decode()
type_tag = t.decode()
if type_tag == "diff":
# Decode via serde — no direct ormsgpack import needed
payload = self.serde.loads_typed((type_tag, v))
diff_channels[channel] = payload # store for chain fetch
elif type_tag != "empty":
result[channel] = self.serde.loads_typed((type_tag, v))
if diff_channels:
result.update(self._load_diff_chains(diff_channels))
return result
```
`_load_diff_chains` issues one SQL query per diff channel (typically just `messages`):
```sql
SELECT version, type, blob
FROM checkpoint_blobs
WHERE thread_id = %s
AND checkpoint_ns = %s
AND channel = %s
AND version <= %s
ORDER BY version ASC
```
In Python, iterate rows in ascending version order: if `type = "diff"`, accumulate the delta; if any other type is encountered, treat it as the base accumulated value and stop. Return `DiffChainValue(base=..., deltas=[...])`.
This results in **at most 2 queries total** for a graph with one `DiffChannel` — existing behaviour for all other channels is unchanged.
**`put()` / `_dump_blobs()`**
No change needed. `_dump_blobs` calls `self.serde.dumps_typed(v)` for each channel value in `new_versions`. When `v` is a `DiffDelta`, the serde produces `("diff", bytes)` which is stored as `type = "diff"` in `checkpoint_blobs`. The `ON CONFLICT DO NOTHING` semantics are preserved.
### SQLite
Deferred. `SqliteSaver` stores the entire checkpoint as a single serialized row — it has no per-channel blob table. Supporting `DiffChannel` on SQLite would require adding a new blobs table, which is a separate migration tracked separately.
---
## Pregel Layer Changes
### `channels_from_checkpoint` — `libs/langgraph/langgraph/pregel/_checkpoint.py`
After constructing each channel from its checkpoint value, call `after_checkpoint` so the channel records its current version:
```python
channels = {}
for k, v in channel_specs.items():
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
ch.after_checkpoint(checkpoint["channel_versions"].get(k))
channels[k] = ch
return channels, managed_specs
```
Existing channels get the no-op `after_checkpoint`. `DiffChannel` uses it to set `_base_version`.
### `PregelLoop._put_checkpoint` — `libs/langgraph/langgraph/pregel/_loop.py`
After `create_checkpoint(self.checkpoint, self.channels, self.step, ...)` returns and `do_checkpoint is True` and `self.channels is not None`, iterate channels and notify:
```python
if do_checkpoint and self.channels:
for k, ch in self.channels.items():
ch.after_checkpoint(self.checkpoint["channel_versions"].get(k))
```
This is called after `create_checkpoint` updates `self.checkpoint["channel_versions"]`, so `get(k)` returns the new version for updated channels and the old version for unchanged ones. `DiffChannel.after_checkpoint` only clears `_pending` when `version != _base_version`, so unchanged channels are no-ops.
---
## Backwards Compatibility
| Scenario | Behaviour |
|---|---|
| Existing graph using `add_messages` (BinaryOperatorAggregate) | Unaffected — no code changes, no data migration |
| New graph with `DiffChannel`, loading old checkpoint blobs | `from_checkpoint` receives a plain `list` → used directly as accumulated value |
| `DiffChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
| `DiffChannel` with `SqliteSaver` | `from_checkpoint` receives a raw `DiffDelta` (SqliteSaver stores channel_values inline), raises `ValueError` with a clear message pointing to supported savers |
| Time-travel / fork to past checkpoint | Chain traversal uses the version at that checkpoint → reconstruction is correct |
| `update_state` | Treated as a normal step: writes are deltas chained to history |
| `Overwrite` value | Resets chain: next blob has `prev_version=None`; reconstruction starts fresh |
---
## Testing Strategy
1. **Unit tests for `DiffChannel`** (`libs/langgraph/tests/`):
- `update` → `checkpoint` → `after_checkpoint` → `checkpoint` lifecycle (2 steps, verify delta isolation)
- `from_checkpoint(DiffChainValue)` correctly replays multi-step chains using the operator
- `from_checkpoint(plain_list)` backwards-compat path
- `Overwrite` creates a root blob (`prev_version=None`) and reconstruction ignores prior chain
- `after_checkpoint` no-ops when version is unchanged
2. **Integration tests with `InMemorySaver`** (`libs/langgraph/tests/`):
- 10-step conversation: verify final loaded state equals full accumulated messages
- Time-travel: fork to step 5, verify only messages 1–5 are present
- Mixed graph: some channels `BinaryOperatorAggregate`, one `DiffChannel` — both reconstruct correctly
3. **Serde tests** (`libs/checkpoint/tests/`):
- `DiffDelta` round-trips through `dumps_typed` / saver storage
- Old `"msgpack"` blob for a channel → `DiffChannel.from_checkpoint` handles it
4. **Postgres integration tests** (`libs/checkpoint-postgres/tests/`):
- Range query reconstructs correct full list after N steps
- Time-travel to checkpoint M reconstructs correct list of M messages
---
## Files Changed
| File | Change |
|---|---|
| `libs/checkpoint/langgraph/checkpoint/base/__init__.py` | Add `DiffDelta`, `DiffChainValue` dataclasses |
| `libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py` | Add `"diff"` branch in `dumps_typed` |
| `libs/checkpoint/langgraph/checkpoint/memory/__init__.py` | Chain traversal in `_load_blobs` |
| `libs/langgraph/langgraph/channels/base.py` | Add no-op `after_checkpoint` method |
| `libs/langgraph/langgraph/channels/diff.py` | **New file** — `DiffChannel` implementation |
| `libs/langgraph/langgraph/channels/__init__.py` | Export `DiffChannel` |
| `libs/langgraph/langgraph/pregel/_checkpoint.py` | Call `after_checkpoint` in `channels_from_checkpoint` |
| `libs/langgraph/langgraph/pregel/_loop.py` | Call `after_checkpoint` after `create_checkpoint` |
| `libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py` | Range-query chain reconstruction in `_load_blobs` |
@@ -4,17 +4,15 @@ import threading
from collections import defaultdict
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, cast
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,
)
@@ -25,11 +23,7 @@ from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
_DeltaCombinedRow,
)
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
Conn = _internal.Conn # For backward compatibility
@@ -436,47 +430,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`.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; the ancestor walk runs in Python.
"""
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"]
def get_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Look up a channel blob by checkpoint ID + channel via checkpoint_blobs."""
with self._cursor() as cur:
cur.execute(
SELECT_DELTA_COMBINED_SQL,
"""
SELECT cb.type, cb.blob
FROM checkpoint_blobs cb
WHERE cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (
SELECT checkpoint->'channel_versions'->>%s
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
)
""",
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
channel,
checkpoint_id,
),
)
rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
rows=cast("list[_DeltaCombinedRow]", rows),
)
row = cur.fetchone()
if row is None:
return NotImplemented
return self.serde.loads_typed((row["type"], row["blob"]))
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
@@ -490,6 +479,13 @@ class PostgresSaver(BasePostgresSaver):
including its configuration, metadata, parent checkpoint (if any),
and pending writes.
"""
with self._cursor() as cur:
channel_values = self._load_blobs(
value["channel_values"],
thread_id=value["thread_id"],
checkpoint_ns=value["checkpoint_ns"],
cur=cur,
)
return CheckpointTuple(
{
"configurable": {
@@ -502,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"],
@@ -4,17 +4,15 @@ import asyncio
from collections import defaultdict
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, cast
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,
)
@@ -25,11 +23,7 @@ from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
_DeltaCombinedRow,
)
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
Conn = _ainternal.Conn # For backward compatibility
@@ -397,45 +391,42 @@ 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`.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; rows are assembled by the shared pure helper on
`BasePostgresSaver`.
"""
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 def aget_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Async look up of a channel blob by checkpoint ID + channel name."""
async with self._cursor() as cur:
await cur.execute(
SELECT_DELTA_COMBINED_SQL,
"""
SELECT cb.type, cb.blob
FROM checkpoint_blobs cb
WHERE cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (
SELECT checkpoint->'channel_versions'->>%s
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
)
""",
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
channel,
checkpoint_id,
),
)
rows = await cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
rows=cast("list[_DeltaCombinedRow]", rows),
)
row = await cur.fetchone()
if row is None:
return NotImplemented
return self.serde.loads_typed((row["type"], row["blob"]))
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
@@ -449,11 +440,19 @@ 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"],
}
},
@@ -461,15 +460,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"],
}
}
@@ -4,16 +4,13 @@ import random
import warnings
from collections.abc import Sequence
from importlib.metadata import version as get_version
from typing import Any, TypedDict, cast
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
@@ -156,62 +153,6 @@ INSERT_CHECKPOINT_WRITES_SQL = """
"""
class _DeltaCombinedRow(TypedDict, total=False):
"""One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables).
Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are
relevant for that kind; irrelevant columns are NULL and typed as `None`.
"""
_kind: str # always present: "p", "w", or "b"
# checkpoint row ("p")
checkpoint_id: str | None
parent_checkpoint_id: str | None
ver: str | None
# write / blob rows ("w", "b")
type: str | None
blob: bytes | None
# write row only ("w")
task_id: str | None
idx: int | None
# blob row only ("b")
version: str | None
# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints,
# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs
# in Python in `_build_delta_channel_writes_history`.
#
# Parameter order: (channel, thread_id, checkpoint_ns,
# thread_id, checkpoint_ns, channel,
# thread_id, checkpoint_ns, channel)
SELECT_DELTA_COMBINED_SQL = """
SELECT 'p'::text AS _kind,
checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver,
NULL::text AS type,
NULL::bytea AS blob,
NULL::text AS task_id,
NULL::int AS idx,
NULL::text AS version
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
UNION ALL
SELECT 'w',
checkpoint_id, NULL, NULL,
type, blob, task_id, idx, NULL
FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
UNION ALL
SELECT 'b',
NULL, NULL, NULL,
type, blob, NULL, NULL, version
FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
"""
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
@@ -244,92 +185,22 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
)
def _load_blobs(
self, blob_values: list[tuple[bytes, bytes, bytes]]
self,
blob_values: list[tuple[bytes, bytes, bytes]],
*,
thread_id: str = "",
checkpoint_ns: str = "",
cur: Any = None,
) -> 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"
}
def _build_delta_channel_writes_history(
self,
*,
channel: str,
target_id: str,
rows: Sequence[_DeltaCombinedRow],
) -> _ChannelWritesHistory:
"""Reconstruct one delta channel's history from the combined UNION ALL rows.
Pure data transform shared by sync (`PostgresSaver`) and async
(`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL`
and feed the tagged 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] = {}
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
blob_by_ver: dict[str, tuple[str, bytes]] = {}
for r in rows:
kind = r["_kind"]
if kind == "p":
cid = cast(str, r["checkpoint_id"])
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
elif kind == "w":
cid = cast(str, r["checkpoint_id"])
writes_by_cid.setdefault(cid, []).append(
cast(
"tuple[str, bytes, str, int]",
(r["type"], r["blob"], r["task_id"], r["idx"]),
)
)
else: # kind == "b"
blob_by_ver[cast(str, r["version"])] = cast(
"tuple[str, bytes]", (r["type"], r["blob"])
)
# newest write first per ancestor (task_id DESC, idx DESC)
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
ancestors: list[str] = []
cur_cid: str | None = parent_of.get(target_id)
while cur_cid is not None:
ancestors.append(cur_cid)
cur_cid = parent_of.get(cur_cid)
if not ancestors:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
collected: list[PendingWrite] = [] # newest first; reversed at the end
for cid in ancestors:
# Collect writes first — they encode the transition FROM this
# ancestor's state to its child's and must be included even if
# this ancestor is also the seed checkpoint.
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))
# Then check seed terminator.
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)
collected.reverse() # oldest → newest
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
result: dict[str, Any] = {}
for k, t, v in blob_values:
channel = k.decode()
type_tag = t.decode()
if type_tag != "empty":
result[channel] = self.serde.loads_typed((type_tag, v))
return result
def _dump_blobs(
self,
+1 -1
View File
@@ -12,7 +12,7 @@ readme = "README.md"
license = "MIT"
license-files = ['LICENSE']
dependencies = [
"langgraph-checkpoint>=4.0.3,<5.0.0",
"langgraph-checkpoint>=2.1.2,<5.0.0",
"orjson>=3.11.5",
"psycopg>=3.2.0",
"psycopg-pool>=3.2.0",
+4 -4
View File
@@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values(
load_checkpoint_tuple = saver._load_checkpoint_tuple
async def patched_load_checkpoint_tuple(value):
def patched_load_checkpoint_tuple(value):
value["checkpoint"].pop("channel_values", None)
return await load_checkpoint_tuple(value)
return load_checkpoint_tuple(value)
monkeypatch.setattr(
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
@@ -385,11 +385,11 @@ async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
@@ -1,9 +1,10 @@
from __future__ import annotations
import copy
import dataclasses
import logging
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import (
from typing import ( # noqa: UP035
Any,
Generic,
Literal,
@@ -18,9 +19,6 @@ 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,
@@ -33,6 +31,24 @@ V = TypeVar("V", int, float, str)
PendingWrite = tuple[str, str, Any]
@dataclasses.dataclass
class DeltaValue:
"""Returned by DeltaChannel.checkpoint(). Represents one step's writes."""
delta: list[Any]
prev_checkpoint_id: (
str | None
) # ID of checkpoint containing previous blob; None = chain root
@dataclasses.dataclass
class DeltaChainValue:
"""Passed to DeltaChannel.from_checkpoint(). Assembled by the pregel layer."""
base: list[Any] | None # starting accumulated value; None = start from empty
deltas: list[list[Any]] # per-step write-sets, ordered oldest → newest
logger = logging.getLogger(__name__)
@@ -124,30 +140,6 @@ 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.
@@ -486,103 +478,41 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Pure storage read used by `_get_channel_writes_history`.
def get_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Look up a single channel blob by checkpoint ID + channel name.
Must return the same value as `get_tuple` but must NOT trigger channel
reconstruction; otherwise the channel-hydration path would re-enter
`_get_channel_writes_history`. Override only if `get_tuple` itself
performs channel hydration.
Returns NotImplemented if this saver does not support efficient
per-channel-version blob lookup. The pregel layer will fall back to
get_tuple() traversal in that case.
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
should override this for O(1) performance.
"""
return self.get_tuple(config)
return NotImplemented
async def _aget_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Async version of `_get_tuple_raw`. See docstring there."""
return await self.aget_tuple(config)
async def aget_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Look up a single channel blob by checkpoint ID + channel name (async).
def _get_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""**Experimental.** Query one channel's writes along the parent chain.
Returns NotImplemented if this saver does not support efficient
per-channel-version blob lookup. The pregel layer will fall back to
aget_tuple() traversal in that case.
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.
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
should override this for O(1) performance.
"""
collected: list[PendingWrite] = [] # newest first; reversed at the end
target_tuple = self._get_tuple_raw(config)
cursor_config: RunnableConfig | None = (
target_tuple.parent_config if target_tuple else None
)
while cursor_config is not None:
tup = self._get_tuple_raw(cursor_config)
if tup is None:
break
# Collect this ancestor's writes FIRST — they encode the
# transition from this ancestor's state to its child's, so
# they must be included whether or not this ancestor is the
# seed terminator.
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)
# Seed terminator: any non-sentinel blob on an ancestor
# establishes the reconstruction base. 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)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
async def _aget_channel_writes_history(
self, config: RunnableConfig, channel: str
) -> _ChannelWritesHistory:
"""Async version of `_get_channel_writes_history`. See docstring there."""
collected: list[PendingWrite] = []
target_tuple = await self._aget_tuple_raw(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_raw(cursor_config)
if tup is None:
break
if tup.pending_writes:
for write in reversed(tup.pending_writes):
if write[1] != channel:
continue
collected.append(write)
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)
cursor_config = tup.parent_config
collected.reverse()
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
return NotImplemented
def get_next_version(self, current: V | None, channel: None) -> V:
"""Generate the next version ID for a channel.
@@ -14,20 +14,16 @@ from typing import Any
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,
)
from langgraph.checkpoint.serde.types import _DeltaSnapshot
logger = logging.getLogger(__name__)
@@ -125,114 +121,50 @@ 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]:
result: dict[str, Any] = {}
for k, ver in versions.items():
kk = (thread_id, checkpoint_ns, k, ver)
channel_values: dict[str, Any] = {}
for k, v in versions.items():
kk = (thread_id, checkpoint_ns, k, v)
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
if vv[0] != "empty":
channel_values[k] = self.serde.loads_typed(vv)
return channel_values
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", "")
def get_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Fast-path blob lookup: checkpoint → channel version → blob."""
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:
if isinstance(blob_value, _DeltaSnapshot):
# Step-based snapshot: the blob is state AT this
# ancestor, but the ancestor's pending_writes
# encode the NEXT step's transition and are NOT
# subsumed by the snapshot — collect them first.
step_writes = self.writes.get(
(thread_id, checkpoint_ns, cp_id), {}
)
for (_task_id, _idx), (
tid,
ch,
serialized,
_,
) in sorted(step_writes.items(), reverse=True):
if ch != channel:
continue
collected.append(
(tid, ch, self.serde.loads_typed(serialized))
)
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
# Pre-delta blob: state AT this ancestor already
# subsumes its pending_writes — skip them.
collected.reverse()
return _ChannelWritesHistory(
seed=blob_value, writes=collected
)
entry = ns_storage.get(checkpoint_id)
if entry is None:
return NotImplemented
checkpoint = self.serde.loads_typed(entry[0])
version = checkpoint["channel_versions"].get(channel)
if version is None:
return NotImplemented
kk = (thread_id, checkpoint_ns, channel, version)
if kk not in self.blobs:
return NotImplemented
vv = self.blobs[kk]
if vv[0] == "empty":
return NotImplemented
return self.serde.loads_typed(vv)
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)
async def aget_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
return self.get_channel_blob(thread_id, checkpoint_ns, checkpoint_id, channel)
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
"""Get a checkpoint tuple from the in-memory storage.
@@ -80,6 +80,8 @@ SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
("langgraph.types", "Overwrite"),
("langgraph.store.base", "Item"),
("langgraph.store.base", "GetOp"),
# DeltaChannel checkpoint value type
("langgraph.checkpoint.base", "DeltaValue"),
}
)
@@ -33,39 +33,24 @@ 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 (
DELTA_SENTINEL,
SendProtocol,
_DeltaSentinel,
_DeltaSnapshot,
)
from langgraph.checkpoint.serde.types import 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""
logger = logging.getLogger(__name__)
# Dedup log warnings across process lifetime; cap bounds state if types are
# dynamically generated (also acts as a circuit breaker on warning volume).
# Dedup is best-effort: racing threads may each emit once for the same key,
# and warnings are silently dropped once _MAX_WARNED_TYPES is reached.
_MAX_WARNED_TYPES = 1000
_warned_unregistered_types: set[tuple[str, str]] = set()
_warned_blocked_types: set[tuple[str, str]] = set()
def _is_delta_value(obj: Any) -> bool:
from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep
def _warn_once(
seen: set[tuple[str, str]], key: tuple[str, str], msg: str, *args: object
) -> None:
if key in seen or len(seen) >= _MAX_WARNED_TYPES:
return
seen.add(key)
logger.warning(msg, *args)
return isinstance(obj, DeltaValue)
class JsonPlusSerializer(SerializerProtocol):
@@ -260,6 +245,8 @@ class JsonPlusSerializer(SerializerProtocol):
return "bytes", obj
elif isinstance(obj, bytearray):
return "bytearray", obj
elif _is_delta_value(obj):
return "delta", _msgpack_enc({"d": obj.delta, "c": obj.prev_checkpoint_id})
else:
try:
return "msgpack", _msgpack_enc(obj)
@@ -282,6 +269,13 @@ class JsonPlusSerializer(SerializerProtocol):
return ormsgpack.unpackb(
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
elif type_ == "delta":
from langgraph.checkpoint.base import DeltaValue # lazy import
raw = ormsgpack.unpackb(
data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
return DeltaValue(delta=raw["d"], prev_checkpoint_id=raw.get("c"))
elif self.pickle_fallback and type_ == "pickle":
return pickle.loads(data_)
else:
@@ -297,16 +291,10 @@ EXT_METHOD_SINGLE_ARG = 3
EXT_PYDANTIC_V1 = 4
EXT_PYDANTIC_V2 = 5
EXT_NUMPY_ARRAY = 6
EXT_DELTA_SNAPSHOT = 7
EXT_DELTA_SENTINEL = 8
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
if isinstance(obj, _DeltaSnapshot):
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
elif isinstance(obj, _DeltaSentinel):
return ormsgpack.Ext(EXT_DELTA_SENTINEL, b"")
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
return ormsgpack.Ext(
EXT_PYDANTIC_V2,
_msgpack_enc(
@@ -561,9 +549,7 @@ def _create_msgpack_ext_hook(
"name": name,
}
)
_warn_once(
_warned_unregistered_types,
key,
logger.warning(
"Deserializing unregistered type %s.%s from checkpoint. "
"This will be blocked in a future version. "
"Set LANGGRAPH_STRICT_MSGPACK=true to block now, or add "
@@ -585,9 +571,7 @@ def _create_msgpack_ext_hook(
"name": name,
}
)
_warn_once(
_warned_blocked_types,
key,
logger.warning(
"Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
"Add to allowed_msgpack_modules to allow: [(%r, %r)]",
module,
@@ -620,15 +604,7 @@ def _create_msgpack_ext_hook(
return False
def ext_hook(code: int, data: bytes) -> Any:
if code == EXT_DELTA_SENTINEL:
return DELTA_SENTINEL
elif code == EXT_DELTA_SNAPSHOT:
return _DeltaSnapshot(
ormsgpack.unpackb(
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
)
)
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
if code == EXT_CONSTRUCTOR_SINGLE_ARG:
try:
tup = ormsgpack.unpackb(
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
@@ -1,7 +1,6 @@
from collections.abc import Sequence
from typing import (
Any,
NamedTuple,
Protocol,
TypeVar,
runtime_checkable,
@@ -15,39 +14,6 @@ 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()
class _DeltaSnapshot(NamedTuple):
"""Snapshot blob for a DeltaChannel with finite snapshot_frequency.
Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code.
The ancestor walk in `_get_channel_writes_history` terminates when it
encounters this type (any non-sentinel blob stops the walk).
`from_checkpoint` reconstructs the channel value directly from `.value`
without replaying writes — the snapshot IS the accumulated state.
"""
value: Any
Value = TypeVar("Value", covariant=True)
Update = TypeVar("Update", contravariant=True)
C = TypeVar("C")
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint"
version = "4.1.0a2"
version = "4.0.2"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
requires-python = ">=3.10"
-9
View File
@@ -29,8 +29,6 @@ from langgraph.checkpoint.serde.jsonplus import (
EXT_METHOD_SINGLE_ARG,
JsonPlusSerializer,
_msgpack_enc,
_warned_blocked_types,
_warned_unregistered_types,
)
@@ -104,13 +102,6 @@ def test_msgpack_method_pathlib_blocked_encrypted_strict(
class TestEncryptedSerializerMsgpackAllowlist:
"""Test msgpack allowlist behavior through EncryptedSerializer."""
@pytest.fixture(autouse=True)
def _reset_warned_types(self) -> None:
# Warning dedup state is process-global; reset per-test so each case
# sees a fresh slate and assertions about warning emission are stable.
_warned_unregistered_types.clear()
_warned_blocked_types.clear()
def test_safe_types_no_warning(self, caplog: pytest.LogCaptureFixture) -> None:
"""Test safe types deserialize without warnings through encryption."""
serde = _make_encrypted_serde()
+24 -22
View File
@@ -35,8 +35,6 @@ from langgraph.checkpoint.serde.jsonplus import (
JsonPlusSerializer,
_msgpack_enc,
_msgpack_ext_hook_to_json,
_warned_blocked_types,
_warned_unregistered_types,
)
from langgraph.store.base import Item
@@ -582,14 +580,6 @@ def test_msgpack_safe_types_no_warning(caplog: pytest.LogCaptureFixture) -> None
assert result is not None
@pytest.fixture(autouse=True)
def _reset_warned_types() -> None:
# Warning dedup state is process-global; reset per-test so each case sees
# a fresh slate and assertions about warning emission are stable.
_warned_unregistered_types.clear()
_warned_blocked_types.clear()
def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
"""Pydantic models not in allowlist should log warning but still deserialize."""
current = _lg_msgpack.STRICT_MSGPACK_ENABLED
@@ -605,12 +595,6 @@ def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) ->
assert "unregistered type" in caplog.text.lower()
assert "allowed_msgpack_modules" in caplog.text
assert result == obj
# Second deserialization of the same type should NOT produce another warning
caplog.clear()
result2 = serde.loads_typed(dumped)
assert "unregistered type" not in caplog.text.lower()
assert result2 == obj
_lg_msgpack.STRICT_MSGPACK_ENABLED = current
@@ -655,6 +639,7 @@ def test_msgpack_allowlist_silences_warning(caplog: pytest.LogCaptureFixture) ->
def test_msgpack_none_blocks_unregistered(caplog: pytest.LogCaptureFixture) -> None:
"""allowed_msgpack_modules=None should block unregistered types."""
serde = JsonPlusSerializer(allowed_msgpack_modules=None)
obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
@@ -672,6 +657,7 @@ def test_msgpack_allowlist_blocks_non_listed(
caplog: pytest.LogCaptureFixture,
) -> None:
"""Allowlists should block unregistered types even if msgpack is enabled."""
serde = JsonPlusSerializer(
allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")]
)
@@ -999,13 +985,29 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
assert result == obj
def test_delta_sentinel_serde_round_trip() -> None:
from langgraph.checkpoint.base import DELTA_SENTINEL
def test_delta_value_serde_round_trip() -> None:
from langgraph.checkpoint.base import DeltaValue
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
assert type_tag == "msgpack"
assert blob # non-empty ext envelope
original = DeltaValue(
delta=[{"type": "human", "content": "hi"}], prev_checkpoint_id="abc-123"
)
type_tag, blob = serde.dumps_typed(original)
assert type_tag == "delta"
loaded = serde.loads_typed((type_tag, blob))
assert loaded is DELTA_SENTINEL
assert isinstance(loaded, DeltaValue)
assert loaded.delta == original.delta
assert loaded.prev_checkpoint_id == "abc-123"
def test_delta_value_serde_chain_root() -> None:
from langgraph.checkpoint.base import DeltaValue
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
original = DeltaValue(delta=[], prev_checkpoint_id=None)
type_tag, blob = serde.dumps_typed(original)
loaded = serde.loads_typed((type_tag, blob))
assert isinstance(loaded, DeltaValue)
assert loaded.prev_checkpoint_id is None
+22 -344
View File
@@ -6,32 +6,19 @@ from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
Checkpoint,
CheckpointMetadata,
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.jsonplus import (
JsonPlusSerializer,
_warned_blocked_types,
_warned_unregistered_types,
)
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
class MemoryPydantic(BaseModel):
foo: str
@pytest.fixture(autouse=True)
def _reset_warned_types() -> None:
# Warning dedup state is process-global; reset per-test so each case sees
# a fresh slate and assertions about warning emission are stable.
_warned_unregistered_types.clear()
_warned_blocked_types.clear()
class TestMemorySaver:
@pytest.fixture(autouse=True)
def setup(self) -> None:
@@ -209,6 +196,8 @@ class TestMemorySaver:
async def test_memory_saver() -> None:
from langgraph.checkpoint.memory import InMemorySaver
memory_saver = InMemorySaver()
assert isinstance(memory_saver, InMemorySaver)
@@ -322,344 +311,33 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
class TestInMemorySaverDeltaChannel:
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
def test_get_channel_blob(self) -> None:
"""get_channel_blob returns the deserialized blob for a checkpoint+channel."""
from langgraph.checkpoint.base import DeltaValue, empty_checkpoint
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0000000000000000"
version = "00000000000000000000000000000001.0000000000000000"
delta = DeltaValue(delta=[{"content": "hi"}], prev_checkpoint_id=None)
saver.blobs[(thread_id, ns, channel, version)] = serde.dumps_typed(delta)
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
cp = empty_checkpoint()
cp["id"] = "cp1"
cp["channel_versions"][channel] = version
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
"cp1": (serde.dumps_typed(cp), serde.dumps_typed({}), None)
}
result = saver._load_blobs(thread_id, ns, {channel: v1})
assert channel in result
assert result[channel] is DELTA_SENTINEL
result = saver.get_channel_blob(thread_id, ns, "cp1", channel)
assert isinstance(result, DeltaValue)
assert result.delta == [{"content": "hi"}]
assert result.prev_checkpoint_id is None
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)."""
def test_get_channel_blob_missing(self) -> None:
"""get_channel_blob returns NotImplemented when checkpoint or channel not found."""
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"}),
"",
assert (
saver.get_channel_blob("t1", "", "no-such-cp", "messages") is NotImplemented
)
# 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
+1 -4
View File
@@ -7,9 +7,6 @@ resolution-markers = [
"python_full_version < '3.11'",
]
[options]
prerelease-mode = "allow"
[[package]]
name = "annotated-types"
version = "0.7.0"
@@ -289,7 +286,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a2"
version = "4.0.2"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -5,5 +5,5 @@ description = "Test for prerelease stuff"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
"langchain-openai==1.1.14"
"langchain-openai==1.0.1"
]
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
"langchain-openai==1.1.14",
"langchain-openai==1.0.0a2",
"langchain-anthropic==1.0.0a5",
"langgraph==1.1.5"
]
@@ -5,7 +5,7 @@ description = "Test for prerelease stuff"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
"langchain-openai==1.1.14",
"langchain-openai==1.0.0a2",
"langgraph==1.1.2",
"langchain_community>=0.3.0",
]
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.4.23"
__version__ = "0.4.22"
+1 -1
View File
@@ -23,7 +23,7 @@ dependencies = [
path = "langgraph_cli/__init__.py"
[project.optional-dependencies]
inmem = [
"langgraph-api>=0.5.35,<0.9.0 ; python_version >= '3.11'",
"langgraph-api>=0.5.35,<0.8.0 ; python_version >= '3.11'",
"langgraph-runtime-inmem>=0.7 ; python_version >= '3.11'",
]
+383 -465
View File
File diff suppressed because it is too large Load Diff
+18
View File
@@ -245,6 +245,15 @@ class _GraphCallbackManager(BaseCallbackManager):
run_id=run_id,
)
def add_handler(
self,
handler: BaseCallbackHandler,
inherit: bool = True, # noqa: FBT001,FBT002
) -> None:
if not isinstance(handler, GraphCallbackHandler):
raise TypeError("handlers must inherit GraphCallbackHandler")
super().add_handler(handler, inherit=inherit)
def copy(
self,
*,
@@ -312,6 +321,15 @@ class _AsyncGraphCallbackManager(BaseCallbackManager):
run_id=run_id,
)
def add_handler(
self,
handler: BaseCallbackHandler,
inherit: bool = True, # noqa: FBT001,FBT002
) -> None:
if not isinstance(handler, GraphCallbackHandler):
raise TypeError("handlers must inherit GraphCallbackHandler")
super().add_handler(handler, inherit=inherit)
def copy(
self,
*,
@@ -119,3 +119,12 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
Returns `True` if the channel was updated, `False` otherwise.
"""
return False
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
"""Called after checkpoint() with the assigned version, and after
from_checkpoint() with the current channel version.
No-op by default. Override in channels that track their own version
for incremental checkpointing (e.g. DeltaChannel).
"""
pass
+9 -16
View File
@@ -22,9 +22,10 @@ __all__ = ("BinaryOperatorAggregate",)
def _strip_extras(t): # type: ignore[no-untyped-def]
"""Strips Annotated, Required and NotRequired from a given type."""
if hasattr(t, "__origin__"):
if t.__origin__ in (Required, NotRequired):
return _strip_extras(t.__args__[0])
return _strip_extras(t.__origin__)
if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
return _strip_extras(t.__args__[0])
return t
@@ -32,22 +33,11 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]:
"""Inspects the given value and returns (is_overwrite, overwrite_value)."""
if isinstance(value, Overwrite):
return True, value.value
if isinstance(value, dict) and len(value) == 1 and OVERWRITE in value:
if isinstance(value, dict) and set(value.keys()) == {OVERWRITE}:
return True, value[OVERWRITE]
return False, None
def _operators_equal(a: Callable, b: Callable) -> bool:
"""Return True if two reducer operators should be considered equal.
Lambdas all share the name '<lambda>' so identity comparison is
unreliable; treat any pairing that includes a lambda as equal.
"""
if a.__name__ == "<lambda>" or b.__name__ == "<lambda>":
return True
return a is b
class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
"""Stores the result of applying a binary operator to the current value and each new value.
@@ -78,8 +68,11 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
self.value = MISSING
def __eq__(self, value: object) -> bool:
return isinstance(value, BinaryOperatorAggregate) and _operators_equal(
self.operator, value.operator
return isinstance(value, BinaryOperatorAggregate) and (
value.operator is self.operator
if value.operator.__name__ != "<lambda>"
and self.operator.__name__ != "<lambda>"
else True
)
@property
+138 -129
View File
@@ -1,183 +1,175 @@
from __future__ import annotations
import collections.abc
import copy as _copy
from collections.abc import Callable, Sequence
from typing import Any, Generic
from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
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, _operators_equal, _strip_extras
from langgraph.errors import (
EmptyChannelError,
ErrorCode,
InvalidUpdateError,
create_error_message,
)
from langgraph.channels.binop import _get_overwrite, _strip_extras
from langgraph.errors import EmptyChannelError
__all__ = ("DeltaChannel",)
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
"""Reducer channel that stores only a sentinel in checkpoint blobs and
reconstructs state by replaying ancestor writes through the reducer.
class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
"""A channel that stores only per-step write deltas in checkpoints.
The reducer receives the current accumulated value and a batch of writes
in one call: `reducer(state, [write1, write2, ...]) -> new_state`.
Reconstructs the full accumulated list at load time by replaying the
chain of deltas through the operator. Use with append-style reducers
(e.g. `add_messages`) on long-running threads to reduce checkpoint
storage from O(N²) to O(N).
Reducers must be deterministic and batching-invariant (associative across
folds): applying two consecutive write batches separately must produce the
same state as applying their concatenation once:
Works with all checkpointers. Savers with a dedicated blob store
(InMemorySaver, PostgresSaver) use an O(1) fast-path per chain step;
all others (SQLite, MongoDB, etc.) fall back to get_tuple traversal.
reducer(reducer(state, xs), ys) == reducer(state, xs + ys)
Use `snapshot_every=N` to cap chain traversal depth at N steps. Every N
steps a full snapshot is written as the chain root; subsequent deltas
chain back to it, so `get_state` / reload never traverses more than N
checkpoints regardless of thread length. Recommended for savers without
a dedicated blob store.
This lets LangGraph replay checkpointed writes in larger batches than they
were originally produced without changing reconstructed state.
Usage::
`snapshot_frequency=None` (default): pure delta; stores only
`DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes.
`snapshot_frequency=N`: `create_checkpoint` writes a full `_DeltaSnapshot`
blob every N steps, bounding replay depth to N.
Parameters:
reducer: `(state, list[writes]) -> new_state`. Must be deterministic
and batching-invariant as described above.
typ: The value type (e.g. `list`, `dict`). Inferred automatically
from the outer type when used inside `Annotated[T, DeltaChannel(...)]`.
snapshot_frequency: Every Nth pregel step writes a snapshot blob.
`None` (default) = pure delta, never snapshot.
class State(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
# Cap reconstruction depth (recommended for SQLite / MongoDB savers):
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages, snapshot_every=50)]
"""
__slots__ = ("value", "reducer", "snapshot_frequency")
value: Value | Any
__slots__ = (
"value",
"operator",
"snapshot_every",
"_pending",
"_base_version",
"_last_checkpoint_id",
"_overwritten",
"_steps_since_snapshot",
)
def __init__(
self,
reducer: Callable[[Any, Sequence[Any]], Any],
typ: type[Value] | None = None,
operator: Callable[[list[Value], Any], list[Value]],
typ: type = list,
*,
snapshot_frequency: int | None = None,
snapshot_every: int | None = None,
) -> None:
if typ is None:
typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel
super().__init__(typ)
self.reducer = reducer
self.snapshot_frequency = snapshot_frequency
typ = _strip_extras(typ)
if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
if typ in (
collections.abc.Sequence,
collections.abc.MutableSequence,
):
typ = list
if typ in (collections.abc.Set, collections.abc.MutableSet):
typ = set
if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
typ = dict
self.typ = typ
self.value: Any = MISSING
super().__init__(typ)
self.operator = operator
self.snapshot_every = snapshot_every
try:
self.value: list[Value] = typ()
except Exception:
self.value = []
self._pending: list[Any] = []
self._base_version: str | None = None
self._last_checkpoint_id: str | None = None
self._overwritten: bool = False
self._steps_since_snapshot: int = 0
def __eq__(self, other: object) -> bool:
if not isinstance(other, DeltaChannel):
return False
if self.snapshot_frequency != other.snapshot_frequency:
if self.snapshot_every != other.snapshot_every:
return False
return _operators_equal(self.reducer, other.reducer)
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
return list[self.typ] # type: ignore[name-defined]
@property
def UpdateType(self) -> Any:
return self.typ
def is_snapshot_step(self, step: int) -> bool:
"""True if pregel should write a snapshot blob at this step."""
return (
self.snapshot_frequency is not None
and step > 0
and step % self.snapshot_frequency == 0
)
return self.typ | list[self.typ] # type: ignore[name-defined]
def copy(self) -> Self:
new = self.__class__(
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
)
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
new.key = self.key
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
new.value = self.value[:]
new._pending = self._pending[:]
new._base_version = self._base_version
new._last_checkpoint_id = self._last_checkpoint_id
new._overwritten = self._overwritten
new._steps_since_snapshot = self._steps_since_snapshot
return new
def from_checkpoint(self, checkpoint: Any) -> Self:
"""Initialize from a stored blob or sentinel.
Blob types (dispatched via serde ext code, not dict key inspection):
* `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes.
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
* plain value (migration from old BinOp blobs): use directly.
"""
new = self.__class__(
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
)
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
new.key = self.key
if checkpoint is MISSING or checkpoint is DELTA_SENTINEL:
new.value = self.typ()
elif isinstance(checkpoint, _DeltaSnapshot):
new.value = checkpoint.value
if checkpoint is MISSING:
new.value = []
elif isinstance(checkpoint, DeltaChainValue):
accumulated: list[Value] = list(checkpoint.base) if checkpoint.base else []
for step_writes in checkpoint.deltas:
for write in step_writes:
accumulated = new.operator(accumulated, write)
new.value = accumulated
# Seed the counter from actual chain depth so rehydration fires at
# the right time regardless of how many prior invocations there were.
new._steps_since_snapshot = len(checkpoint.deltas)
elif isinstance(checkpoint, DeltaValue):
# Should never reach here — the pregel layer assembles DeltaValues
# into DeltaChainValue before calling from_checkpoint.
raise AssertionError(
"DeltaChannel.from_checkpoint received a raw DeltaValue. "
"This is a bug in the pregel layer — chain assembly should have "
"occurred before from_checkpoint was called."
)
else:
new.value = checkpoint
# Backwards compat: plain list from old BinaryOperatorAggregate checkpoint.
new.value = list(checkpoint)
new._pending = []
new._base_version = None # set by the subsequent after_checkpoint() call
new._overwritten = False
return new
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
"""Apply ancestor writes oldest-to-newest via a single reducer call.
If any write is an Overwrite, the last one in the sequence acts as
the reset point: its value becomes the new base and only writes
after it are passed to the reducer.
"""
values = [v for _, _, v in writes]
if not values:
return
base = self.value
start = 0
for i, v in enumerate(values):
is_ow, ow_value = _get_overwrite(v)
if is_ow:
base = _copy.copy(ow_value) if ow_value is not None else self.typ()
start = i + 1
remaining = values[start:]
self.value = self.reducer(base, remaining) if remaining else base
def update(self, values: Sequence[Any]) -> bool:
if not values:
return False
overwrite_idx: int | None = None
for i, v in enumerate(values):
is_ow, _ = _get_overwrite(v)
if is_ow:
if overwrite_idx is not None:
seen_overwrite = False
for value in values:
is_overwrite, overwrite_value = _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)
overwrite_idx = i
if overwrite_idx is not None:
_, overwrite_value = _get_overwrite(values[overwrite_idx])
base = (
_copy.copy(overwrite_value)
if overwrite_value is not None
else self.typ()
)
remaining = [v for i, v in enumerate(values) if i != overwrite_idx]
self.value = self.reducer(base, remaining) if remaining else base
return True
base = self.typ() if self.value is MISSING else self.value
self.value = self.reducer(base, list(values))
self.value = (
list(overwrite_value) if overwrite_value is not None else []
)
self._pending = list(self.value)
self._overwritten = True
seen_overwrite = True
elif not seen_overwrite:
self.value = self.operator(self.value, value)
self._pending.append(value)
return True
def get(self) -> Any:
def get(self) -> list[Value]:
if self.value is MISSING:
raise EmptyChannelError()
return self.value
@@ -186,12 +178,29 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
return self.value is not MISSING
def checkpoint(self) -> Any:
"""Return stored representation: always `DELTA_SENTINEL`.
if (
self.snapshot_every is not None
and self._steps_since_snapshot >= self.snapshot_every
):
# Emit a full snapshot to cap chain depth at snapshot_every.
# The saver stores this as a plain (non-diff) blob, so future
# deltas will chain back to it and traversal depth resets to 1.
return list(self.value)
return DeltaValue(
delta=self._pending[:],
prev_checkpoint_id=None if self._overwritten else self._last_checkpoint_id,
)
Snapshot decisions are made by `create_checkpoint` in pregel (which
has the step number) via `is_snapshot_step`. `checkpoint()` is only
called for non-snapshot steps or when no checkpointer is available.
"""
if self.value is MISSING:
return MISSING
return DELTA_SENTINEL
def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None:
if version != self._base_version:
if self._base_version is None:
pass # First call after from_checkpoint — anchor without counting a step.
elif self.snapshot_every is not None:
if self._steps_since_snapshot >= self.snapshot_every:
self._steps_since_snapshot = 0
else:
self._steps_since_snapshot += 1
self._base_version = version
self._last_checkpoint_id = checkpoint_id
self._pending = []
self._overwritten = False
-46
View File
@@ -244,52 +244,6 @@ def add_messages(
return merged
def _messages_delta_reducer(
state: list[AnyMessage], writes: list[list[AnyMessage]]
) -> list[AnyMessage]:
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
Processes all writes in one pass — dedup by ID, `RemoveMessage`
tombstoning — without calling `add_messages`. Assumes writes contain
already-typed `BaseMessage` objects (no raw-dict coercion).
This reducer is batching-invariant, as required by `DeltaChannel`:
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
Use `add_messages` as the reducer for `BinaryOperatorAggregate` or
anywhere raw message dicts / strings need to be coerced first.
Example::
from typing import Annotated
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import _messages_delta_reducer
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
"""
from itertools import chain
index: dict[str, int] = {m.id: i for i, m in enumerate(state) if m.id is not None}
result: list[AnyMessage | None] = list(state)
for msg in chain.from_iterable(
[w] if isinstance(w, BaseMessage) else w for w in writes
):
mid = msg.id
if mid is None:
result.append(msg)
elif isinstance(msg, RemoveMessage):
if mid in index:
result[index[mid]] = None
del index[mid]
elif mid in index:
result[index[mid]] = msg
else:
index[mid] = len(result)
result.append(msg)
return [m for m in result if m is not None]
@deprecated(
"MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
category=None,
-16
View File
@@ -48,7 +48,6 @@ 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.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 (
@@ -1083,7 +1082,6 @@ 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]] = [
@@ -1669,20 +1667,6 @@ 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]
item = item.__class__(
item.reducer,
origin,
snapshot_frequency=item.snapshot_frequency,
)
return item
elif isclass(item) and issubclass(item, BaseChannel):
# ex, Annotated[int, EphemeralValue, SomeOtherAnnotation]
+183 -109
View File
@@ -1,22 +1,189 @@
from __future__ import annotations
from collections.abc import Callable, Mapping
import logging
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any, cast
from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
DeltaChainValue,
DeltaValue,
)
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.types import _DeltaSnapshot
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
GetNextVersion = Callable[[Any, None], Any]
logger = logging.getLogger(__name__)
_MISSING_SENTINEL = object()
def _assemble_delta_channels(
checkpoint: Checkpoint,
config: RunnableConfig,
checkpointer: BaseCheckpointSaver,
) -> dict[str, Any]:
"""Resolve any DeltaValue entries in checkpoint channel_values to DeltaChainValue.
Returns a dict of only the channels that needed assembly (others are untouched).
Tries get_channel_blob fast-path first; falls back to get_tuple traversal.
"""
thread_id = str(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
current_checkpoint_id = checkpoint.get("id")
assembled: dict[str, Any] = {}
for channel, value in checkpoint["channel_values"].items():
if not isinstance(value, DeltaValue):
continue
chain_deltas: list[list[Any]] = []
base: list[Any] | None = None
cursor: DeltaValue = value
# Pre-seed with current checkpoint ID to guard against self-referential chains.
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
while True:
chain_deltas.append(cursor.delta)
prev_id = cursor.prev_checkpoint_id
if prev_id is None:
break # chain root
if prev_id in visited:
logger.warning(
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
prev_id,
channel,
)
break
visited.add(prev_id)
# Fast path: saver has a dedicated blob store.
blob = checkpointer.get_channel_blob(
thread_id, checkpoint_ns, prev_id, channel
)
if blob is not NotImplemented:
if isinstance(blob, DeltaValue):
cursor = blob
continue
else:
base = blob # plain list = snapshot root
break
# Fallback: load the full checkpoint and extract channel value.
parent_config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": prev_id,
}
}
parent_tuple = checkpointer.get_tuple(parent_config)
if parent_tuple is None:
logger.warning(
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
prev_id,
channel,
)
break
prev_val = parent_tuple.checkpoint["channel_values"].get(
channel, _MISSING_SENTINEL
)
if prev_val is _MISSING_SENTINEL:
break
elif isinstance(prev_val, DeltaValue):
cursor = prev_val
else:
base = prev_val
break
chain_deltas.reverse()
assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas)
return assembled
async def _aassemble_delta_channels(
checkpoint: Checkpoint,
config: RunnableConfig,
checkpointer: BaseCheckpointSaver,
) -> dict[str, Any]:
"""Async version of _assemble_delta_channels."""
thread_id = str(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
current_checkpoint_id = checkpoint.get("id")
assembled: dict[str, Any] = {}
for channel, value in checkpoint["channel_values"].items():
if not isinstance(value, DeltaValue):
continue
chain_deltas: list[list[Any]] = []
base: list[Any] | None = None
cursor: DeltaValue = value
visited: set[str] = {current_checkpoint_id} if current_checkpoint_id else set()
while True:
chain_deltas.append(cursor.delta)
prev_id = cursor.prev_checkpoint_id
if prev_id is None:
break
if prev_id in visited:
logger.warning(
"DeltaChannel chain cycle at checkpoint %r for channel %r; breaking",
prev_id,
channel,
)
break
visited.add(prev_id)
blob = await checkpointer.aget_channel_blob(
thread_id, checkpoint_ns, prev_id, channel
)
if blob is not NotImplemented:
if isinstance(blob, DeltaValue):
cursor = blob
continue
else:
base = blob
break
parent_config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": prev_id,
}
}
parent_tuple = await checkpointer.aget_tuple(parent_config)
if parent_tuple is None:
logger.warning(
"DeltaChannel chain broken: checkpoint %r not found for channel %r",
prev_id,
channel,
)
break
prev_val = parent_tuple.checkpoint["channel_values"].get(
channel, _MISSING_SENTINEL
)
if prev_val is _MISSING_SENTINEL:
break
elif isinstance(prev_val, DeltaValue):
cursor = prev_val
else:
base = prev_val
break
chain_deltas.reverse()
assembled[channel] = DeltaChainValue(base=base, deltas=chain_deltas)
return assembled
def empty_checkpoint() -> Checkpoint:
@@ -37,87 +204,35 @@ def create_checkpoint(
*,
id: str | None = None,
updated_channels: set[str] | None = None,
get_next_version: GetNextVersion | None = None,
force_delta_snapshot: bool = False,
) -> Checkpoint:
"""Create a checkpoint for the given channels.
For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a
`_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor
walk to at most N steps. Snapshots are eager: even if the channel had no
write this step, a version bump is forced (via `get_next_version`) so the
blob is stored by `put()`. Without `get_next_version` (e.g. static
contexts), snapshot steps gracefully fall back to sentinel.
`force_delta_snapshot` writes available `DeltaChannel` values as snapshots
regardless of `snapshot_frequency`. This is used by `durability="exit"`,
where intermediate writes are not stored as ancestor `checkpoint_writes`.
"""
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
if channels is None:
values = checkpoint["channel_values"]
channel_versions = checkpoint["channel_versions"]
else:
values = {}
channel_versions = dict(checkpoint["channel_versions"])
for k in channels:
if k not in channel_versions:
if k not in checkpoint["channel_versions"]:
continue
ch = channels[k]
if (
isinstance(ch, DeltaChannel)
and (force_delta_snapshot or ch.is_snapshot_step(step))
and ch.is_available()
):
# Eager snapshot: bump version if not already written this step
# so put() includes this channel in new_versions and stores blob.
if get_next_version is not None and (
updated_channels is None or k not in updated_channels
):
channel_versions[k] = get_next_version(channel_versions[k], None)
values[k] = _DeltaSnapshot(ch.get())
else:
v = ch.checkpoint()
if v is not MISSING:
values[k] = v
v = channels[k].checkpoint()
if v is not MISSING:
values[k] = v
return Checkpoint(
v=LATEST_VERSION,
ts=ts,
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=channel_versions,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
updated_channels=None if updated_channels is None else sorted(updated_channels),
)
def _needs_replay(spec: BaseChannel, stored: object) -> bool:
"""True if `spec` is a `DeltaChannel` and the stored blob is a sentinel,
requiring an ancestor walk to reconstruct.
`_DeltaSnapshot` blobs and plain values (migration) resolve directly via
`from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay.
"""
if not isinstance(spec, DeltaChannel):
return False
return stored is MISSING or stored is DELTA_SENTINEL
def channels_from_checkpoint(
specs: Mapping[str, BaseChannel | ManagedValueSpec],
checkpoint: Checkpoint,
*,
saver: BaseCheckpointSaver | None = None,
config: RunnableConfig | None = None,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Hydrate channels from a checkpoint.
For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])`
is sufficient. `DeltaChannel` is the exception: sentinel blobs require an
ancestor walk via `saver._get_channel_writes_history`. The walk terminates
at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain
value, so read depth is bounded by `snapshot_frequency`.
"""
"""Get channels from a checkpoint."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -125,51 +240,10 @@ def channels_from_checkpoint(
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 _needs_replay(spec, stored) and saver is not None and config is not None:
delta_spec = cast(DeltaChannel, spec)
history = saver._get_channel_writes_history(config, k)
replay_ch = delta_spec.from_checkpoint(history.seed)
replay_ch.replay_writes(history.writes)
ch = replay_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 _needs_replay(spec, stored) and saver is not None and config is not None:
delta_spec = cast(DeltaChannel, spec)
history = await saver._aget_channel_writes_history(config, k)
replay_ch = delta_spec.from_checkpoint(history.seed)
replay_ch.replay_writes(history.writes)
ch = replay_ch
else:
ch = spec.from_checkpoint(stored)
for k, v in channel_specs.items():
ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
ch.after_checkpoint(checkpoint["channel_versions"].get(k), checkpoint.get("id"))
channels[k] = ch
return channels, managed_specs
+39 -40
View File
@@ -68,7 +68,6 @@ from langgraph.callbacks import (
GraphResumeEvent,
)
from langgraph.channels.base import BaseChannel
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.constants import TAG_HIDDEN
from langgraph.errors import (
@@ -93,7 +92,8 @@ from langgraph.pregel._algo import (
task_path_str,
)
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
_aassemble_delta_channels,
_assemble_delta_channels,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -190,8 +190,6 @@ class PregelLoop:
_migrate_checkpoint: Callable[[Checkpoint], None] | None
submit: Submit
channels: Mapping[str, BaseChannel]
# Only set on AsyncPregelLoop; sync loops keep this as None.
_delta_write_futs: list[Any] | None = None
managed: ManagedValueMapping
checkpoint: Checkpoint
checkpoint_id_saved: str
@@ -410,7 +408,7 @@ class PregelLoop:
task = self.tasks.get(task_id)
else:
task = None
fut = self.submit(
self.submit(
self.checkpointer_put_writes,
config,
writes_to_save,
@@ -418,16 +416,12 @@ class PregelLoop:
task_path_str(task.path) if task else "",
)
else:
fut = self.submit(
self.submit(
self.checkpointer_put_writes,
config,
writes_to_save,
task_id,
)
if self._delta_write_futs is not None and any(
isinstance(self.specs.get(c), DeltaChannel) for c, _ in writes_to_save
):
self._delta_write_futs.append(fut)
# output writes
if hasattr(self, "tasks"):
self.output_writes(task_id, writes)
@@ -839,18 +833,8 @@ class PregelLoop:
# parent. For forks (source=update/fork), use the fork's parent
# checkpoint ID since the fork was created after the subgraph's
# checkpoints from the original execution.
#
# Only gate on is_time_traveling (not is_replaying). When the
# client resumes with an explicit checkpoint_id that happens to
# point at the current head (e.g. LangGraph Studio sending
# `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
# is_replaying is True but is_time_traveling is False. In that
# case subgraphs should load their latest checkpoint normally,
# not go through ReplayState's before-bound lookup which would
# miss subgraph checkpoints created during processing of the
# current parent step.
replay_state: ReplayState | None = None
if is_time_traveling:
if self.is_replaying:
replay_checkpoint_id = self.checkpoint["id"]
if (
self.checkpoint_metadata.get("source")
@@ -898,11 +882,13 @@ class PregelLoop:
self.step,
id=self.checkpoint["id"] if exiting else None,
updated_channels=self.updated_channels,
get_next_version=self.checkpointer_get_next_version
if do_checkpoint
else None,
force_delta_snapshot=exiting and self.durability == "exit",
)
if do_checkpoint and self.channels:
for k, ch in self.channels.items():
ch.after_checkpoint(
self.checkpoint["channel_versions"].get(k),
self.checkpoint.get("id"),
)
# sanitize TASK channel in the checkpoint before saving (durability=="exit")
if TASKS in self.checkpoint["channel_values"] and any(
isinstance(channel, UntrackedValue) for channel in self.channels.values()
@@ -1284,11 +1270,21 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
else []
)
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
# Assemble any DeltaChannel chains before constructing channel objects.
if self.checkpointer is not None:
assembled = _assemble_delta_channels(
self.checkpoint, self.checkpoint_config, self.checkpointer
)
if assembled:
self.checkpoint = {
**self.checkpoint,
"channel_values": {
**self.checkpoint["channel_values"],
**assembled,
},
}
self.channels, self.managed = channels_from_checkpoint(
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
self.specs, self.checkpoint
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
@@ -1383,11 +1379,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
# Drain DeltaChannel write futures before committing the checkpoint so
# DELTA_SENTINEL blobs are never saved ahead of their backing writes.
if self._delta_write_futs:
futs, self._delta_write_futs = self._delta_write_futs, []
await asyncio.gather(*futs)
try:
if prev is not None:
await prev
@@ -1493,15 +1484,23 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if saved.pending_writes is not None
else []
)
self._delta_write_futs = []
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
)
self.channels, self.managed = await achannels_from_checkpoint(
self.specs,
self.checkpoint,
saver=self.checkpointer,
config=self.checkpoint_config,
if self.checkpointer is not None:
assembled = await _aassemble_delta_channels(
self.checkpoint, self.checkpoint_config, self.checkpointer
)
if assembled:
self.checkpoint = {
**self.checkpoint,
"channel_values": {
**self.checkpoint["channel_values"],
**assembled,
},
}
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
)
self.stack.push(self._suppress_interrupt)
self.status = "input"
+14 -6
View File
@@ -14,7 +14,7 @@ from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
from pydantic import BaseModel
from langgraph._internal._constants import NS_SEP
from langgraph._internal._constants import NS_END, NS_SEP
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
from langgraph.pregel.protocol import StreamChunk
from langgraph.types import Command
@@ -132,15 +132,23 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
**kwargs: Any,
) -> Any:
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
task_checkpoint_ns = cast(str, metadata["langgraph_checkpoint_ns"])
checkpoint_ns = (
f"{task_checkpoint_ns.rsplit(NS_END, 1)[0]}{NS_END}"
if NS_END in task_checkpoint_ns
else task_checkpoint_ns
)
ns = tuple(task_checkpoint_ns.split(NS_SEP))[:-1]
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return
stream_metadata = dict(metadata)
stream_metadata["langgraph_checkpoint_ns"] = checkpoint_ns
# Preserve backwards-compatible streamed checkpoint metadata shape.
stream_metadata["checkpoint_ns"] = checkpoint_ns
if tags:
if filtered_tags := [t for t in tags if not t.startswith("seq:step")]:
metadata["tags"] = filtered_tags
self.metadata[run_id] = (ns, metadata)
stream_metadata["tags"] = filtered_tags
self.metadata[run_id] = (ns, stream_metadata)
def on_llm_new_token(
self,
+56 -31
View File
@@ -122,7 +122,8 @@ from langgraph.pregel._algo import (
)
from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import (
achannels_from_checkpoint,
_aassemble_delta_channels,
_assemble_delta_channels,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -1050,17 +1051,23 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
checkpoint = saved.checkpoint
if isinstance(self.checkpointer, BaseCheckpointSaver):
assembled = _assemble_delta_channels(
checkpoint, saved.config, self.checkpointer
)
if assembled:
checkpoint = {
**checkpoint,
"channel_values": {**checkpoint["channel_values"], **assembled},
}
channels, managed = channels_from_checkpoint(
self.channels,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1173,17 +1180,23 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
channels, managed = await achannels_from_checkpoint(
checkpoint = saved.checkpoint
if isinstance(self.checkpointer, BaseCheckpointSaver):
assembled = await _aassemble_delta_channels(
checkpoint, saved.config, self.checkpointer
)
if assembled:
checkpoint = {
**checkpoint,
"channel_values": {**checkpoint["channel_values"], **assembled},
}
channels, managed = channels_from_checkpoint(
self.channels,
saved.checkpoint,
saver=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
config=saved.config,
checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1529,9 +1542,20 @@ class Pregel(
saved = checkpointer.get_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
if saved:
assembled = _assemble_delta_channels(
base_checkpoint, saved.config, checkpointer
)
if assembled:
base_checkpoint = {
**base_checkpoint,
"channel_values": {
**base_checkpoint["channel_values"],
**assembled,
},
}
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
@@ -1550,11 +1574,6 @@ 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]
@@ -1980,9 +1999,20 @@ class Pregel(
saved = await checkpointer.aget_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
if saved:
assembled = await _aassemble_delta_channels(
base_checkpoint, saved.config, checkpointer
)
if assembled:
base_checkpoint = {
**base_checkpoint,
"channel_values": {
**base_checkpoint["channel_values"],
**assembled,
},
}
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
@@ -1998,14 +2028,9 @@ class Pregel(
)
if saved:
checkpoint_config = patch_configurable(config, saved.config[CONF])
channels, managed = await achannels_from_checkpoint(
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]
# no values, just clear all tasks
+3 -3
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "1.2.0a2"
version = "1.1.7a2"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.10"
@@ -24,8 +24,8 @@ classifiers = [
'Programming Language :: Python :: 3.13',
]
dependencies = [
"langchain-core>=1.3.2,<2",
"langgraph-checkpoint>=4.0.3,<5.0.0",
"langchain-core==1.3.0a2",
"langgraph-checkpoint>=2.1.0,<5.0.0",
"langgraph-sdk>=0.3.0,<0.4.0",
"langgraph-prebuilt>=1.0.9,<1.1.0",
"xxhash>=3.5.0",
+316 -445
View File
@@ -1,34 +1,18 @@
import operator
from collections.abc import Sequence
from typing import Annotated
import pytest
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import NotRequired, TypedDict
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 import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer
from langgraph.graph.state import _get_channel
from langgraph.types import Overwrite
pytestmark = pytest.mark.anyio
# ---------------------------------------------------------------------------
# Core channel primitives
# ---------------------------------------------------------------------------
def test_last_value() -> None:
channel = LastValue(int).from_checkpoint(MISSING)
assert channel.ValueType is int
@@ -111,57 +95,95 @@ def test_untracked_value() -> None:
assert channel.ValueType is dict
assert channel.UpdateType is dict
# UntrackedValue should start empty
with pytest.raises(EmptyChannelError):
channel.get()
# Should be able to update with a value
test_data = {"session": "test", "temp": "dir"}
channel.update([test_data])
assert channel.get() == test_data
# Update with new value
new_data = {"session": "updated", "temp": "newdir"}
channel.update([new_data])
assert channel.get() == new_data
# On checkpoint, UntrackedValue should return MISSING
checkpoint = channel.checkpoint()
assert checkpoint is MISSING
# Creating from checkpoint with MISSING should start empty
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
with pytest.raises(EmptyChannelError):
new_channel.get()
# ---------------------------------------------------------------------------
# DeltaChannel — message reducer
# ---------------------------------------------------------------------------
def test_delta_channel_basic_two_steps() -> None:
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DeltaValue
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
ch.after_checkpoint(None)
# Step 1: one message added
ch.update([HumanMessage(content="hi", id="h1")])
d1 = ch.checkpoint()
assert d1 is DELTA_SENTINEL
assert isinstance(d1, DeltaValue)
assert len(d1.delta) == 1
assert d1.prev_checkpoint_id is None # first ever step
ch.after_checkpoint("v1", checkpoint_id="cid1")
# Step 2: another message
ch.update([AIMessage(content="hello", id="a1")])
d2 = ch.checkpoint()
assert d2 is DELTA_SENTINEL
assert d2.prev_checkpoint_id == "cid1"
assert len(d2.delta) == 1
ch.after_checkpoint("v2")
# 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."""
spec = DeltaChannel(_messages_delta_reducer, list)
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")),
]
def test_delta_channel_after_checkpoint_no_op_when_unchanged() -> None:
from langchain_core.messages import HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
ch.after_checkpoint(None)
ch.update([HumanMessage(content="hi", id="h1")])
ch.after_checkpoint("v1")
# Same version: no-op
ch.after_checkpoint("v1")
assert ch._base_version == "v1"
assert ch._pending == []
def test_delta_channel_from_checkpoint_chain() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import DeltaChainValue
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
chain = DeltaChainValue(
base=None,
deltas=[
[HumanMessage(content="hi", id="h1")],
[AIMessage(content="hello", id="a1")],
[HumanMessage(content="bye", id="h2")],
],
)
ch = spec.from_checkpoint(chain)
msgs = ch.get()
assert len(msgs) == 3
assert msgs[0].content == "hi"
@@ -170,484 +192,333 @@ def test_delta_channel_from_checkpoint_writes_list() -> None:
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
spec = DeltaChannel(_messages_delta_reducer, list)
from langchain_core.messages import HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
# Old BinaryOperatorAggregate checkpoint: plain list
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:
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
ch.update([HumanMessage(content="old", id="h1")])
def test_delta_channel_overwrite_resets_chain() -> None:
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import DeltaValue
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.after_checkpoint(None)
ch.update([HumanMessage(content="old", id="h1")])
ch.after_checkpoint("v1")
# Overwrite should create a root blob (prev_checkpoint_id=None)
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
d = ch.checkpoint()
assert d is DELTA_SENTINEL
assert len(ch.get()) == 1
assert ch.get()[0].content == "new"
assert isinstance(d, DeltaValue)
assert d.prev_checkpoint_id is None # chain root
assert len(d.delta) == 1
assert d.delta[0].content == "new"
def test_delta_channel_remove_message_and_replay() -> None:
"""RemoveMessage must round-trip correctly when writes are replayed."""
spec = DeltaChannel(_messages_delta_reducer, list)
def test_delta_channel_assembly_fallback_via_get_tuple() -> None:
"""Assembly falls back to get_tuple for savers without get_channel_blob."""
from unittest.mock import MagicMock
from langgraph.checkpoint.base import (
CheckpointTuple,
DeltaChainValue,
DeltaValue,
empty_checkpoint,
)
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
from langgraph.pregel._checkpoint import _assemble_delta_channels
msg1 = {"type": "human", "content": "hello"}
msg2 = {"type": "ai", "content": "world"}
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_values"]["messages"] = [msg1]
cp2 = empty_checkpoint()
cp2["id"] = "cp2"
cp2["channel_values"]["messages"] = DeltaValue(
delta=[msg2], prev_checkpoint_id="cp1"
)
saver = MagicMock()
saver.get_channel_blob.return_value = NotImplemented
saver.get_tuple.return_value = CheckpointTuple(
config={
"configurable": {
"thread_id": "t1",
"checkpoint_ns": "",
"checkpoint_id": "cp1",
}
},
checkpoint=cp1,
metadata={},
parent_config=None,
pending_writes=[],
)
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
assembled = _assemble_delta_channels(cp2, config, saver)
assert "messages" in assembled
chain = assembled["messages"]
assert isinstance(chain, DeltaChainValue)
assert chain.base == [msg1]
assert chain.deltas == [[msg2]]
from langchain_core.messages import AIMessage, HumanMessage
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(chain)
result = ch.get()
assert len(result) == 2
assert isinstance(result[0], HumanMessage) and result[0].content == "hello"
assert isinstance(result[1], AIMessage) and result[1].content == "world"
def test_delta_channel_remove_message_delta_and_replay() -> None:
"""RemoveMessage stored in a delta must round-trip correctly through the chain."""
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(MISSING)
ch.after_checkpoint(None)
# Step 1: add two messages
ch.update([HumanMessage(content="hi", id="h1")])
ch.update([AIMessage(content="hello", id="a1")])
d1 = ch.checkpoint()
assert isinstance(d1, DeltaValue)
ch.after_checkpoint("v1", checkpoint_id="cid1")
assert ch.get() == [
HumanMessage(content="hi", id="h1"),
AIMessage(content="hello", id="a1"),
]
# Step 2: remove the AI message
ch.update([RemoveMessage(id="a1")])
d2 = ch.checkpoint()
assert isinstance(d2, DeltaValue)
assert d2.prev_checkpoint_id == "cid1"
assert any(isinstance(w, RemoveMessage) for w in d2.delta)
ch.after_checkpoint("v2", checkpoint_id="cid2")
assert ch.get() == [HumanMessage(content="hi", id="h1")]
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")),
]
)
# Replay the full chain from scratch — must reproduce the post-remove state
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
ch2 = spec.from_checkpoint(chain)
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."""
spec = DeltaChannel(_messages_delta_reducer, list)
ch = spec.from_checkpoint(MISSING)
def test_delta_channel_update_by_id_delta_and_replay() -> None:
"""Updating a message by ID stored in a delta must round-trip correctly."""
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import DeltaChainValue, DeltaValue
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(MISSING)
ch.after_checkpoint(None)
# Step 1: add a message
ch.update([HumanMessage(content="original", id="h1")])
d1 = ch.checkpoint()
assert isinstance(d1, DeltaValue)
ch.after_checkpoint("v1", checkpoint_id="cid1")
# Step 2: update the same message by ID
ch.update([HumanMessage(content="updated", id="h1")])
d2 = ch.checkpoint()
assert isinstance(d2, DeltaValue)
assert d2.prev_checkpoint_id == "cid1"
ch.after_checkpoint("v2", checkpoint_id="cid2")
assert ch.get() == [HumanMessage(content="updated", id="h1")]
ch2 = spec.from_checkpoint(DELTA_SENTINEL)
ch2.replay_writes(
[
("t0", "messages", HumanMessage(content="original", id="h1")),
("t1", "messages", HumanMessage(content="updated", id="h1")),
]
)
# Replay the full chain — must produce the updated message, not the original
chain = DeltaChainValue(base=None, deltas=[d1.delta, d2.delta])
ch2 = spec.from_checkpoint(chain)
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."""
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
assert ch.checkpoint() is DELTA_SENTINEL
def test_delta_channel_snapshot_every_emits_plain_list() -> None:
"""snapshot_every=N causes a plain-list snapshot after N steps; next deltas chain to it."""
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.base import DeltaValue
ch.update([HumanMessage(content="hi", id="h1")])
assert ch.checkpoint() is DELTA_SENTINEL
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
SNAP = 3
spec = DeltaChannel(add_messages, snapshot_every=SNAP)
ch = spec.from_checkpoint(MISSING)
# First after_checkpoint anchors _base_version without counting a step.
ch.after_checkpoint("v0", checkpoint_id="cid0")
# Steps 1..SNAP: each should stay as DeltaValue; counter increments each step.
for i in range(1, SNAP + 1):
ch.update([HumanMessage(content=f"m{i}", id=f"h{i}")])
ckpt = ch.checkpoint()
assert isinstance(ckpt, DeltaValue), f"expected DeltaValue at step {i}"
ch.after_checkpoint(f"v{i}", checkpoint_id=f"cid{i}")
# Step SNAP+1: _steps_since_snapshot == SNAP → snapshot fires
ch.update([HumanMessage(content="snap", id="hsnap")])
snap = ch.checkpoint()
assert isinstance(snap, list), "expected plain-list snapshot at snapshot_every step"
assert len(snap) == SNAP + 1
# After snapshot, counter resets — next step is DeltaValue again
ch.after_checkpoint("vsnap", checkpoint_id="cidsnap")
ch.update([HumanMessage(content="post", id="hpost")])
post = ch.checkpoint()
assert isinstance(post, DeltaValue)
assert post.prev_checkpoint_id == "cidsnap"
# ---------------------------------------------------------------------------
# DeltaChannel — snapshot frequency
# ---------------------------------------------------------------------------
def test_delta_channel_snapshot_every_end_to_end() -> None:
"""Graph with snapshot_every: get_state returns correct accumulated value after snapshot."""
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
def test_delta_channel_snapshot_step_based() -> None:
"""Snapshots fire on every Nth step regardless of whether the channel was written.
With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot
blob — even if the channel had no write that step (eager snapshot). This
bounds the ancestor walk to at most N steps on any read.
"""
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(_messages_delta_reducer, snapshot_frequency=5)
]
other: str
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
def node_a(state: State) -> dict:
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
def node_b(state: State) -> dict:
return {"other": "y"}
g = StateGraph(State)
g.add_node("a", node_a)
g.add_node("b", node_b)
g.add_edge(START, "a")
g.add_edge("a", "b")
saver = InMemorySaver()
graph = g.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "t1"}}
for i in range(6):
graph.invoke(
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
config,
)
msg_blob_values = [
saver.serde.loads_typed((type_tag, blob))
for k, (type_tag, blob) in saver.blobs.items()
if k[2] == "messages" and type_tag == "msgpack" and blob
]
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
state = graph.get_state(config)
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
def test_delta_channel_snapshot_fires_even_when_not_written() -> None:
"""Eager snapshot: _DeltaSnapshot stored at snapshot step even when the
channel had no write that step (node_b doesn't touch messages).
"""
class State(TypedDict):
messages: Annotated[
list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3)
]
tick: int
def writer(state: State) -> dict:
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
def ticker(state: State) -> dict:
return {"tick": state["tick"] + 1}
g = StateGraph(State)
g.add_node("writer", writer)
g.add_node("ticker", ticker)
g.add_edge(START, "writer")
g.add_edge("writer", "ticker")
saver = InMemorySaver()
graph = g.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "t1"}}
for i in range(5):
graph.invoke(
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0},
config,
)
msg_blobs = {
k: saver.serde.loads_typed((t, b))
for k, (t, b) in saver.blobs.items()
if k[2] == "messages" and t == "msgpack" and b
}
snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)}
assert snapshots, (
"eager snapshots must fire even on steps where messages wasn't written"
)
state = graph.get_state(config)
assert len(state.values["messages"]) == 10 # 5 human + 5 AI
# ---------------------------------------------------------------------------
# DeltaChannel — end-to-end (InMemorySaver)
# ---------------------------------------------------------------------------
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer, list)]
n = {"v": 0}
counter = {"n": 0}
def respond(state: State) -> dict:
n["v"] += 1
return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]}
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": "t1"}}
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "snap-test"}}
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
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
# Run 5 turns — snapshot fires after 2 steps, then again after 2 more
for i in range(5):
graph.invoke({"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config)
state = graph.get_state(config)
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
msgs = state.values["messages"]
# 5 human + 5 AI = 10 total
assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}"
# ---------------------------------------------------------------------------
# DeltaChannel — dict reducer
# ---------------------------------------------------------------------------
def test_delta_channel_assembly_fast_path_returns_delta_value() -> None:
"""get_channel_blob returning a DeltaValue continues chain traversal (fast-path)."""
from unittest.mock import MagicMock
def _delta_channel_with_type(op, typ):
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
return _get_channel("_test", Annotated[typ, DeltaChannel(op)])
def test_delta_channel_dict_reducer_fresh_channel() -> None:
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
def merge_dicts(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING)
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."""
def merge_dicts(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
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(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
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}),
]
from langgraph.checkpoint.base import (
DeltaChainValue,
DeltaValue,
empty_checkpoint,
)
assert ch.get() == {"a": 1, "b": 2, "c": 3}
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
from langgraph.pregel._checkpoint import _assemble_delta_channels
def test_delta_channel_dict_reducer_with_deletions() -> None:
"""Dict reducer that treats None values as deletions works end-to-end."""
msg1 = {"type": "human", "content": "one"}
msg2 = {"type": "ai", "content": "two"}
msg3 = {"type": "human", "content": "three"}
def merge_files(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
for k, v in w.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"}])
ch.update([{"file1.py": None, "file3.py": "content3"}])
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
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"}),
]
# cp3 → cp2 (DeltaValue) → cp1 (base list)
dv_cp2 = DeltaValue(delta=[msg2], prev_checkpoint_id="cp1")
cp3 = empty_checkpoint()
cp3["id"] = "cp3"
cp3["channel_values"]["messages"] = DeltaValue(
delta=[msg3], prev_checkpoint_id="cp2"
)
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
saver = MagicMock()
def _get_blob(thread_id, ns, checkpoint_id, channel):
if checkpoint_id == "cp2":
return dv_cp2 # DeltaValue — chain continues
if checkpoint_id == "cp1":
return [msg1] # plain list — chain root
return NotImplemented
saver.get_channel_blob.side_effect = _get_blob
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
assembled = _assemble_delta_channels(cp3, config, saver)
chain = assembled["messages"]
assert isinstance(chain, DeltaChainValue)
assert chain.base == [msg1]
assert chain.deltas == [[msg2], [msg3]]
spec = DeltaChannel(add_messages)
ch = spec.from_checkpoint(chain)
# add_messages converts dicts to message objects; check by type and content
result = ch.get()
assert len(result) == 3
assert result[0].content == "one"
assert result[1].content == "two"
assert result[2].content == "three"
def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
"""Overwrite(dict) in update() must preserve dict shape, not coerce to list."""
def test_delta_channel_assembly_broken_chain_logs_warning() -> None:
"""If a prev_checkpoint_id points to a missing checkpoint, log a warning and use partial chain."""
from unittest.mock import MagicMock
def merge_dicts(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
from langgraph.checkpoint.base import DeltaValue, empty_checkpoint
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}
from langgraph.pregel._checkpoint import _assemble_delta_channels
def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
"""Overwrite(dict) embedded in replayed writes must reconstruct as dict."""
def merge_dicts(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
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}),
]
cp = empty_checkpoint()
cp["id"] = "cp2"
cp["channel_values"]["messages"] = DeltaValue(
delta=["msg2"], prev_checkpoint_id="cp-missing"
)
assert ch.get() == {"x": 10, "y": 20, "z": 30}
saver = MagicMock()
saver.get_channel_blob.return_value = NotImplemented
saver.get_tuple.return_value = None # checkpoint not found
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`."""
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
def merge_dicts(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
assembled = _assemble_delta_channels(cp, config, saver)
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}
# Should still assemble — with partial chain (just the current delta, base=None)
assert "messages" in assembled
from langgraph.checkpoint.base import DeltaChainValue
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel."""
def merge_files(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
for k, v in w.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)
saved = saver.get_tuple(config)
assert saved is not None
assert saved.checkpoint["channel_values"]["files"] 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",
}
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(state: dict, writes: list) -> dict:
result = dict(state)
for w in writes:
result.update(w)
return result
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}
# ---------------------------------------------------------------------------
# DeltaChannel — 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(_messages_delta_reducer, list)
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(_messages_delta_reducer, list)
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(state, writes):
return writes[-1] if writes else state
spec = DeltaChannel(replace, list)
ch = spec.from_checkpoint(None)
ch.replay_writes([("t0", "x", "after")])
assert ch.get() == "after"
chain = assembled["messages"]
assert isinstance(chain, DeltaChainValue)
assert chain.base is None
assert chain.deltas == [["msg2"]]
@@ -1,46 +1,37 @@
"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff.
"""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
Part 1 — baseline (original): DeltaChannel(inf) vs add_messages (BinOp).
Part 2 — snapshot_frequency sweep: shows the storage/read-latency tradeoff
across frequencies [1, 5, 10, 50, inf] at scale.
Simulates realistic multi-turn conversations with paragraph-length messages
(~100 tokens each) scaling up to 1M-token-equivalent histories.
Key insight:
snapshot_frequency=inf → O(N) storage, O(N) read depth (pure delta)
snapshot_frequency=N → O(N²/N) storage, O(N) read depth bounded by freq
snapshot_frequency=1 → O(N²) storage, O(1) read depth (full snapshot)
Token estimates: 1 token ≈ 4 chars; each turn ≈ 200 tokens (human + AI).
A 1M-token conversation ≈ 5,000 turns of realistic messages.
"""
from __future__ import annotations
import contextlib
import math
import os
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 _messages_delta_reducer, add_messages
from langgraph.graph.message import add_messages
try:
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.sqlite import SqliteSaver
_POSTGRES_AVAILABLE = True
_POSTGRES_URI = os.environ.get(
"LANGGRAPH_BENCH_POSTGRES_URI",
"postgres://postgres@localhost:5432/postgres?sslmode=disable",
)
_SQLITE_AVAILABLE = True
except ImportError:
_POSTGRES_AVAILABLE = False
_SQLITE_AVAILABLE = False
SNAPSHOT_EVERY = 50
# ---------------------------------------------------------------------------
# Realistic message payload (~100 tokens / ~400 chars each)
@@ -49,9 +40,7 @@ except ImportError:
_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?"
"and whether we need to refactor the {component} layer before proceeding."
)
_AI_TEMPLATE = (
@@ -118,21 +107,11 @@ class BinaryState(TypedDict):
class DeltaState(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
def _make_delta_state(snapshot_frequency: int | float) -> type:
"""Create a TypedDict with DeltaChannel at the given snapshot_frequency."""
channel = DeltaChannel(
_messages_delta_reducer, snapshot_frequency=snapshot_frequency
)
# Use the functional TypedDict form so the Annotated type is stored as an
# already-evaluated object rather than a forward-reference string (which
# would fail when get_type_hints tries to resolve 'snapshot_frequency').
return TypedDict( # type: ignore[return-value]
f"DeltaState_freq{snapshot_frequency}",
{"messages": Annotated[list, channel]},
)
class DeltaSnapshotState(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)]
# ---------------------------------------------------------------------------
@@ -178,8 +157,9 @@ def _run_turns(
"""Run n_turns conversation turns.
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
Read latency is the average of 5 get_state calls after the full history
is built — forces state rehydration including ancestor replay if needed.
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"}}
@@ -192,16 +172,16 @@ def _run_turns(
)
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
blob_bytes = (
_total_blob_bytes(graph.checkpointer)
if isinstance(graph.checkpointer, MemorySaver)
else -1
)
if isinstance(graph.checkpointer, MemorySaver):
blob_bytes = _total_blob_bytes(graph.checkpointer)
else:
blob_bytes = -1
return write_elapsed, read_elapsed, blob_bytes
@@ -214,6 +194,7 @@ def _fmt_bytes(n: int) -> str:
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"
@@ -223,248 +204,146 @@ def _approx_tokens(n_turns: int) -> str:
# ---------------------------------------------------------------------------
# Checkpointer factories
# Benchmark matrix
# ---------------------------------------------------------------------------
@contextlib.contextmanager
def _pg_saver(thread_id: str = "bench"):
"""Context manager that yields a fresh PostgresSaver and cleans up after."""
with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver:
saver.setup()
with saver._cursor() as cur:
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
yield saver
with saver._cursor() as cur:
for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,))
# Turn counts chosen to span from a short session to a long-running agent conversation.
# Storage and time complexity differences are clearly visible by 500 turns.
# Extrapolation: 5,000 turns × ~200 tokens/turn ≈ 1M tokens (Claude's full context window).
TURN_COUNTS = [50, 100, 200, 500]
def _checkpointers() -> list[tuple[str, Any]]:
"""Return (label, saver_or_None) pairs for available checkpointers."""
result: list[tuple[str, Any]] = [("InMemory", None)]
if _POSTGRES_AVAILABLE:
try:
import psycopg
def _checkpointer_factories() -> list[tuple[str, Any]]:
"""Return (label, context_manager_or_none) pairs for available checkpointers."""
factories: list[tuple[str, Any]] = [("InMemory", None)]
if _SQLITE_AVAILABLE:
import tempfile
psycopg.connect(_POSTGRES_URI).close()
result.append(("Postgres", "postgres"))
except Exception:
pass
return result
factories.append(("SQLite", tempfile.NamedTemporaryFile(suffix=".db")))
return factories
# ---------------------------------------------------------------------------
# Part 1: baseline DeltaChannel(inf) vs add_messages
# ---------------------------------------------------------------------------
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()
BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500]
DELTA_ONLY_TURN_COUNTS = [1000]
checkpointers: list[tuple[str, Any]] = [("InMemory (fast-path)", None)]
if _SQLITE_AVAILABLE:
checkpointers.append(("SQLite (get_tuple fallback)", "sqlite"))
for cp_label, cp_hint in checkpointers:
print(f"--- Checkpointer: {cp_label} ---")
_run_benchmark_for_checkpointer(cp_hint)
def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
W = 72
def _run_benchmark_for_checkpointer(cp_hint: Any) -> None:
import contextlib
import tempfile
@contextlib.contextmanager
def _make_saver():
if cp_hint is None:
return contextlib.nullcontext(None)
return _pg_saver()
yield None
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 BASELINE_TURN_COUNTS:
W = 120
print("=" * W)
header = (
f"{'turns':>6} {'ctx size':>10} "
f"{'add_msgs (bytes)':>18} {'delta (bytes)':>15} {'delta+snap (bytes)':>18} "
f"{'storage saved':>14} "
f"{'read: add_msgs':>14} {'read: delta+snap':>16}"
)
print(header)
print("-" * W)
results = []
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))
s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
def _bytes_or_na(v: Any) -> str:
if v is None or 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"
print(f"\n [{cp_label}] Storage (blob bytes)")
print(
f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}"
)
print(" " + "-" * (W - 2))
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"
# For non-InMemory savers, blob_bytes are unavailable (-1); use read times only
if b_bytes < 0 or s_bytes < 0:
b_bytes_str = "n/a"
d_bytes_str = "n/a"
s_bytes_str = "n/a"
storage_ratio_str = "n/a"
else:
ratio = b_bytes / d_bytes if d_bytes else float("inf")
ratio_str = f"{ratio:.0f}x"
storage_ratio = b_bytes / s_bytes if s_bytes else float("inf")
b_bytes_str = _fmt_bytes(b_bytes)
d_bytes_str = _fmt_bytes(d_bytes)
s_bytes_str = _fmt_bytes(s_bytes)
storage_ratio_str = f"{storage_ratio:.1f}x"
results.append((turns, b_bytes, s_bytes, b_rt, s_rt, storage_ratio))
print(
f" {turns:>6} {_approx_tokens(turns):>10} "
f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} {ratio_str:>8}"
f"{turns:>6} {_approx_tokens(turns):>10} "
f"{b_bytes_str:>18} {d_bytes_str:>15} {s_bytes_str:>18} "
f"{storage_ratio_str:>14} "
f"{b_rt * 1000:>12.1f}ms {s_rt * 1000:>14.1f}ms"
)
print(f"\n [{cp_label}] Read latency (avg of 5 get_state calls)")
print(f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}")
print(" " + "-" * (W - 2))
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
print("=" * W)
print()
if results:
best = results[-1]
turns, b_bytes, s_bytes, b_rt, s_rt, ratio = best
print(f"Key findings at max scale ({turns} turns):")
print(
f" {turns:>6} {_approx_tokens(turns):>10} "
f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}"
f" Storage: {_fmt_bytes(b_bytes)} (add_messages) → {_fmt_bytes(s_bytes)} (DeltaChannel+snapshot) — {ratio:.0f}x reduction"
)
def run_baseline_benchmark() -> None:
print()
print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency")
print("=" * 72)
for cp_label, cp_hint in _checkpointers():
_run_baseline_for_checkpointer(cp_label, cp_hint)
print()
# ---------------------------------------------------------------------------
# Part 2: snapshot_frequency sweep
# ---------------------------------------------------------------------------
# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta.
SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf]
# Turn counts for the sweep — high enough to show storage divergence.
SWEEP_TURN_COUNTS = [50, 100, 500]
def _freq_label(freq: int | float) -> str:
if freq == math.inf:
return "inf"
return str(int(freq))
def _run_sweep_for_checkpointer(cp_label: str, cp_hint: Any) -> None:
def _make_saver():
if cp_hint is None:
return contextlib.nullcontext(None)
return _pg_saver()
# Collect results: {turns: {freq_label: (write_s, read_s, bytes)}}
results: dict[int, dict[str, tuple[float, float, int]]] = {}
for turns in SWEEP_TURN_COUNTS:
results[turns] = {}
for freq in SNAPSHOT_FREQUENCIES:
state_cls = _make_delta_state(freq)
with _make_saver() as saver:
wt, rt, bb = _run_turns(turns, state_cls, saver)
results[turns][_freq_label(freq)] = (wt, rt, bb)
freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES]
col_w = 12
header = f" {'turns':>6} {'ctx':>10}" + "".join(
f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels
)
print(f"\n [{cp_label}] Storage (blob bytes) — lower is better")
print(header)
print(" " + "-" * (len(header) - 2))
for turns in SWEEP_TURN_COUNTS:
row = f" {turns:>6} {_approx_tokens(turns):>10}"
for label in freq_labels:
_, _, bb = results[turns][label]
row += f" {_fmt_bytes(bb) if bb >= 0 else 'n/a':>{col_w}}"
print(row)
print(f"\n [{cp_label}] Read latency (avg of 5 get_state) — lower is better")
print(header)
print(" " + "-" * (len(header) - 2))
for turns in SWEEP_TURN_COUNTS:
row = f" {turns:>6} {_approx_tokens(turns):>10}"
for label in freq_labels:
_, rt, _ = results[turns][label]
row += f" {f'{rt * 1000:.1f}ms':>{col_w}}"
print(row)
print(
f"\n [{cp_label}] Per-invoke write latency (total / turns) — lower is better"
)
print(header)
print(" " + "-" * (len(header) - 2))
for turns in SWEEP_TURN_COUNTS:
row = f" {turns:>6} {_approx_tokens(turns):>10}"
for label in freq_labels:
wt, _, _ = results[turns][label]
row += f" {f'{(wt / turns) * 1000:.1f}ms':>{col_w}}"
print(row)
def run_snapshot_freq_benchmark() -> None:
print()
print("Part 2 — DeltaChannel snapshot_frequency sweep")
print("Lower freq → fewer snapshots → less storage but deeper read replay")
print("=" * 80)
for cp_label, cp_hint in _checkpointers():
_run_sweep_for_checkpointer(cp_label, cp_hint)
print()
print(
f" Read latency: {b_rt * 1000:.1f}ms (add_messages) vs {s_rt * 1000:.1f}ms (DeltaChannel+snapshot)"
)
print()
print("Legend:")
print(
" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)"
" add_msgs = Annotated[list, add_messages] — current default, O(N²) storage"
)
print(
" delta = DeltaChannel(add_messages) — O(N) storage, unbounded chain at read"
)
print(
f" delta+snap = DeltaChannel(add_messages, snapshot_every={SNAPSHOT_EVERY}) — O(N) storage, O(1) read depth"
)
print(" freq=N snapshot every N writes; read walks at most N ancestor writes")
print(" freq=inf pure delta; read walks entire ancestor chain")
print()
# ---------------------------------------------------------------------------
# Pytest entry points
# Pytest entry point
# ---------------------------------------------------------------------------
@pytest.mark.skip(
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
)
def test_delta_channel_baseline_benchmark(capsys: Any) -> None:
"""DeltaChannel(inf) uses less storage than add_messages at scale."""
def test_delta_channel_benchmark(capsys: Any) -> None:
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
with capsys.disabled():
run_baseline_benchmark()
run_benchmark()
for turns in [25, 50]:
# Correctness assertion: DeltaChannel must use less storage at scale.
for turns in [100, 200]:
_, _, b_bytes = _run_turns(turns, BinaryState)
_, _, d_bytes = _run_turns(turns, DeltaState)
_, _, s_bytes = _run_turns(turns, DeltaSnapshotState)
assert d_bytes < b_bytes, (
f"DeltaChannel should use less storage at {turns} turns, "
f"got delta={d_bytes} binary={b_bytes}"
)
@pytest.mark.skip(
reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py"
)
def test_snapshot_freq_benchmark(capsys: Any) -> None:
"""snapshot_frequency trades storage for bounded read depth."""
with capsys.disabled():
run_snapshot_freq_benchmark()
# Correctness: results at all frequencies should agree on final state.
n_turns = 20
states: dict[str, list] = {}
for freq in SNAPSHOT_FREQUENCIES:
state_cls = _make_delta_state(freq)
graph = _make_graph(state_cls)
config = {"configurable": {"thread_id": "correctness"}}
for i in range(n_turns):
graph.invoke(
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
config,
)
state = graph.get_state(config)
states[_freq_label(freq)] = [m.id for m in state.values["messages"]]
ref = states["inf"]
for label, msg_ids in states.items():
assert msg_ids == ref, (
f"freq={label} produced different message IDs than freq=inf"
assert s_bytes < b_bytes, (
f"DeltaChannel+snapshot should use less storage at {turns} turns, "
f"got snapshot={s_bytes} binary={b_bytes}"
)
@@ -473,6 +352,5 @@ def test_snapshot_freq_benchmark(capsys: Any) -> None:
# ---------------------------------------------------------------------------
if __name__ == "__main__":
run_baseline_benchmark()
run_snapshot_freq_benchmark()
run_benchmark()
sys.exit(0)
@@ -1,613 +0,0 @@
"""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 langchain_core.messages import AIMessage, HumanMessage
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
from langgraph.graph.message import _messages_delta_reducer, add_messages
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 _list_concat(state: list, writes: list) -> list:
result = list(state)
for w in writes:
result.extend(w if isinstance(w, list) else [w])
return result
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(_list_concat)]
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}"
# ---------------------------------------------------------------------------
# 9. Migration from `add_messages` → `DeltaChannel(_messages_delta_reducer)`
#
# `add_messages` is the primary real-world use case: it creates a
# BinaryOperatorAggregate with dedup-by-ID and RemoveMessage semantics.
# After swapping the annotation to DeltaChannel, pre-migration blobs
# (plain lists of Message objects) must be used directly as the seed.
# ---------------------------------------------------------------------------
def _add_messages_graph(checkpointer: Any) -> Any:
class MessagesState(TypedDict):
messages: Annotated[list, add_messages]
return (
StateGraph(MessagesState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def _delta_messages_graph(checkpointer: Any) -> Any:
class DeltaMessagesState(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
return (
StateGraph(DeltaMessagesState)
.add_node("noop", _noop)
.add_edge(START, "noop")
.add_edge("noop", END)
.compile(checkpointer=checkpointer)
)
def test_add_messages_to_delta_migration_preserves_message_history() -> None:
"""Migration from `add_messages` to `DeltaChannel(_messages_delta_reducer)`
preserves message ordering and IDs at both the tip and settled ancestor
boundaries.
The pre-migration blob is a plain list of Message objects; DeltaChannel
must use it directly as the seed without walking ancestors past it.
"""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "add-messages-migration"}}
pre_graph = _add_messages_graph(checkpointer)
pre_graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
pre_graph.invoke({"messages": [AIMessage(content="hi", id="a1")]}, config)
pre_graph.invoke({"messages": [HumanMessage(content="thanks", id="h2")]}, config)
pre_tip = pre_graph.get_state(config)
assert [m.id for m in pre_tip.values["messages"]] == ["h1", "a1", "h2"]
delta_graph = _delta_messages_graph(checkpointer)
# Tip: latest checkpoint has a full list blob — must use it directly.
snap = delta_graph.get_state(config)
assert [m.id for m in snap.values["messages"]] == ["h1", "a1", "h2"], (
f"tip hydration mismatch: got {[m.id for m in snap.values['messages']]}"
)
# Settled ancestor boundaries must also match.
pre_settled = [
[m.id for m in s.values.get("messages", [])]
for s in pre_graph.get_state_history(config)
if s.next == ("__start__",)
]
delta_settled = [
[m.id for m in s.values.get("messages", [])]
for s in delta_graph.get_state_history(config)
if s.next == ("__start__",)
]
assert delta_settled == pre_settled, (
f"settled boundary mismatch after migration: "
f"pre={pre_settled}, delta={delta_settled}"
)
async def test_add_messages_to_delta_migration_preserves_message_history_async() -> (
None
):
"""Async variant of the add_messages migration test."""
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "add-messages-migration-async"}}
pre_graph = _add_messages_graph(checkpointer)
await pre_graph.ainvoke(
{"messages": [HumanMessage(content="hello", id="h1")]}, config
)
await pre_graph.ainvoke({"messages": [AIMessage(content="hi", id="a1")]}, config)
delta_graph = _delta_messages_graph(checkpointer)
snap = await delta_graph.aget_state(config)
assert [m.id for m in snap.values["messages"]] == ["h1", "a1"], (
f"async tip hydration mismatch: got {[m.id for m in snap.values['messages']]}"
)
@@ -275,70 +275,3 @@ def test_graph_callbacks_accept_base_callback_manager() -> None:
assert "__interrupt__" in first
assert len(graph_handler.interrupt_events) == 1
def test_non_graph_handler_via_add_handler_does_not_crash() -> None:
"""Non-GraphCallbackHandler added via add_handler should not raise.
Libraries like opentelemetry-instrumentation-langchain monkey-patch
BaseCallbackManager.__init__ and inject handlers via add_handler().
These handlers inherit from BaseCallbackHandler, not
GraphCallbackHandler. They must be silently accepted — graph lifecycle
events will simply not be dispatched to them.
"""
from langgraph.callbacks import _GraphCallbackManager
manager = _GraphCallbackManager()
plain_handler = _LangChainCustomEventHandler()
manager.add_handler(plain_handler, inherit=True)
assert plain_handler in manager.handlers
def test_non_graph_handler_does_not_receive_lifecycle_events() -> None:
"""Non-GraphCallbackHandler added alongside a GraphCallbackHandler
should not interfere with lifecycle event dispatch."""
graph = _build_interrupt_graph()
graph_handler = _GraphEventHandler()
plain_handler = _LangChainCustomEventHandler()
config = {
"configurable": {"thread_id": "graph-callback-mixed-handlers"},
"callbacks": [plain_handler, graph_handler],
}
first = graph.invoke({"answer": None}, config)
assert "__interrupt__" in first
assert len(graph_handler.interrupt_events) == 1
assert plain_handler.events == []
resumed = graph.invoke(Command(resume="done"), config)
assert resumed == {"answer": "done"}
assert len(graph_handler.resume_events) == 1
assert plain_handler.events == []
@pytest.mark.anyio
@NEEDS_CONTEXTVARS
async def test_non_graph_handler_does_not_receive_lifecycle_events_async() -> None:
"""Async variant: non-GraphCallbackHandler should not interfere."""
graph = _build_interrupt_graph()
graph_handler = _GraphEventHandler()
plain_handler = _LangChainCustomEventHandler()
config = {
"configurable": {"thread_id": "graph-callback-mixed-handlers-async"},
"callbacks": [plain_handler, graph_handler],
}
first = await graph.ainvoke({"answer": None}, config)
assert "__interrupt__" in first
assert len(graph_handler.interrupt_events) == 1
assert plain_handler.events == []
resumed = await graph.ainvoke(Command(resume="done"), config)
assert resumed == {"answer": "done"}
assert len(graph_handler.resume_events) == 1
assert plain_handler.events == []
+30 -96
View File
@@ -16,7 +16,7 @@ from typing import Annotated, Any, Literal, get_type_hints
import pytest
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, RemoveMessage
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.runnables import (
RunnableConfig,
RunnableLambda,
@@ -25,7 +25,6 @@ from langchain_core.runnables import (
from langchain_core.runnables.graph import Edge
from langgraph.cache.base import BaseCache
from langgraph.checkpoint.base import (
DELTA_SENTINEL,
BaseCheckpointSaver,
Checkpoint,
CheckpointMetadata,
@@ -42,7 +41,6 @@ from typing_extensions import NotRequired, TypedDict
from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
@@ -51,7 +49,7 @@ from langgraph.config import get_stream_writer
from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand
from langgraph.func import entrypoint, task
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
from langgraph.graph.message import MessagesState, add_messages
from langgraph.pregel import (
NodeBuilder,
Pregel,
@@ -9406,9 +9404,15 @@ def test_fork_does_not_apply_pending_writes(
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(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
@@ -9442,9 +9446,15 @@ async def test_delta_channel_end_to_end_inmemory() -> None:
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(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
counter = {"n": 0}
@@ -9494,9 +9504,15 @@ async def test_delta_channel_time_travel() -> None:
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(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
return {"messages": [AIMessage(content="reply", id="ai-1")]}
@@ -9535,9 +9551,15 @@ async def test_delta_channel_remove_message_end_to_end() -> None:
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(_messages_delta_reducer)]
messages: Annotated[list, DeltaChannel(add_messages)]
def update_msg(state: State) -> dict:
# re-send h1 with updated content
@@ -9565,91 +9587,3 @@ async def test_delta_channel_update_by_id_end_to_end() -> None:
assert "h1" in ids # h1 persists (updated, not duplicated)
assert "h2" in ids
assert ids.count("h1") == 1, "h1 must not be duplicated"
async def test_delta_channel_durability_exit_stores_snapshot() -> None:
"""DeltaChannel must reload from a durability='exit' checkpoint."""
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
def respond(state: State) -> dict:
return {"messages": [AIMessage(content="reply", id="ai1")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "delta-exit-test"}}
result = graph.invoke(
{"messages": [HumanMessage(content="hello", id="h1")]},
config,
durability="exit",
)
assert [m.content for m in result["messages"]] == ["hello", "reply"]
state = graph.get_state(config)
assert [m.content for m in state.values["messages"]] == ["hello", "reply"]
async def test_delta_channel_async_write_ordering() -> None:
"""In async mode, DeltaChannel write futures are awaited before the checkpoint
is committed, so aput_writes always precedes aput for sentinel checkpoints."""
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
def respond(state: State) -> dict:
i = len(state["messages"])
return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]}
order: list[str] = []
original_aput_writes = InMemorySaver.aput_writes
original_aput = InMemorySaver.aput
async def tracked_aput_writes(self, config, writes, task_id, task_path=""):
result = await original_aput_writes(self, config, writes, task_id, task_path)
order.append("aput_writes")
return result
async def tracked_aput(self, config, checkpoint, metadata, new_versions):
has_sentinel = any(
v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values()
)
order.append("aput_sentinel" if has_sentinel else "aput_other")
return await original_aput(self, config, checkpoint, metadata, new_versions)
InMemorySaver.aput_writes = tracked_aput_writes
InMemorySaver.aput = tracked_aput
try:
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "async-ordering-test"}}
for i in range(3):
await graph.ainvoke(
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config
)
# Every aput_sentinel must be preceded by at least one aput_writes
for i, event in enumerate(order):
if event == "aput_sentinel":
preceding = order[:i]
assert "aput_writes" in preceding, (
f"aput_sentinel at {i} had no preceding aput_writes: {order}"
)
last_write_idx = max(
j for j, e in enumerate(order[:i]) if e == "aput_writes"
)
assert last_write_idx < i, (
f"aput_writes at {last_write_idx} should precede aput_sentinel at {i}: {order}"
)
finally:
InMemorySaver.aput_writes = original_aput_writes
InMemorySaver.aput = original_aput
state = await graph.aget_state(config)
assert len(state.values["messages"]) == 6 # 3 human + 3 AI
-30
View File
@@ -6101,36 +6101,6 @@ async def test_parent_command(
)
async def test_delta_channel_durability_exit_stores_snapshot_async() -> None:
"""DeltaChannel must reload from an async durability='exit' checkpoint."""
from langchain_core.messages import AIMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import _messages_delta_reducer
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
async def respond(state: State) -> dict:
return {"messages": [AIMessage(content="reply", id="ai1")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "delta-exit-async-test"}}
result = await graph.ainvoke(
{"messages": [HumanMessage(content="hello", id="h1")]},
config,
durability="exit",
)
assert [m.content for m in result["messages"]] == ["hello", "reply"]
state = await graph.aget_state(config)
assert [m.content for m in state.values["messages"]] == ["hello", "reply"]
@NEEDS_CONTEXTVARS
async def test_interrupt_subgraph(async_checkpointer: BaseCheckpointSaver) -> None:
class State(TypedDict):
@@ -0,0 +1,185 @@
"""Sweep snapshot_every values to find the storage vs. time-travel tradeoff.
Run directly: python tests/test_rehydrate_sweep.py
Run via pytest: pytest tests/test_rehydrate_sweep.py -s
"""
from __future__ import annotations
import sys
import time
from typing import Annotated, Any
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
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
REHYDRATE_SWEEP = [5, 10, 25, 50, 100, None] # None = no rehydration (pure diff)
TURN_COUNTS = [50, 100, 250, 500]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_state(snapshot_every: int | None) -> type:
channel = DeltaChannel(add_messages, snapshot_every=snapshot_every)
return TypedDict("S", {"messages": Annotated[list, channel]})
def _make_graph(state_cls: type) -> Any:
def human_node(state: Any) -> dict:
return {}
def ai_node(state: Any) -> dict:
last = state["messages"][-1]
return {"messages": [AIMessage(content=f"reply-to-{last.id}")]}
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=MemorySaver())
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 _measure_time_travel_ms(graph: Any, config: dict) -> float:
"""Time how long it takes to get state at the very first checkpoint (worst case)."""
history = list(graph.get_state_history(config))
if not history:
return 0.0
oldest = history[-1]
t0 = time.perf_counter()
graph.get_state(oldest.config)
return (time.perf_counter() - t0) * 1000
def _run(n_turns: int, snapshot_every: int | None) -> tuple[float, int, float]:
"""Returns (write_ms, blob_bytes, time_travel_ms)."""
state_cls = _make_state(snapshot_every)
graph = _make_graph(state_cls)
saver: MemorySaver = graph.checkpointer # type: ignore[assignment]
config = {"configurable": {"thread_id": "sweep"}}
t0 = time.perf_counter()
for i in range(n_turns):
graph.invoke(
{"messages": [HumanMessage(content=f"msg-{i}", id=f"h{i}")]}, config
)
write_ms = (time.perf_counter() - t0) * 1000
blob_bytes = _total_blob_bytes(saver)
tt_ms = _measure_time_travel_ms(graph, config)
return write_ms, blob_bytes, tt_ms
# ---------------------------------------------------------------------------
# ASCII sparkline
# ---------------------------------------------------------------------------
def _sparkline(values: list[float], width: int = 20) -> str:
bars = " ▁▂▃▄▅▆▇█"
lo, hi = min(values), max(values)
span = hi - lo or 1
chars = [bars[round((v - lo) / span * (len(bars) - 1))] for v in values]
return "".join(chars).ljust(width)
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def run_sweep() -> None:
label = {v: (str(v) if v is not None else "None(∞)") for v in REHYDRATE_SWEEP}
print()
print("snapshot_every sweep — storage vs time-travel cost")
print("=" * 90)
for turns in TURN_COUNTS:
print(f"\n--- {turns} turns ---")
col_w = 12
header = (
f"{'snapshot_every':>18} "
f"{'blob_bytes':>{col_w}} "
f"{'write_ms':>{col_w}} "
f"{'time_travel_ms':>{col_w}}"
)
print(header)
print("-" * 60)
tt_vals: list[float] = []
byte_vals: list[int] = []
write_vals: list[float] = []
rows: list[tuple] = []
for rv in REHYDRATE_SWEEP:
write_ms, blob_bytes, tt_ms = _run(turns, rv)
rows.append((rv, blob_bytes, write_ms, tt_ms))
byte_vals.append(blob_bytes)
write_vals.append(write_ms)
tt_vals.append(tt_ms)
for rv, blob_bytes, write_ms, tt_ms in rows:
print(
f"{label[rv]:>18} "
f"{blob_bytes:>{col_w},} "
f"{write_ms:>{col_w}.1f} "
f"{tt_ms:>{col_w}.2f}"
)
print()
print(
f" bytes spark: [{_sparkline(byte_vals)}] "
f"lo={min(byte_vals):,} hi={max(byte_vals):,}"
)
print(
f" time-travel spark: [{_sparkline(tt_vals)}] "
f"lo={min(tt_vals):.2f}ms hi={max(tt_vals):.2f}ms"
)
print(
f" write spark: [{_sparkline(write_vals)}] "
f"lo={min(write_vals):.1f}ms hi={max(write_vals):.1f}ms"
)
print()
print("=" * 90)
print(
"snapshot_every=None means pure diff (no snapshots) — "
"lowest storage, highest time-travel cost."
)
print(
"Lower snapshot_every = more frequent full snapshots = "
"faster time-travel, more storage."
)
print()
def test_rehydrate_sweep(capsys: Any) -> None:
with capsys.disabled():
run_sweep()
if __name__ == "__main__":
run_sweep()
sys.exit(0)
-64
View File
@@ -1113,70 +1113,6 @@ def test_subgraph_interrupt_replay_from_parent_then_resume(
]
def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Resume with Command(resume=...) plus the current head checkpoint_id
in config. The subgraph must continue from the interrupted node, not
restart from scratch. Explicit checkpoint_id triggers is_replaying but
this is a resume, not a time-travel, so ReplayState should not apply."""
called: list[str] = []
def step_a(state: State) -> State:
called.append("step_a")
return {"value": ["sub_a"]}
def ask_human(state: State) -> State:
called.append("ask_human")
answer = interrupt("Provide input:")
return {"value": [f"human:{answer}"]}
def step_b(state: State) -> State:
called.append("step_b")
return {"value": ["sub_b"]}
subgraph = (
StateGraph(State)
.add_node("step_a", step_a)
.add_node("ask_human", ask_human)
.add_node("step_b", step_b)
.add_edge(START, "step_a")
.add_edge("step_a", "ask_human")
.add_edge("ask_human", "step_b")
.compile(checkpointer=True)
)
graph = (
StateGraph(State)
.add_node("subgraph_node", subgraph)
.add_edge(START, "subgraph_node")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
# Run until interrupt fires in subgraph
graph.invoke({"value": []}, config)
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"]
called.clear()
resume_config = {
"configurable": {
"thread_id": "1",
"checkpoint_id": head_checkpoint_id,
"checkpoint_ns": "",
}
}
result = graph.invoke(Command(resume="answer"), resume_config)
assert called == ["ask_human", "step_b"]
assert "__interrupt__" not in result
assert result["value"] == ["sub_a", "human:answer", "sub_b"]
def test_subgraph_replay_loads_accumulated_state_then_resume(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
+8 -24
View File
@@ -7,9 +7,6 @@ resolution-markers = [
"python_full_version < '3.11'",
]
[options]
prerelease-mode = "allow"
[[package]]
name = "aiosqlite"
version = "0.22.1"
@@ -1351,11 +1348,10 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.2"
version = "1.3.0a2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
{ name = "langchain-protocol" },
{ name = "langsmith" },
{ name = "packaging" },
{ name = "pydantic" },
@@ -1364,26 +1360,14 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.14"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/05/bf/efb5e2ed832e4d6d45590e25a9e5191986b291b543bc6a807b48bee070b0/langchain_protocol-0.0.14.tar.gz", hash = "sha256:bc1e8553122e6ede310280462d5813023a172ff2785ccbbdec54d43f3a15e5f2", size = 5862, upload-time = "2026-04-29T16:40:18.657Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c2/e9/06c47ecb2aff08f83dfa30058da3bf86be64862c19569043ed5331bbeecd/langchain_protocol-0.0.14-py3-none-any.whl", hash = "sha256:ffc35089779bd8ca217015180cef5e660fc3b074efdaa0f2e95df73583f1a047", size = 6984, upload-time = "2026-04-29T16:40:17.841Z" },
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
]
[[package]]
name = "langgraph"
version = "1.2.0a2"
version = "1.1.7a2"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -1455,7 +1439,7 @@ test = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "../sdk-py" },
@@ -1564,7 +1548,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a2"
version = "4.0.2"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },
@@ -1722,7 +1706,7 @@ inmem = [
requires-dist = [
{ name = "click", specifier = ">=8.1.7" },
{ name = "httpx", specifier = ">=0.24.0" },
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.9.0" },
{ name = "langgraph-api", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.5.35,<0.8.0" },
{ name = "langgraph-runtime-inmem", marker = "python_full_version >= '3.11' and extra == 'inmem'", specifier = ">=0.7" },
{ name = "langgraph-sdk", marker = "python_full_version >= '3.11'", specifier = ">=0.1.0" },
{ name = "pathspec", specifier = ">=0.11.0" },
@@ -1758,7 +1742,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.0.10"
version = "1.0.9"
source = { editable = "../prebuilt" }
dependencies = [
{ name = "langchain-core" },
+7 -23
View File
@@ -614,7 +614,6 @@ class _InjectedArgs:
store: str | None
runtime: str | None
all_injected_keys: set[str]
_optional_state_args: set[str]
class ToolNode(RunnableCallable):
@@ -808,7 +807,6 @@ class ToolNode(RunnableCallable):
context=runtime.context,
store=runtime.store,
stream_writer=runtime.stream_writer,
tools=list(self.tools_by_name.values()),
execution_info=runtime.execution_info,
server_info=runtime.server_info,
)
@@ -843,7 +841,6 @@ class ToolNode(RunnableCallable):
context=runtime.context,
store=runtime.store,
stream_writer=runtime.stream_writer,
tools=list(self.tools_by_name.values()),
execution_info=runtime.execution_info,
server_info=runtime.server_info,
)
@@ -1336,7 +1333,7 @@ class ToolNode(RunnableCallable):
return tool_call
tool_call_copy: ToolCall = copy(tool_call)
injected_args: dict[str, Any] = {}
injected_args = {}
# Inject state
if injected.state:
@@ -1364,20 +1361,14 @@ class ToolNode(RunnableCallable):
# Extract state values
if isinstance(state, dict):
for tool_arg, state_field in injected.state.items():
if not state_field:
injected_args[tool_arg] = state
elif state_field in state:
injected_args[tool_arg] = state[state_field]
elif tool_arg not in injected._optional_state_args:
raise KeyError(state_field)
injected_args[tool_arg] = (
state[state_field] if state_field else state
)
else:
for tool_arg, state_field in injected.state.items():
if not state_field:
injected_args[tool_arg] = state
elif hasattr(state, state_field):
injected_args[tool_arg] = getattr(state, state_field)
elif tool_arg not in injected._optional_state_args:
raise AttributeError(state_field)
injected_args[tool_arg] = (
getattr(state, state_field) if state_field else state
)
# Inject store
if injected.store:
@@ -1578,7 +1569,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
- `context`: Runtime context (shared with `Runtime`)
- `store`: `BaseStore` instance for persistent storage (shared with `Runtime`)
- `stream_writer`: `StreamWriter` for streaming output (shared with `Runtime`)
- `tools`: List of all available `BaseTool` instances
No `Annotated` wrapper is needed - just use `runtime: ToolRuntime`
as a parameter.
@@ -1621,7 +1611,6 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
context: ContextT
config: RunnableConfig
stream_writer: StreamWriter
tools: list[BaseTool]
tool_call_id: str | None
store: BaseStore | None
execution_info: ExecutionInfo | None = None
@@ -1870,7 +1859,6 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
store_arg: str | None = None
runtime_arg: str | None = None
all_injected_keys: set[str] = set()
_optional_state_args: set[str] = set()
for name, type_ in all_annotations.items():
# Track all InjectedToolArg-annotated params (including custom subclasses)
@@ -1885,9 +1873,6 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
if state_inj := _get_injection_from_type(type_, InjectedState):
if isinstance(state_inj, InjectedState) and state_inj.field:
state_args[name] = state_inj.field
field_info = full_schema.model_fields.get(name)
if field_info and not field_info.is_required():
_optional_state_args.add(name)
else:
state_args[name] = None
@@ -1904,5 +1889,4 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
store=store_arg,
runtime=runtime_arg,
all_injected_keys=all_injected_keys,
_optional_state_args=_optional_state_args,
)
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-prebuilt"
version = "1.0.10"
version = "1.0.9"
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
authors = []
requires-python = ">=3.10"
@@ -69,7 +69,6 @@ def _create_config_with_runtime(store=None, state=None):
context={},
store=store,
stream_writer=None,
tools=[],
tool_call_id="test_id",
)
return {
+8 -29
View File
@@ -2016,8 +2016,8 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
assert tool_message.tool_call_id == "call_dynamic_2"
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
def test_tool_runtime_forwards_execution_info_and_server_info() -> None:
"""Test that execution_info and server_info are forwarded from Runtime to ToolRuntime."""
from langgraph.runtime import ExecutionInfo, ServerInfo
exec_info = ExecutionInfo(
@@ -2043,15 +2043,9 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
"""Tool that captures runtime info."""
captured["execution_info"] = runtime.execution_info
captured["server_info"] = runtime.server_info
captured["tools"] = runtime.tools
return "ok"
@dec_tool
def other_tool(y: int) -> str:
"""Another tool available to the runtime."""
return str(y)
node = ToolNode([info_tool, other_tool])
node = ToolNode([info_tool])
tool_call = {
"name": "info_tool",
"args": {"x": 1},
@@ -2060,21 +2054,17 @@ def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
}
msg = AIMessage("", tool_calls=[tool_call])
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
result = node.invoke({"messages": [msg]}, config=config)
node.invoke({"messages": [msg]}, config=config)
assert result["messages"][-1].content == "ok"
assert captured["execution_info"] is exec_info
assert captured["execution_info"].thread_id == "t-1"
assert captured["execution_info"].task_id == "tk-1"
assert captured["server_info"] is server_info
assert captured["server_info"].assistant_id == "asst-1"
assert [tool.name for tool in captured["tools"]] == ["info_tool", "other_tool"]
async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async() -> (
None
):
"""Test that execution_info, server_info, and tools are forwarded in async path."""
async def test_tool_runtime_forwards_execution_info_and_server_info_async() -> None:
"""Test that execution_info and server_info are forwarded in async path."""
from langgraph.runtime import ExecutionInfo, ServerInfo
exec_info = ExecutionInfo(
@@ -2100,15 +2090,9 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
"""Async tool that captures runtime info."""
captured["execution_info"] = runtime.execution_info
captured["server_info"] = runtime.server_info
captured["tools"] = runtime.tools
return "ok"
@dec_tool
async def other_tool_async(y: int) -> str:
"""Another async tool available to the runtime."""
return str(y)
node = ToolNode([info_tool_async, other_tool_async])
node = ToolNode([info_tool_async])
tool_call = {
"name": "info_tool_async",
"args": {"x": 1},
@@ -2117,17 +2101,12 @@ async def test_tool_runtime_forwards_execution_info_server_info_and_tools_async(
}
msg = AIMessage("", tool_calls=[tool_call])
config: RunnableConfig = {"configurable": {"__pregel_runtime": mock_runtime}}
result = await node.ainvoke({"messages": [msg]}, config=config)
await node.ainvoke({"messages": [msg]}, config=config)
assert result["messages"][-1].content == "ok"
assert captured["execution_info"] is exec_info
assert captured["execution_info"].thread_id == "t-2"
assert captured["server_info"] is server_info
assert captured["server_info"].graph_id == "graph-2"
assert [tool.name for tool in captured["tools"]] == [
"info_tool_async",
"other_tool_async",
]
# --- InjectedToolArg security tests ---
+6 -6
View File
@@ -249,7 +249,7 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.0"
version = "1.3.0a2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
@@ -261,14 +261,14 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
]
[[package]]
name = "langgraph"
version = "1.1.9"
version = "1.1.7a2"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
@@ -281,7 +281,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "." },
{ name = "langgraph-sdk", editable = "../sdk-py" },
@@ -490,7 +490,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.0.10"
version = "1.0.9"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
+26 -4
View File
@@ -1,8 +1,30 @@
from langgraph_sdk.auth import Auth
from langgraph_sdk.client import get_client, get_sync_client
from langgraph_sdk.encryption import Encryption
from langgraph_sdk.encryption.types import EncryptionContext
from __future__ import annotations
import importlib
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from langgraph_sdk.auth import Auth
from langgraph_sdk.client import get_client, get_sync_client
from langgraph_sdk.encryption import Encryption
from langgraph_sdk.encryption.types import EncryptionContext
__version__ = "0.3.13"
__all__ = ["Auth", "Encryption", "EncryptionContext", "get_client", "get_sync_client"]
_LAZY: dict[str, str] = {
"Auth": "langgraph_sdk.auth",
"get_client": "langgraph_sdk.client",
"get_sync_client": "langgraph_sdk.client",
"Encryption": "langgraph_sdk.encryption",
"EncryptionContext": "langgraph_sdk.encryption.types",
}
def __getattr__(name: str) -> object:
if name in _LAZY:
mod = importlib.import_module(_LAZY[name])
return getattr(mod, name)
msg = f"module {__name__!r} has no attribute {name!r}"
raise AttributeError(msg)
+6 -6
View File
@@ -262,7 +262,7 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.3.0"
version = "1.3.0a2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
@@ -274,14 +274,14 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/92/fe/20190232d9b513242899dbb0c2bb77e31b4d61e343743adbe90ebc2603d2/langchain_core-1.3.0.tar.gz", hash = "sha256:14a39f528bf459aa3aa40d0a7f7f1bae7520d435ef991ae14a4ceb74d8c49046", size = 860755, upload-time = "2026-04-17T14:51:38.298Z" }
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f8/e2/dbfa347aa072a6dc4cd38d6f9ebfc730b4c14c258c47f480f4c5c546f177/langchain_core-1.3.0-py3-none-any.whl", hash = "sha256:baf16ee028475df177b9ab8869a751c79406d64a6f12125b93802991b566cced", size = 515140, upload-time = "2026-04-17T14:51:36.274Z" },
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
]
[[package]]
name = "langgraph"
version = "1.1.9"
version = "1.1.7a2"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
@@ -294,7 +294,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
{ name = "langchain-core", specifier = "==1.3.0a2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "." },
@@ -413,7 +413,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.0.10"
version = "1.0.9"
source = { editable = "../prebuilt" }
dependencies = [
{ name = "langchain-core" },