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
25 changed files with 3731 additions and 57 deletions
+1
View File
@@ -100,3 +100,4 @@ dmypy.json
.turbo
.editorconfig
.scratch
.worktrees/
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 15 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` |
@@ -430,6 +430,43 @@ class PostgresSaver(BasePostgresSaver):
with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
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 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
)
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
),
)
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:
"""
Convert a database row into a CheckpointTuple object.
@@ -442,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": {
@@ -454,7 +498,7 @@ class PostgresSaver(BasePostgresSaver):
**value["checkpoint"],
"channel_values": {
**(value["checkpoint"].get("channel_values") or {}),
**self._load_blobs(value["channel_values"]),
**channel_values,
},
},
value["metadata"],
@@ -391,6 +391,43 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
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 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
)
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
),
)
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:
"""
Convert a database row into a CheckpointTuple object.
@@ -403,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"],
}
},
@@ -415,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"],
}
}
@@ -185,15 +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"
}
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,
@@ -371,3 +371,47 @@ async def test_get_checkpoint_no_channel_values(
checkpoint = await saver.aget_tuple(config)
assert checkpoint.checkpoint["channel_values"] == {}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
pytest.importorskip(
"langgraph.channels.delta", reason="langgraph core not installed"
)
from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
async with _saver(saver_name) as saver:
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "diff-channel-test-1"}}
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
await graph.ainvoke(
{"messages": [HumanMessage(content="there", id="h2")]}, config
)
state = await graph.aget_state(config)
msgs = state.values["messages"]
assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
assert msgs[0].content == "hi"
assert msgs[1].content == "reply-1"
assert msgs[2].content == "there"
assert msgs[3].content == "reply-3"
@@ -1,6 +1,7 @@
from __future__ import annotations
import copy
import dataclasses
import logging
from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence
from typing import ( # noqa: UP035
@@ -28,6 +29,26 @@ from langgraph.checkpoint.serde.types import (
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__)
@@ -457,6 +478,42 @@ class BaseCheckpointSaver(Generic[V]):
"""
raise NotImplementedError
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.
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 NotImplemented
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).
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.
Savers with a dedicated blob store (InMemorySaver, PostgresSaver)
should override this for O(1) performance.
"""
return NotImplemented
def get_next_version(self, current: V | None, channel: None) -> V:
"""Generate the next version ID for a channel.
@@ -126,12 +126,46 @@ class InMemorySaver(
channel_values: dict[str, Any] = {}
for k, v in versions.items():
kk = (thread_id, checkpoint_ns, k, v)
if kk in self.blobs:
vv = self.blobs[kk]
if vv[0] != "empty":
channel_values[k] = self.serde.loads_typed(vv)
if kk not in self.blobs:
continue
vv = self.blobs[kk]
if vv[0] != "empty":
channel_values[k] = self.serde.loads_typed(vv)
return channel_values
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, {})
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)
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"),
}
)
@@ -47,6 +47,12 @@ EMPTY_BYTES = b""
logger = logging.getLogger(__name__)
def _is_delta_value(obj: Any) -> bool:
from langgraph.checkpoint.base import DeltaValue # lazy import avoids circular dep
return isinstance(obj, DeltaValue)
class JsonPlusSerializer(SerializerProtocol):
"""Serializer that uses ormsgpack, with optional fallbacks.
@@ -239,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)
@@ -261,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:
+28
View File
@@ -983,3 +983,31 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
# No blocking should occur - inner is serialized as dict, not ext
assert "blocked" not in caplog.text.lower()
assert result == obj
def test_delta_value_serde_round_trip() -> None:
from langgraph.checkpoint.base import DeltaValue
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
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 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
+33
View File
@@ -308,3 +308,36 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
assert direct is not None
expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict()
assert direct.checkpoint["channel_values"]["foo"] == expected
class TestInMemorySaverDeltaChannel:
def test_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"
version = "00000000000000000000000000000001.0000000000000000"
delta = DeltaValue(delta=[{"content": "hi"}], prev_checkpoint_id=None)
saver.blobs[(thread_id, ns, channel, version)] = serde.dumps_typed(delta)
cp = empty_checkpoint()
cp["id"] = "cp1"
cp["channel_versions"][channel] = version
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp), serde.dumps_typed({}), None)
}
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_blob_missing(self) -> None:
"""get_channel_blob returns NotImplemented when checkpoint or channel not found."""
saver = InMemorySaver()
assert (
saver.get_channel_blob("t1", "", "no-such-cp", "messages") is NotImplemented
)
@@ -1,6 +1,7 @@
from langgraph.channels.any_value import AnyValue
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 (
@@ -20,6 +21,7 @@ __all__ = (
"UntrackedValue",
"EphemeralValue",
"BinaryOperatorAggregate",
"DeltaChannel",
"NamedBarrierValue",
"NamedBarrierValueAfterFinish",
# topics
@@ -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
+206
View File
@@ -0,0 +1,206 @@
from __future__ import annotations
import collections.abc
from collections.abc import Callable, Sequence
from typing import Any, Generic
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, _strip_extras
from langgraph.errors import EmptyChannelError
__all__ = ("DeltaChannel",)
class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]):
"""A channel that stores only per-step write deltas in checkpoints.
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).
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.
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.
Usage::
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",
"operator",
"snapshot_every",
"_pending",
"_base_version",
"_last_checkpoint_id",
"_overwritten",
"_steps_since_snapshot",
)
def __init__(
self,
operator: Callable[[list[Value], Any], list[Value]],
typ: type = list,
*,
snapshot_every: int | None = None,
) -> None:
typ = _strip_extras(typ)
if typ in (
collections.abc.Sequence,
collections.abc.MutableSequence,
):
typ = list
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_every != other.snapshot_every:
return False
if (
self.operator.__name__ != "<lambda>"
and other.operator.__name__ != "<lambda>"
):
return self.operator is other.operator
return True
@property
def ValueType(self) -> Any:
return list[self.typ] # type: ignore[name-defined]
@property
def UpdateType(self) -> Any:
return self.typ | list[self.typ] # type: ignore[name-defined]
def copy(self) -> Self:
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
new.key = self.key
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:
new = DeltaChannel(self.operator, self.typ, snapshot_every=self.snapshot_every)
new.key = self.key
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:
# 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 update(self, values: Sequence[Any]) -> bool:
if not values:
return False
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)
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) -> list[Value]:
if self.value is MISSING:
raise EmptyChannelError()
return self.value
def is_available(self) -> bool:
return self.value is not MISSING
def checkpoint(self) -> Any:
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,
)
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
+180 -8
View File
@@ -1,9 +1,17 @@
from __future__ import annotations
import logging
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Any
from langgraph.checkpoint.base import Checkpoint
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
DeltaChainValue,
DeltaValue,
)
from langgraph.checkpoint.base.id import uuid6
from langgraph._internal._typing import MISSING
@@ -12,6 +20,171 @@ from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 4
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:
return Checkpoint(
@@ -67,13 +240,12 @@ def channels_from_checkpoint(
channel_specs[k] = v
else:
managed_specs[k] = v
return (
{
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
for k, v in channel_specs.items()
},
managed_specs,
)
channels: dict[str, BaseChannel] = {}
for k, 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
def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
+33
View File
@@ -92,6 +92,8 @@ from langgraph.pregel._algo import (
task_path_str,
)
from langgraph.pregel._checkpoint import (
_aassemble_delta_channels,
_assemble_delta_channels,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -881,6 +883,12 @@ class PregelLoop:
id=self.checkpoint["id"] if exiting else None,
updated_channels=self.updated_channels,
)
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()
@@ -1262,6 +1270,19 @@ 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
)
@@ -1466,6 +1487,18 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.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
)
+54 -10
View File
@@ -122,6 +122,8 @@ from langgraph.pregel._algo import (
)
from langgraph.pregel._call import identifier
from langgraph.pregel._checkpoint import (
_aassemble_delta_channels,
_assemble_delta_channels,
channels_from_checkpoint,
copy_checkpoint,
create_checkpoint,
@@ -1049,13 +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,
checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1168,13 +1180,23 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
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,
checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1520,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 {}
)
@@ -1966,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 {}
)
+405
View File
@@ -117,3 +117,408 @@ def test_untracked_value() -> None:
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
with pytest.raises(EmptyChannelError):
new_channel.get()
def test_delta_channel_basic_two_steps() -> None:
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import 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 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.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_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"
assert msgs[1].content == "hello"
assert msgs[2].content == "bye"
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
from langchain_core.messages import HumanMessage
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
# Old BinaryOperatorAggregate checkpoint: plain list
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_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 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_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")]
# 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_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")]
# 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_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
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"
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
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=2)]
counter = {"n": 0}
def respond(state: State) -> dict:
counter["n"] += 1
return {
"messages": [
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
]
}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "snap-test"}}
# 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)
msgs = state.values["messages"]
# 5 human + 5 AI = 10 total
assert len(msgs) == 10, f"expected 10 messages, got {len(msgs)}: {msgs}"
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
from langgraph.checkpoint.base import (
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": "one"}
msg2 = {"type": "ai", "content": "two"}
msg3 = {"type": "human", "content": "three"}
# 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"
)
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_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
from langgraph.checkpoint.base import DeltaValue, empty_checkpoint
from langgraph.pregel._checkpoint import _assemble_delta_channels
cp = empty_checkpoint()
cp["id"] = "cp2"
cp["channel_values"]["messages"] = DeltaValue(
delta=["msg2"], prev_checkpoint_id="cp-missing"
)
saver = MagicMock()
saver.get_channel_blob.return_value = NotImplemented
saver.get_tuple.return_value = None # checkpoint not found
config = {"configurable": {"thread_id": "t1", "checkpoint_ns": ""}}
assembled = _assemble_delta_channels(cp, config, saver)
# Should still assemble — with partial chain (just the current delta, base=None)
assert "messages" in assembled
from langgraph.checkpoint.base import DeltaChainValue
chain = assembled["messages"]
assert isinstance(chain, DeltaChainValue)
assert chain.base is None
assert chain.deltas == [["msg2"]]
@@ -0,0 +1,356 @@
"""Benchmark: DeltaChannel vs BinaryOperatorAggregate storage and time.
Run directly: python tests/test_delta_channel_benchmark.py
Run via pytest: pytest tests/test_delta_channel_benchmark.py -s
Simulates realistic multi-turn conversations with paragraph-length messages
(~100 tokens each) scaling up to 1M-token-equivalent histories.
Token estimates: 1 token 4 chars; each turn 200 tokens (human + AI).
A 1M-token conversation 5,000 turns of realistic messages.
"""
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
try:
from langgraph.checkpoint.sqlite import SqliteSaver
_SQLITE_AVAILABLE = True
except ImportError:
_SQLITE_AVAILABLE = False
SNAPSHOT_EVERY = 50
# ---------------------------------------------------------------------------
# Realistic message payload (~100 tokens / ~400 chars each)
# ---------------------------------------------------------------------------
_HUMAN_TEMPLATE = (
"I need help understanding the implications of {topic} on our system architecture. "
"Specifically, I'm concerned about how this interacts with our existing {concern} "
"and whether we need to refactor the {component} layer before proceeding."
)
_AI_TEMPLATE = (
"Great question about {topic}. The key insight here is that {concern} introduces "
"a subtle ordering dependency that most teams overlook until they hit it in production. "
"For your {component} layer specifically, I'd recommend starting with a careful audit "
"of the interface boundaries before making any structural changes. This will give you "
"a clear picture of the blast radius and let you sequence the migration safely."
)
_TOPICS = [
"distributed tracing",
"eventual consistency",
"schema migration",
"backpressure handling",
"idempotency guarantees",
"cache invalidation",
"connection pooling",
"rate limiting",
"circuit breaking",
"observability pipelines",
]
_CONCERNS = [
"concurrency model",
"retry semantics",
"state management",
"error propagation",
"latency budget",
]
_COMPONENTS = [
"persistence",
"routing",
"ingestion",
"aggregation",
"serialization",
]
def _human_content(i: int) -> str:
return _HUMAN_TEMPLATE.format(
topic=_TOPICS[i % len(_TOPICS)],
concern=_CONCERNS[i % len(_CONCERNS)],
component=_COMPONENTS[i % len(_COMPONENTS)],
)
def _ai_content(i: int) -> str:
return _AI_TEMPLATE.format(
topic=_TOPICS[i % len(_TOPICS)],
concern=_CONCERNS[i % len(_CONCERNS)],
component=_COMPONENTS[i % len(_COMPONENTS)],
)
# ---------------------------------------------------------------------------
# State definitions
# ---------------------------------------------------------------------------
class BinaryState(TypedDict):
messages: Annotated[list, add_messages]
class DeltaState(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
class DeltaSnapshotState(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages, snapshot_every=SNAPSHOT_EVERY)]
# ---------------------------------------------------------------------------
# Graph factory
# ---------------------------------------------------------------------------
def _make_graph(state_cls: type, checkpointer: Any = None) -> Any:
def human_node(state: Any) -> dict:
return {}
def ai_node(state: Any) -> dict:
i = len(state["messages"]) // 2
return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]}
g = StateGraph(state_cls)
g.add_node("human", human_node)
g.add_node("ai", ai_node)
g.add_edge("human", "ai")
g.add_edge("ai", END)
g.set_entry_point("human")
return g.compile(checkpointer=checkpointer or MemorySaver())
# ---------------------------------------------------------------------------
# Measurement helpers
# ---------------------------------------------------------------------------
def _total_blob_bytes(saver: MemorySaver) -> int:
total = 0
for (_, _, _, _), (type_tag, blob) in saver.blobs.items():
if blob is not None:
total += len(blob)
return total
def _run_turns(
n_turns: int,
state_cls: type,
checkpointer: Any = None,
) -> tuple[float, float, int]:
"""Run n_turns conversation turns.
Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes).
blob_bytes is -1 for savers without in-memory blob stores (e.g. SQLite).
Read latency is measured as the time to invoke the graph with no new
messages after the full history is built this forces state rehydration.
"""
graph = _make_graph(state_cls, checkpointer)
config = {"configurable": {"thread_id": "bench"}}
t0 = time.perf_counter()
for i in range(n_turns):
graph.invoke(
{"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]},
config,
)
write_elapsed = time.perf_counter() - t0
# Measure read/rehydration: get_state forces the channel to rebuild
t1 = time.perf_counter()
for _ in range(5):
graph.get_state(config)
read_elapsed = (time.perf_counter() - t1) / 5
if isinstance(graph.checkpointer, MemorySaver):
blob_bytes = _total_blob_bytes(graph.checkpointer)
else:
blob_bytes = -1
return write_elapsed, read_elapsed, blob_bytes
def _fmt_bytes(n: int) -> str:
if n >= 1_000_000:
return f"{n / 1_000_000:.1f} MB"
if n >= 1_000:
return f"{n / 1_000:.1f} KB"
return f"{n} B"
def _approx_tokens(n_turns: int) -> str:
# ~100 tokens human + ~100 tokens AI per turn
tokens = n_turns * 200
if tokens >= 1_000_000:
return f"~{tokens / 1_000_000:.1f}M tok"
if tokens >= 1_000:
return f"~{tokens / 1_000:.0f}K tok"
return f"~{tokens} tok"
# ---------------------------------------------------------------------------
# Benchmark matrix
# ---------------------------------------------------------------------------
# Turn counts chosen to 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 _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
factories.append(("SQLite", tempfile.NamedTemporaryFile(suffix=".db")))
return factories
def run_benchmark() -> None:
print()
print(
"DeltaChannel vs add_messages (BinaryOperatorAggregate) — checkpoint storage & latency"
)
print("Simulating realistic multi-turn conversations up to ~1M-token histories")
print("(5,000 turns × ~200 tokens/turn ≈ 1M tokens — Claude's full context window)")
print()
checkpointers: list[tuple[str, Any]] = [("InMemory (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_benchmark_for_checkpointer(cp_hint: Any) -> None:
import contextlib
import tempfile
@contextlib.contextmanager
def _make_saver():
if cp_hint is None:
yield None
else:
with tempfile.NamedTemporaryFile(suffix=".db") as f:
with SqliteSaver.from_conn_string(f.name) as saver:
yield saver
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)
with _make_saver() as saver:
s_wt, s_rt, s_bytes = _run_turns(turns, DeltaSnapshotState, saver)
# 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:
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"{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("=" * 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" Storage: {_fmt_bytes(b_bytes)} (add_messages) → {_fmt_bytes(s_bytes)} (DeltaChannel+snapshot) — {ratio:.0f}x reduction"
)
print(
f" Read latency: {b_rt * 1000:.1f}ms (add_messages) vs {s_rt * 1000:.1f}ms (DeltaChannel+snapshot)"
)
print()
print("Legend:")
print(
" 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()
# ---------------------------------------------------------------------------
# Pytest entry point
# ---------------------------------------------------------------------------
def test_delta_channel_benchmark(capsys: Any) -> None:
"""Storage grows O(N²) for add_messages, O(N) for DeltaChannel."""
with capsys.disabled():
run_benchmark()
# Correctness assertion: DeltaChannel must use less storage at scale.
for turns in [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}"
)
assert s_bytes < b_bytes, (
f"DeltaChannel+snapshot should use less storage at {turns} turns, "
f"got snapshot={s_bytes} binary={b_bytes}"
)
# ---------------------------------------------------------------------------
# Script entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
run_benchmark()
sys.exit(0)
+187
View File
@@ -9400,3 +9400,190 @@ def test_fork_does_not_apply_pending_writes(
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121}
async def test_delta_channel_end_to_end_inmemory() -> None:
"""Full graph run: DeltaChannel accumulates correctly across multiple turns."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
n = len(state["messages"])
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-test-1"}}
# Turn 1
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
# Turn 2
graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config)
# Turn 3
graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# 3 human + 3 AI = 6 total
assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}"
assert msgs[0].content == "hello"
assert msgs[2].content == "world"
assert msgs[4].content == "bye"
assert msgs[1].content == "reply-1"
assert msgs[3].content == "reply-3"
assert msgs[5].content == "reply-5"
async def test_delta_channel_time_travel() -> None:
"""Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas."""
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
counter = {"n": 0}
def respond(state: State) -> dict:
counter["n"] += 1
return {
"messages": [
AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}")
]
}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_edge(START, "respond")
saver = InMemorySaver()
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "diff-time-travel"}}
# Run 2 turns: h1→ai-1, h2→ai-2
graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config)
graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config)
# Find the checkpoint after turn 1 (2 messages: h1 + ai-1)
history = list(graph.get_state_history(config))
after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2)
assert len(after_turn1.values["messages"]) == 2
assert after_turn1.values["messages"][0].content == "h1"
assert after_turn1.values["messages"][1].content == "ai-1"
# Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N)
# NOT 5 messages (turn-2 deltas must not bleed into the resumed run)
result = graph.invoke(
{"messages": [HumanMessage(content="h3", id="h3")]},
after_turn1.config,
)
msgs = result["messages"]
# Should be: h1, ai-1, h3, ai-N — 4 messages total
assert len(msgs) == 4, (
f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}"
)
assert msgs[0].content == "h1"
assert msgs[1].content == "ai-1"
assert msgs[2].content == "h3"
async def test_delta_channel_remove_message_end_to_end() -> None:
"""RemoveMessage inside a DeltaChannel graph must persist and reload correctly."""
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def respond(state: State) -> dict:
return {"messages": [AIMessage(content="reply", id="ai-1")]}
def delete_first(state: State) -> dict:
# removes the first message
return {"messages": [RemoveMessage(id=state["messages"][0].id)]}
builder = StateGraph(State)
builder.add_node("respond", respond)
builder.add_node("delete_first", delete_first)
builder.add_edge(START, "respond")
builder.add_edge("respond", "delete_first")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-remove-test"}}
graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# h1 was removed, only ai-1 should remain
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
assert msgs[0].id == "ai-1"
# A subsequent turn must reconstruct from the checkpoint correctly
graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
# ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed
# more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2
# just assert h1 is still gone
assert all(m.id != "h1" for m in msgs), (
"h1 should still be absent after second turn"
)
async def test_delta_channel_update_by_id_end_to_end() -> None:
"""Updating a message by ID via DeltaChannel must persist and reload correctly."""
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
class State(TypedDict):
messages: Annotated[list, DeltaChannel(add_messages)]
def update_msg(state: State) -> dict:
# re-send h1 with updated content
return {"messages": [HumanMessage(content="updated", id="h1")]}
builder = StateGraph(State)
builder.add_node("update_msg", update_msg)
builder.add_edge(START, "update_msg")
graph = builder.compile(checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "diff-update-id-test"}}
graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}"
assert msgs[0].content == "updated"
assert msgs[0].id == "h1"
# Second turn: verify the updated state is the base for further accumulation
graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config)
state = graph.get_state(config)
msgs = state.values["messages"]
ids = [m.id for m in msgs]
assert "h1" in ids # h1 persists (updated, not duplicated)
assert "h2" in ids
assert ids.count("h1") == 1, "h1 must not be duplicated"
@@ -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)
+7 -19
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):
@@ -1334,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:
@@ -1362,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:
@@ -1866,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)
@@ -1881,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
@@ -1900,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,
)
+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)