From 86baa5d08eee6f1efe89d4c023e785d7c5c8e77a Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Tue, 5 May 2026 15:30:57 -0400 Subject: [PATCH] feat(checkpoint-sqlite): override get_delta_channel_history with streaming walk (#7702) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds a sqlite-specific override of `BaseCheckpointSaver.get_delta_channel_history` (and async). Before this PR, `SqliteSaver` / `AsyncSqliteSaver` inherited the default impl, which calls `get_tuple` once per ancestor — N round-trips, full pending-writes fetch per step regardless of channel relevance. The override mirrors the postgres two-stage shape (ancestor walk + per-channel UNION ALL writes fetch) but adapted for sqlite: - **No JSONB** → stage 1 streams the cursor row-by-row in `checkpoint_id` DESC order. The merged walk advances one row at a time, deserializing only on-path checkpoints and dropping each before advancing — peak in-flight is one deserialized checkpoint, no `fetchall()` materialization. - **No separate blob table** → `channel_values` lives inline in the checkpoint blob, so seeds come back from stage 1 with no second fetch. - **Single merged walk (not K independent walks)**: each visited cid is deserialized exactly once, regardless of how many channels are still seeking their seed. - **Stage 2** stays per-channel UNION ALL to avoid over-fetching writes when channels have different chain depths — same rationale as postgres. `AsyncSqliteSaver.get_delta_channel_history` bridges to its async form via `run_coroutine_threadsafe`, matching the same cross-thread guard used by `get_tuple` / `delete_thread`. ## Tests - New `tests/test_delta_channel_migration.py`: covers the `BinaryOperatorAggregate -> DeltaChannel` migration path on sqlite (sync round-trip, sync continuation with post-migration delta folding, async round-trip). Mirrors `libs/langgraph/tests/test_delta_channel_migration.py` (which covered `InMemorySaver`); without these, the override's behavior on pre-migration threads was unverified — the override has to identify a plain accumulated `channel_values[ch]` at a pre-migration ancestor as a valid `seed`, not just `_DeltaSnapshot` sentinels. - Existing `tests/test_get_delta_channel_history.py` (7 tests) continues to pass and now exercises the optimized override end-to-end (previously hit the inherited default impl). - `make format`, `make lint`, `make test`: clean. 97/97 in the non-flaky sqlite suite (the one ignored test, `test_async_asearch_refresh_ttl`, is a known TTL-store timing flake on a separate module unrelated to this PR). ## Benchmarks ### `get_delta_channel_history` micro-bench (override vs inherited default impl) 1000-turn synthetic threads with sentinel snapshots + per-step writes; `bench_sqlite_delta_history.py`. Per-call latency in microseconds. | Scenario | min | median | mean | |---|---:|---:|---:| | S1 single channel, root-only snapshot | **4.60x** | **4.90x** | **5.13x** | | S2 mixed cadence (every-50 + root-only), 2 channels | **6.08x** | **6.37x** | **6.84x** | | S3 K=8 channels, root-only snapshot | 1.23x | 1.27x | 0.90x | S2 wins biggest because per-channel UNION ALL avoids over-fetching writes for the shallow channel. S3 is the worst case for sqlite (8 channels all walking to root, 1000 deserializations either way) — the override still wins on min/median. ### Long-running thread mem/storage bench (delta vs no-delta) `bench_sqlite_delta_memory.py`. `delta` mode uses `DeltaChannel` + the override; `no_delta` uses `Annotated[list, _messages_delta_reducer]` (full state in every blob). Same workload, file-backed sqlite. Latency measured untraced (30 iterations); peak heap measured separately under tracemalloc. | Scenario | Turns | Storage Δ | Peak heap Δ | Read latency Δ | |---|---:|---|---|---| | K=1, freq=50 | 200 | **-96%** (942 KB vs 25.1 MB) | +21% (504 KB vs 418 KB) | **+13%** | | K=1, freq=50 | 500 | **-98%** (2.9 MB vs 152.3 MB) | +20% (1.2 MB vs 1.0 MB) | **-6%** (delta wins) | | K=3, freq=50 uniform | 200 | **-98%** (1.7 MB vs 73.5 MB) | +7% (1.3 MB vs 1.2 MB) | **+10%** | | K=3, freq=50 uniform | 500 | **-99%** (6.0 MB vs 452.5 MB) | +7% (3.3 MB vs 3.0 MB) | **+6%** | | K=3, freq=mixed | 200 | **-98%** (1.4 MB vs 73.5 MB) | +5% (1.3 MB vs 1.2 MB) | +190% (5.1 ms vs 1.7 ms abs) | | K=3, freq=mixed | 500 | **-99%** (4.1 MB vs 452.5 MB) | +8% (3.3 MB vs 3.0 MB) | +377% (20.9 ms vs 4.4 ms abs) | - **Storage**: -96 to -99% on long threads (a 500-turn K=3 thread shrinks from 452 MB to 6 MB on disk). This is the headline win. - **Peak heap**: within +5 to +21% of the no-delta path — the streaming cursor + merged walk + drop-after-deserialize keep peak in-flight at one checkpoint at a time. - **Read latency**: equivalent-ish (within ~15%) on uniform-cadence scenarios; at K=1/500 turns delta even wins by 6%. The mixed-cadence rows have one channel with `snapshot_frequency=1000` walking to root on a 500-turn thread — by configuration. Absolute mixed-delta latency is still 5-21 ms per read. Bench scripts (not committed; workspace-root convention matches other `bench_*.py` files): - `bench_sqlite_delta_history.py` - `bench_sqlite_delta_memory.py` ## Test plan - [x] `cd libs/checkpoint-sqlite && make format` clean - [x] `cd libs/checkpoint-sqlite && make lint` clean - [x] `cd libs/checkpoint-sqlite && make test` — 97 passed (1 known flake unrelated) - [x] `tests/test_get_delta_channel_history.py` — 7/7 (now exercises the override) - [x] `tests/test_delta_channel_migration.py` — 3/3 (new) --- .../langgraph/checkpoint/sqlite/__init__.py | 91 ++++++++- .../langgraph/checkpoint/sqlite/_delta.py | 172 ++++++++++++++++++ .../langgraph/checkpoint/sqlite/aio.py | 109 ++++++++++- libs/checkpoint-sqlite/pyproject.toml | 4 +- .../tests/test_delta_channel_migration.py | 170 +++++++++++++++++ libs/checkpoint-sqlite/uv.lock | 2 +- libs/langgraph/uv.lock | 2 +- libs/prebuilt/uv.lock | 2 +- 8 files changed, 545 insertions(+), 7 deletions(-) create mode 100644 libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py create mode 100644 libs/checkpoint-sqlite/tests/test_delta_channel_migration.py diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 2684ce164..6ca2448d0 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -4,7 +4,7 @@ import json import random import sqlite3 import threading -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from contextlib import closing, contextmanager from typing import Any, cast @@ -16,12 +16,19 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, + DeltaChannelHistory, SerializerProtocol, get_checkpoint_id, get_checkpoint_metadata, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.checkpoint.sqlite._delta import ( + DELTA_STAGE1_SQL, + build_delta_channels_writes_history, + build_delta_stage2_sql, + step_walk_with_row, +) from langgraph.checkpoint.sqlite.utils import search_where _AIO_ERROR_MSG = ( @@ -493,6 +500,88 @@ class SqliteSaver(BaseCheckpointSaver[str]): (str(thread_id),), ) + def get_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`. + + Two-stage query: + + * Stage 1 (paged): newest-first slice of `checkpoints` returning + `(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per + ancestor. Sqlite has no JSONB, so we ship the full serialized + checkpoint blob and inspect `channel_values` in Python. Pages + newest-first by `checkpoint_id` with a `< cursor` predicate; + page size is `DELTA_PAGE_SIZE`. Stops paging when every channel + has found its seed or the chain is exhausted. + + * Stage 2 (per-channel UNION ALL): one branch per channel reading + `writes` filtered to that channel's specific `chain_cids`. No + separate seed-blob fetch — sqlite stores `channel_values` inline + in the checkpoint blob, so seeds come back from stage 1. + """ + if not channels: + return {} + channels = list(channels) + thread_id = str(config["configurable"]["thread_id"]) + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id is None: + target = self.get_tuple(config) + if target is None: + return {ch: {"writes": []} for ch in channels} + checkpoint_id = target.config["configurable"]["checkpoint_id"] + + chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} + seed_val_by_ch: dict[str, Any] = {} + walk_state: dict[str, Any] = {} + seeded: set[str] = set() + + with self.cursor(transaction=False) as cur: + cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)) + for row in cur: + cid, parent_cid, type_tag, blob = row + if step_walk_with_row( + cid=cid, + parent_cid=parent_cid, + type_tag=type_tag, + blob=blob, + target_id=checkpoint_id, + serde=self.serde, + chain_by_ch=chain_by_ch, + seed_val_by_ch=seed_val_by_ch, + walk_state=walk_state, + seeded=seeded, + channels=channels, + ): + break + + channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] + stage2_sql = build_delta_stage2_sql( + chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain], + ) + if stage2_sql: + stage2_params: list[Any] = [] + for ch in channels_with_chain: + stage2_params.extend( + [thread_id, checkpoint_ns, ch, *chain_by_ch[ch]] + ) + cur.execute(stage2_sql, stage2_params) + stage2_rows = cast( + "list[tuple[str, str, str, int, str, bytes]]", cur.fetchall() + ) + else: + stage2_rows = [] + + return build_delta_channels_writes_history( + channels=channels, + chain_by_ch=chain_by_ch, + seed_val_by_ch=seed_val_by_ch, + seeded=seeded, + stage2_rows=stage2_rows, + serde=self.serde, + ) + async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the database asynchronously. diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py new file mode 100644 index 000000000..1fe617ff7 --- /dev/null +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py @@ -0,0 +1,172 @@ +"""Shared helpers for `get_delta_channel_history` on sqlite savers. + +Mirrors the two-stage shape of `BasePostgresSaver` (ancestor walk + +per-channel UNION ALL writes fetch), but adapted for sqlite's +constraints. The structural differences: + +* No JSONB — to inspect `channel_values` for a checkpoint we must + deserialize the full blob. Stage 1 streams the cursor row-by-row and + deserializes only the rows the merged walk visits, freeing each blob + before advancing. +* No separate blob table — `channel_values` lives inline in the + checkpoint, so seeds come back from stage 1 with no second fetch. +* Single merged walk (not K independent walks): each visited cid is + deserialized exactly once, regardless of how many channels are still + seeking their seed. + +The streaming design keeps peak in-flight memory at roughly one +deserialized checkpoint at a time, instead of holding the entire +ancestor chain's worth of raw blobs as a `fetchall()`-materialized list. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite + +# Stage 1 streams ancestors of `target_cid` newest-first. The `<=` +# predicate keeps target itself in the stream so we can read its +# `parent_checkpoint_id` from the first row without a separate lookup; +# the caller skips target's own writes/seed (matches the +# `BaseCheckpointSaver` contract). +DELTA_STAGE1_SQL = ( + "SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint " + "FROM checkpoints " + "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? " + "ORDER BY checkpoint_id DESC" +) + + +def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str: + """Stage-2 per-channel UNION ALL fetching writes from `writes`. + + One branch per channel with a non-empty chain. Each branch inlines its + own `IN (?, ?, ...)` placeholder list because sqlite has no array-bind + equivalent of postgres's `= ANY(%s)`. Caller passes parameters in + matching order: `[thread_id, checkpoint_ns, channel, *chain_cids]` per + branch. + + Returns an empty string when no channel has a chain (caller skips + executing in that case). Per-channel UNION ALL avoids the over-fetch + of a single `channel = ANY(channels)` filter when channels have + different chain depths — same rationale as postgres. + """ + branches: list[str] = [] + for n in chain_lens: + cid_placeholders = ",".join("?" * n) + branches.append( + "SELECT checkpoint_id, channel, task_id, idx, type, value " + "FROM writes " + "WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ? " + f"AND checkpoint_id IN ({cid_placeholders})" + ) + return " UNION ALL ".join(branches) + + +def step_walk_with_row( + *, + cid: str, + parent_cid: str | None, + type_tag: str, + blob: bytes, + target_id: str, + serde: Any, + chain_by_ch: dict[str, list[str]], + seed_val_by_ch: dict[str, Any], + walk_state: dict[str, Any], + seeded: set[str], + channels: Sequence[str], +) -> bool: + """Process one streamed stage-1 row in the merged ancestor walk. + + The cursor returns (cid, parent_cid, type, blob) rows in + `checkpoint_id` DESC order starting at target. The first row is + target itself; we read its parent_cid to seed the walk and otherwise + skip it (target's own writes/seed are not part of the contract). + + For each subsequent row, if `cid` matches the walk's current + position, we deserialize the blob, append the cid to every + not-yet-seeded channel's chain, and check `channel_values` for + seeds. The deserialized checkpoint is dropped before advancing — no + cross-row cache, so peak in-flight is one deserialized checkpoint. + + Off-path rows (different branch on the same thread) advance the + cursor without doing any work. + + Returns True when every requested channel is seeded — the caller + can stop iterating and close the cursor. + """ + if "started" not in walk_state: + if cid == target_id: + walk_state["started"] = True + walk_state["cur_cid"] = parent_cid + walk_state["active"] = {ch for ch in channels if ch not in seeded} + # Not target yet (or target not present): keep streaming. + return False + active: set[str] = walk_state["active"] + if not active: + return True + if cid != walk_state["cur_cid"]: + # Off-path row from a sibling branch — skip without deserializing. + return False + for ch in active: + chain_by_ch[ch].append(cid) + ckpt = serde.loads_typed((type_tag, blob)) + channel_values: Mapping[str, Any] = ckpt.get("channel_values") or {} + for ch in [ch for ch in active if ch in channel_values]: + seed_val_by_ch[ch] = channel_values[ch] + seeded.add(ch) + active.discard(ch) + del ckpt, channel_values + walk_state["cur_cid"] = parent_cid + return not active + + +def build_delta_channels_writes_history( + *, + channels: Sequence[str], + chain_by_ch: Mapping[str, list[str]], + seed_val_by_ch: Mapping[str, Any], + seeded: set[str], + stage2_rows: Sequence[tuple[str, str, str, int, str, bytes]], + serde: Any, +) -> dict[str, DeltaChannelHistory]: + """Demux stage-2 rows per channel; produce per-channel histories. + + Stage-2 rows are `(checkpoint_id, channel, task_id, idx, type, value)`. + Final write order is oldest→newest globally and `(task_id, idx)` within + a checkpoint, matching the contract on `DeltaChannelHistory.writes`. + + `seed` is omitted when the walk reached a true root with no snapshot + found (channel never entered `seeded`); consumers treat absence as + "start empty". + """ + writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = { + ch: {} for ch in channels + } + for cid, ch, task_id, idx, type_tag, value_blob in stage2_rows: + writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append( + (type_tag, value_blob, task_id, idx) + ) + for cid_map in writes_by_ch_by_cid.values(): + for ws in cid_map.values(): + ws.sort(key=lambda w: (w[2], w[3])) + + result: dict[str, DeltaChannelHistory] = {} + for ch in channels: + chain_cids = chain_by_ch.get(ch, []) + cid_writes = writes_by_ch_by_cid.get(ch, {}) + collected: list[PendingWrite] = [] + # Chain is newest-first; iterate oldest-first for the public order. + for cid in reversed(chain_cids): + for type_tag, value_blob, task_id, _idx in cid_writes.get(cid, []): + collected.append( + (task_id, ch, serde.loads_typed((type_tag, value_blob))) + ) + entry: DeltaChannelHistory = {"writes": collected} + if ch in seeded: + entry["seed"] = seed_val_by_ch[ch] + result[ch] = entry + return result diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index e20b11f4c..368428c68 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -4,7 +4,7 @@ import asyncio import json import random import threading -from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from typing import Any, TypeVar, cast @@ -17,12 +17,19 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, + DeltaChannelHistory, SerializerProtocol, get_checkpoint_id, get_checkpoint_metadata, ) from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.checkpoint.sqlite._delta import ( + DELTA_STAGE1_SQL, + build_delta_channels_writes_history, + build_delta_stage2_sql, + step_walk_with_row, +) from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=Callable) @@ -272,6 +279,29 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): self.adelete_thread(thread_id), self.loop ).result() + def get_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Sync bridge to `aget_delta_channel_history`. + + Mirrors the same cross-thread guard as `get_tuple` / + `delete_thread` — calling from the loop thread raises rather than + deadlocking. + """ + try: + if asyncio.get_running_loop() is self.loop: + raise asyncio.InvalidStateError( + "Synchronous calls to AsyncSqliteSaver are only allowed from a " + "different thread. From the main thread, use the async interface. " + "For example, use `await checkpointer.aget_delta_channel_history(...)`." + ) + except RuntimeError: + pass + return asyncio.run_coroutine_threadsafe( + self.aget_delta_channel_history(config=config, channels=channels), + self.loop, + ).result() + async def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -589,6 +619,83 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]): ) await self.conn.commit() + async def aget_delta_channel_history( + self, *, config: RunnableConfig, channels: Sequence[str] + ) -> Mapping[str, DeltaChannelHistory]: + """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`. + + See `SqliteSaver.get_delta_channel_history` for design notes; this + is the async equivalent using `aiosqlite` cursors. Stage 1 pages + the parent chain newest-first and Python-deserializes each + checkpoint blob to find per-channel snapshots; stage 2 fetches + only the relevant writes via per-channel UNION ALL. + """ + if not channels: + return {} + channels = list(channels) + await self.setup() + thread_id = str(config["configurable"]["thread_id"]) + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id is None: + target = await self.aget_tuple(config) + if target is None: + return {ch: {"writes": []} for ch in channels} + checkpoint_id = target.config["configurable"]["checkpoint_id"] + + chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels} + seed_val_by_ch: dict[str, Any] = {} + walk_state: dict[str, Any] = {} + seeded: set[str] = set() + + async with self.lock, self.conn.cursor() as cur: + await cur.execute( + DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id) + ) + async for row in cur: + cid, parent_cid, type_tag, blob = row + if step_walk_with_row( + cid=cid, + parent_cid=parent_cid, + type_tag=type_tag, + blob=blob, + target_id=checkpoint_id, + serde=self.serde, + chain_by_ch=chain_by_ch, + seed_val_by_ch=seed_val_by_ch, + walk_state=walk_state, + seeded=seeded, + channels=channels, + ): + break + + channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] + stage2_sql = build_delta_stage2_sql( + chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain], + ) + if stage2_sql: + stage2_params: list[Any] = [] + for ch in channels_with_chain: + stage2_params.extend( + [thread_id, checkpoint_ns, ch, *chain_by_ch[ch]] + ) + await cur.execute(stage2_sql, stage2_params) + stage2_rows = cast( + "list[tuple[str, str, str, int, str, bytes]]", + await cur.fetchall(), + ) + else: + stage2_rows = [] + + return build_delta_channels_writes_history( + channels=channels, + chain_by_ch=chain_by_ch, + seed_val_by_ch=seed_val_by_ch, + seeded=seeded, + stage2_rows=stage2_rows, + serde=self.serde, + ) + def get_next_version(self, current: str | None, channel: None) -> str: """Generate the next version ID for a channel. diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index 39bfc07aa..40d05a6b4 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "langgraph-checkpoint-sqlite" -version = "3.0.3" +version = "3.1.0a1" description = "Library with a SQLite implementation of LangGraph checkpoint saver." authors = [] requires-python = ">=3.10" @@ -12,7 +12,7 @@ readme = "README.md" license = "MIT" license-files = ['LICENSE'] dependencies = [ - "langgraph-checkpoint>=3,<5.0.0", + "langgraph-checkpoint>=4.1.0a4,<5.0.0", "aiosqlite>=0.20", "sqlite-vec>=0.1.6", ] diff --git a/libs/checkpoint-sqlite/tests/test_delta_channel_migration.py b/libs/checkpoint-sqlite/tests/test_delta_channel_migration.py new file mode 100644 index 000000000..c71f3b12c --- /dev/null +++ b/libs/checkpoint-sqlite/tests/test_delta_channel_migration.py @@ -0,0 +1,170 @@ +"""Sqlite-specific migration smoke tests: BinaryOperatorAggregate -> DeltaChannel. + +Mirrors `libs/langgraph/tests/test_delta_channel_migration.py` (which +covers `InMemorySaver` + a third-party fallback to the base default +impl). This file exercises the same migration scenario through the +sqlite-specific `SqliteSaver.get_delta_channel_history` override — +specifically that the streaming ancestor walk finds a pre-migration +plain `channel_values[ch]` entry and surfaces it as the `seed`, with +post-migration writes folding on top through the reducer. + +Pre-migration checkpoints under `BinaryOperatorAggregate` carry the +full accumulated value at every settled super-step boundary. The +override has to identify those as "real" seeds (not `_DeltaSnapshot` +sentinels) — the saver layer is intentionally delta-agnostic and just +returns whatever is stored in `channel_values[ch]`. +""" + +from __future__ import annotations + +import operator +from typing import Annotated, Any + +import pytest +from langchain_core.runnables import RunnableConfig + +# `langgraph` core isn't a dep of `langgraph-checkpoint-sqlite`. Skip the +# whole module rather than importerror-ing in the standalone CI shape. +pytest.importorskip("langgraph.channels.delta", reason="langgraph core not installed") +pytest.importorskip("langgraph.channels.binop", reason="langgraph core not installed") +pytest.importorskip("langgraph.graph", reason="langgraph core not installed") + +from langgraph.channels.binop import BinaryOperatorAggregate # type: ignore[import-untyped] # noqa: E402,I001 +from langgraph.channels.delta import DeltaChannel # type: ignore[import-untyped] # noqa: E402 +from langgraph.graph import END, START, StateGraph # type: ignore[import-untyped] # noqa: E402 +from typing_extensions import TypedDict # noqa: E402 + +from langgraph.checkpoint.sqlite import SqliteSaver # noqa: E402 +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver # noqa: E402 + +pytestmark = pytest.mark.anyio + + +def _noop(_state: Any) -> dict: + return {} + + +def _list_concat(state: list, writes: list) -> list: + result = list(state) + for w in writes: + result.extend(w if isinstance(w, list) else [w]) + return result + + +def _binop_graph(checkpointer: Any) -> Any: + class BinopState(TypedDict): + items: Annotated[list, BinaryOperatorAggregate(list, operator.add)] + + return ( + StateGraph(BinopState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _delta_graph(checkpointer: Any) -> Any: + class DeltaState(TypedDict): + items: Annotated[list, DeltaChannel(_list_concat)] + + return ( + StateGraph(DeltaState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _drive(graph: Any, config: RunnableConfig, tag: str, n: int) -> None: + for i in range(n): + graph.invoke({"items": [f"{tag}{i}"]}, config) + + +async def _adrive(graph: Any, config: RunnableConfig, tag: str, n: int) -> None: + for i in range(n): + await graph.ainvoke({"items": [f"{tag}{i}"]}, config) + + +def _settled_boundaries(history: list) -> list[tuple[RunnableConfig, list]]: + """`(config, items)` for every checkpoint with `next == ('__start__',)` + — the stable inter-invoke boundaries that round-trip predictably. + """ + return [ + (s.config, list(s.values.get("items", []))) + for s in history + if s.next == ("__start__",) + ] + + +def test_migration_preserves_pre_migration_state_sync() -> None: + """Drive 3 invokes under `BinaryOperatorAggregate`, swap the + annotation to `DeltaChannel` on the same sqlite-backed thread, and + verify every settled pre-migration boundary round-trips exactly. + + The override's streaming walk must identify the plain accumulated + list at each pre-migration ancestor as a valid `seed` even though + no `_DeltaSnapshot` was ever written there. + """ + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "mig-sync"}} + + binop = _binop_graph(saver) + _drive(binop, config, "u", 3) + + pre_boundaries = _settled_boundaries(list(binop.get_state_history(config))) + assert len(pre_boundaries) >= 2, "expected multiple settled boundaries" + + delta = _delta_graph(saver) + for cfg, items in pre_boundaries: + snap = delta.get_state(cfg) + assert list(snap.values.get("items", [])) == items, ( + f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: " + f"expected {items}, got {snap.values.get('items', [])}" + ) + + +def test_migration_continued_thread_folds_deltas_on_seed_sync() -> None: + """After migration, driving one more super-step extends the + pre-migration accumulated state via the delta reducer — the seed + plus a single new write. + """ + with SqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "mig-continue-sync"}} + + binop = _binop_graph(saver) + _drive(binop, config, "u", 3) + + pre_history = list(binop.get_state_history(config)) + pre_boundaries = _settled_boundaries(pre_history) + # Latest settled boundary — the leaf pre-migration state. + leaf_cfg, leaf_items = pre_boundaries[0] + assert leaf_items, "expected non-empty pre-migration leaf" + + delta = _delta_graph(saver) + delta.invoke({"items": ["after-migration"]}, leaf_cfg) + new_state = delta.get_state(config).values["items"] + assert new_state[: len(leaf_items)] == leaf_items + assert "after-migration" in new_state + + +async def test_migration_preserves_pre_migration_state_async() -> None: + """Async equivalent of the basic-migration round-trip check on + `AsyncSqliteSaver`.""" + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + config: RunnableConfig = {"configurable": {"thread_id": "mig-async"}} + + binop = _binop_graph(saver) + await _adrive(binop, config, "u", 3) + + pre_history = [s async for s in binop.aget_state_history(config)] + pre_boundaries = _settled_boundaries(pre_history) + assert len(pre_boundaries) >= 2 + + delta = _delta_graph(saver) + for cfg, items in pre_boundaries: + snap = await delta.aget_state(cfg) + assert list(snap.values.get("items", [])) == items, ( + f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}" + ) diff --git a/libs/checkpoint-sqlite/uv.lock b/libs/checkpoint-sqlite/uv.lock index 7c41f741f..46121451a 100644 --- a/libs/checkpoint-sqlite/uv.lock +++ b/libs/checkpoint-sqlite/uv.lock @@ -320,7 +320,7 @@ test = [ [[package]] name = "langgraph-checkpoint-sqlite" -version = "3.0.3" +version = "3.1.0a1" source = { editable = "." } dependencies = [ { name = "aiosqlite" }, diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index 8ddfcf534..7ee81a465 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -1658,7 +1658,7 @@ test = [ [[package]] name = "langgraph-checkpoint-sqlite" -version = "3.0.3" +version = "3.1.0a1" source = { editable = "../checkpoint-sqlite" } dependencies = [ { name = "aiosqlite" }, diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index a794dedc9..4f00624e6 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -464,7 +464,7 @@ test = [ [[package]] name = "langgraph-checkpoint-sqlite" -version = "3.0.3" +version = "3.1.0a1" source = { editable = "../checkpoint-sqlite" } dependencies = [ { name = "aiosqlite" },