mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-25 19:15:11 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8e77945f40 | ||
|
|
728fb748be |
+18
-3
@@ -54,6 +54,12 @@ async def build_delta_chain(
|
|||||||
stored: list[RunnableConfig] = []
|
stored: list[RunnableConfig] = []
|
||||||
parent_cfg: RunnableConfig | None = None
|
parent_cfg: RunnableConfig | None = None
|
||||||
|
|
||||||
|
# Track supersteps since the channel's last snapshot, mirroring Pregel's
|
||||||
|
# `counters_since_delta_snapshot` bookkeeping: every step increments the
|
||||||
|
# counter; a snapshot step resets it (and carries no counter entry, so
|
||||||
|
# the walk locates the seed exactly `supersteps` hops back).
|
||||||
|
supersteps_since_snapshot = 0
|
||||||
|
|
||||||
for step in range(total_steps):
|
for step in range(total_steps):
|
||||||
config: RunnableConfig = {
|
config: RunnableConfig = {
|
||||||
"configurable": {
|
"configurable": {
|
||||||
@@ -68,11 +74,19 @@ async def build_delta_chain(
|
|||||||
|
|
||||||
channel_values: dict[str, Any] = {}
|
channel_values: dict[str, Any] = {}
|
||||||
channel_versions: dict[str, int] = {}
|
channel_versions: dict[str, int] = {}
|
||||||
|
supersteps_since_snapshot += 1
|
||||||
|
counters: dict[str, tuple[int, int]] = {}
|
||||||
if step in snapshot_set:
|
if step in snapshot_set:
|
||||||
channel_values[channel] = _DeltaSnapshot(
|
channel_values[channel] = _DeltaSnapshot(
|
||||||
write_value_fn(step),
|
write_value_fn(step),
|
||||||
)
|
)
|
||||||
channel_versions[channel] = step + 1
|
channel_versions[channel] = step + 1
|
||||||
|
supersteps_since_snapshot = 0
|
||||||
|
else:
|
||||||
|
counters[channel] = (
|
||||||
|
supersteps_since_snapshot,
|
||||||
|
supersteps_since_snapshot,
|
||||||
|
)
|
||||||
|
|
||||||
cp = Checkpoint(
|
cp = Checkpoint(
|
||||||
v=1,
|
v=1,
|
||||||
@@ -84,9 +98,10 @@ async def build_delta_chain(
|
|||||||
updated_channels=None,
|
updated_channels=None,
|
||||||
)
|
)
|
||||||
new_versions = dict(channel_versions)
|
new_versions = dict(channel_versions)
|
||||||
parent_cfg = await saver.aput(
|
md = generate_metadata(step=step)
|
||||||
config, cp, generate_metadata(step=step), new_versions
|
if counters:
|
||||||
)
|
md["counters_since_delta_snapshot"] = counters
|
||||||
|
parent_cfg = await saver.aput(config, cp, md, new_versions)
|
||||||
stored.append(parent_cfg)
|
stored.append(parent_cfg)
|
||||||
|
|
||||||
# Write a pending write for non-snapshot steps so the walk has
|
# Write a pending write for non-snapshot steps so the walk has
|
||||||
|
|||||||
+92
-2
@@ -87,6 +87,10 @@ async def test_history_multi_channel(
|
|||||||
|
|
||||||
from langgraph.checkpoint.conformance.test_utils import generate_metadata
|
from langgraph.checkpoint.conformance.test_utils import generate_metadata
|
||||||
|
|
||||||
|
# Per-channel supersteps since last snapshot ("a" snapshots at step 1,
|
||||||
|
# "b" at step 3). Drives the per-channel walk depth from the head.
|
||||||
|
s_a = 0
|
||||||
|
s_b = 0
|
||||||
for step in range(5):
|
for step in range(5):
|
||||||
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
||||||
if parent_cfg:
|
if parent_cfg:
|
||||||
@@ -95,12 +99,21 @@ async def test_history_multi_channel(
|
|||||||
]
|
]
|
||||||
cv: dict = {}
|
cv: dict = {}
|
||||||
cvs: dict = {}
|
cvs: dict = {}
|
||||||
|
s_a += 1
|
||||||
|
s_b += 1
|
||||||
|
counters: dict = {}
|
||||||
if step == 1:
|
if step == 1:
|
||||||
cv["a"] = _DeltaSnapshot("snap_a")
|
cv["a"] = _DeltaSnapshot("snap_a")
|
||||||
cvs["a"] = step + 1
|
cvs["a"] = step + 1
|
||||||
|
s_a = 0
|
||||||
|
else:
|
||||||
|
counters["a"] = (s_a, s_a)
|
||||||
if step == 3:
|
if step == 3:
|
||||||
cv["b"] = _DeltaSnapshot("snap_b")
|
cv["b"] = _DeltaSnapshot("snap_b")
|
||||||
cvs["b"] = step + 1
|
cvs["b"] = step + 1
|
||||||
|
s_b = 0
|
||||||
|
else:
|
||||||
|
counters["b"] = (s_b, s_b)
|
||||||
cp = Checkpoint(
|
cp = Checkpoint(
|
||||||
v=1,
|
v=1,
|
||||||
id=str(uuid6(clock_seq=-1)),
|
id=str(uuid6(clock_seq=-1)),
|
||||||
@@ -110,7 +123,10 @@ async def test_history_multi_channel(
|
|||||||
versions_seen={},
|
versions_seen={},
|
||||||
updated_channels=None,
|
updated_channels=None,
|
||||||
)
|
)
|
||||||
parent_cfg = await saver.aput(config, cp, generate_metadata(step=step), cvs)
|
md = generate_metadata(step=step)
|
||||||
|
if counters:
|
||||||
|
md["counters_since_delta_snapshot"] = counters
|
||||||
|
parent_cfg = await saver.aput(config, cp, md, cvs)
|
||||||
configs.append(parent_cfg)
|
configs.append(parent_cfg)
|
||||||
await saver.aput_writes(parent_cfg, [("a", step), ("b", step)], str(uuid4()))
|
await saver.aput_writes(parent_cfg, [("a", step), ("b", step)], str(uuid4()))
|
||||||
|
|
||||||
@@ -148,7 +164,12 @@ async def test_history_walk_to_root_no_seed(
|
|||||||
)
|
)
|
||||||
head = configs[-1]
|
head = configs[-1]
|
||||||
result = await saver.aget_delta_channel_history(config=head, channels=["ch"])
|
result = await saver.aget_delta_channel_history(config=head, channels=["ch"])
|
||||||
|
# The head's supersteps counter (4) runs one hop past the real root —
|
||||||
|
# the snapshot checkpoint was never persisted (implicit empty baseline).
|
||||||
|
# No seed, and the full chain (steps 0..2) replays from empty.
|
||||||
assert "seed" not in result["ch"], f"Expected no seed, got {result['ch']}"
|
assert "seed" not in result["ch"], f"Expected no seed, got {result['ch']}"
|
||||||
|
values = [w[2] for w in result["ch"]["writes"]]
|
||||||
|
assert values == [0, 1, 2], f"Expected full-chain writes [0,1,2], got {values}"
|
||||||
|
|
||||||
|
|
||||||
async def test_history_migration_plain_value_as_seed(
|
async def test_history_migration_plain_value_as_seed(
|
||||||
@@ -170,6 +191,10 @@ async def test_history_migration_plain_value_as_seed(
|
|||||||
configs: list = []
|
configs: list = []
|
||||||
parent_cfg = None
|
parent_cfg = None
|
||||||
|
|
||||||
|
# The channel was a non-delta channel through step 1 (plain value, no
|
||||||
|
# delta counter), then migrated to DeltaChannel at step 2. Supersteps
|
||||||
|
# count from the migration boundary: step 2 -> 1, step 3 -> 2, so the
|
||||||
|
# head (step 3) walks 2 hops back to the plain-value seed at step 1.
|
||||||
for step in range(4):
|
for step in range(4):
|
||||||
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
||||||
if parent_cfg:
|
if parent_cfg:
|
||||||
@@ -182,6 +207,9 @@ async def test_history_migration_plain_value_as_seed(
|
|||||||
if step == 1:
|
if step == 1:
|
||||||
cv["ch"] = [10, 20, 30]
|
cv["ch"] = [10, 20, 30]
|
||||||
cvs["ch"] = step + 1
|
cvs["ch"] = step + 1
|
||||||
|
md = generate_metadata(step=step)
|
||||||
|
if step >= 2:
|
||||||
|
md["counters_since_delta_snapshot"] = {"ch": (step - 1, step - 1)}
|
||||||
cp = Checkpoint(
|
cp = Checkpoint(
|
||||||
v=1,
|
v=1,
|
||||||
id=str(uuid6(clock_seq=-1)),
|
id=str(uuid6(clock_seq=-1)),
|
||||||
@@ -191,7 +219,7 @@ async def test_history_migration_plain_value_as_seed(
|
|||||||
versions_seen={},
|
versions_seen={},
|
||||||
updated_channels=None,
|
updated_channels=None,
|
||||||
)
|
)
|
||||||
parent_cfg = await saver.aput(config, cp, generate_metadata(step=step), cvs)
|
parent_cfg = await saver.aput(config, cp, md, cvs)
|
||||||
configs.append(parent_cfg)
|
configs.append(parent_cfg)
|
||||||
if step != 1:
|
if step != 1:
|
||||||
await saver.aput_writes(parent_cfg, [("ch", step)], str(uuid4()))
|
await saver.aput_writes(parent_cfg, [("ch", step)], str(uuid4()))
|
||||||
@@ -208,6 +236,67 @@ async def test_history_migration_plain_value_as_seed(
|
|||||||
assert values == [2], f"Expected [2], got {values}"
|
assert values == [2], f"Expected [2], got {values}"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_history_migration_skips_seed_checkpoint_writes(
|
||||||
|
saver: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""A migrated plain-value seed already incorporates its own checkpoint's
|
||||||
|
writes, so those writes must NOT be re-emitted in the history.
|
||||||
|
|
||||||
|
This is the discriminating case for non-additive reducers (e.g. an
|
||||||
|
even-only filter): re-applying the seed checkpoint's writes on top of
|
||||||
|
the seed would corrupt the reconstructed value. The saver contract is
|
||||||
|
to skip them — only writes strictly after the seed are returned.
|
||||||
|
"""
|
||||||
|
from langgraph.checkpoint.base import Checkpoint
|
||||||
|
from langgraph.checkpoint.base.id import uuid6
|
||||||
|
|
||||||
|
from langgraph.checkpoint.conformance.test_utils import generate_metadata
|
||||||
|
|
||||||
|
tid = str(uuid4())
|
||||||
|
configs: list = []
|
||||||
|
parent_cfg = None
|
||||||
|
|
||||||
|
for step in range(4):
|
||||||
|
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
||||||
|
if parent_cfg:
|
||||||
|
config["configurable"]["checkpoint_id"] = parent_cfg["configurable"][
|
||||||
|
"checkpoint_id"
|
||||||
|
]
|
||||||
|
cv: dict = {}
|
||||||
|
cvs: dict = {}
|
||||||
|
# Step 1: pre-delta plain accumulated value AND its own writes.
|
||||||
|
if step == 1:
|
||||||
|
cv["ch"] = [2, 4]
|
||||||
|
cvs["ch"] = step + 1
|
||||||
|
md = generate_metadata(step=step)
|
||||||
|
if step >= 2:
|
||||||
|
md["counters_since_delta_snapshot"] = {"ch": (step - 1, step - 1)}
|
||||||
|
cp = Checkpoint(
|
||||||
|
v=1,
|
||||||
|
id=str(uuid6(clock_seq=-1)),
|
||||||
|
ts="",
|
||||||
|
channel_values=cv,
|
||||||
|
channel_versions=cvs,
|
||||||
|
versions_seen={},
|
||||||
|
updated_channels=None,
|
||||||
|
)
|
||||||
|
parent_cfg = await saver.aput(config, cp, md, cvs)
|
||||||
|
configs.append(parent_cfg)
|
||||||
|
# The seed checkpoint (step 1) carries writes that the plain value
|
||||||
|
# already subsumes; later steps carry post-migration delta writes.
|
||||||
|
write_value = 99 if step == 1 else step * 10
|
||||||
|
await saver.aput_writes(parent_cfg, [("ch", write_value)], str(uuid4()))
|
||||||
|
|
||||||
|
head = configs[-1]
|
||||||
|
result = await saver.aget_delta_channel_history(config=head, channels=["ch"])
|
||||||
|
assert result["ch"].get("seed") == [2, 4], (
|
||||||
|
f"Expected plain seed [2, 4], got {result['ch'].get('seed')}"
|
||||||
|
)
|
||||||
|
values = [w[2] for w in result["ch"]["writes"]]
|
||||||
|
# Seed checkpoint's own write (99) is skipped; step 2's write (20) kept.
|
||||||
|
assert values == [20], f"Expected only post-seed writes [20], got {values}"
|
||||||
|
|
||||||
|
|
||||||
ALL_DELTA_CHANNEL_HISTORY_TESTS = [
|
ALL_DELTA_CHANNEL_HISTORY_TESTS = [
|
||||||
test_history_returns_writes_oldest_first,
|
test_history_returns_writes_oldest_first,
|
||||||
test_history_seed_is_nearest_snapshot,
|
test_history_seed_is_nearest_snapshot,
|
||||||
@@ -216,6 +305,7 @@ ALL_DELTA_CHANNEL_HISTORY_TESTS = [
|
|||||||
test_history_empty_channels_returns_empty,
|
test_history_empty_channels_returns_empty,
|
||||||
test_history_walk_to_root_no_seed,
|
test_history_walk_to_root_no_seed,
|
||||||
test_history_migration_plain_value_as_seed,
|
test_history_migration_plain_value_as_seed,
|
||||||
|
test_history_migration_skips_seed_checkpoint_writes,
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
|
|||||||
CheckpointMetadata,
|
CheckpointMetadata,
|
||||||
CheckpointTuple,
|
CheckpointTuple,
|
||||||
DeltaChannelHistory,
|
DeltaChannelHistory,
|
||||||
|
_parse_supersteps_since_last_snapshot_by_channel,
|
||||||
get_checkpoint_id,
|
get_checkpoint_id,
|
||||||
get_serializable_checkpoint_metadata,
|
get_serializable_checkpoint_metadata,
|
||||||
)
|
)
|
||||||
@@ -28,9 +29,11 @@ from langgraph.checkpoint.postgres import _internal
|
|||||||
from langgraph.checkpoint.postgres.base import (
|
from langgraph.checkpoint.postgres.base import (
|
||||||
_DELTA_PAGE_SIZE,
|
_DELTA_PAGE_SIZE,
|
||||||
BasePostgresSaver,
|
BasePostgresSaver,
|
||||||
_build_delta_stage1_sql,
|
_advance_shared_chain,
|
||||||
_build_delta_stage2_sql,
|
_build_delta_fetch_sql,
|
||||||
|
_build_delta_walk_sql,
|
||||||
_DeltaStage2Row,
|
_DeltaStage2Row,
|
||||||
|
_ingest_walk_page,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||||
|
|
||||||
@@ -446,107 +449,137 @@ class PostgresSaver(BasePostgresSaver):
|
|||||||
) -> Mapping[str, DeltaChannelHistory]:
|
) -> Mapping[str, DeltaChannelHistory]:
|
||||||
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
|
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
|
||||||
|
|
||||||
Two-stage query, both stages cover ALL requested channels:
|
Reconstructs each delta channel's state for the target checkpoint by
|
||||||
|
walking the parent chain `supersteps` hops (read from the target's
|
||||||
|
`counters_since_delta_snapshot` metadata) to its seed snapshot, then
|
||||||
|
collecting the writes between the seed and the target.
|
||||||
|
|
||||||
* Stage 1 (paged): dynamic SELECT over `checkpoints` with K parallel
|
Two passes (see the `# Multi-channel two-pass` comment in `base.py`):
|
||||||
JSONB key lookups (one column pair per channel) — no subquery, no
|
|
||||||
aggregation. Pages newest-first by `checkpoint_id` with a cursor;
|
|
||||||
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
|
* WALK (paged): dynamic SELECT over `checkpoints` with K parallel
|
||||||
`checkpoint_writes` filtered to that channel's specific
|
JSONB lookups for `channel_versions[ch]` (the seed blob version),
|
||||||
`chain_cids`, plus one branch per channel that has a seed reading
|
following `parent_checkpoint_id` newest-first. Page size is
|
||||||
`checkpoint_blobs` for that channel + version. Avoids the
|
`_DELTA_PAGE_SIZE`; stops once the shared chain reaches the deepest
|
||||||
over-fetch of a single `channel = ANY(channels)` filter when
|
requested `supersteps`, the root is reached, or the chain is
|
||||||
channels have different chain depths.
|
exhausted.
|
||||||
|
|
||||||
|
* FETCH (per-channel UNION ALL): one branch per channel reading
|
||||||
|
`checkpoint_writes` for its `chain_cids`, plus one branch per
|
||||||
|
channel that located a seed reading `checkpoint_blobs` at the
|
||||||
|
seed's version.
|
||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
channels = list(channels)
|
channels = list(channels)
|
||||||
thread_id = config["configurable"]["thread_id"]
|
thread_id = config["configurable"]["thread_id"]
|
||||||
|
if not thread_id:
|
||||||
|
raise ValueError("empty thread ID")
|
||||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
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"]
|
|
||||||
|
|
||||||
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
|
# Resolve the target checkpoint id + its metadata (for supersteps).
|
||||||
# Python after each page. Stops as soon as every channel has its seed.
|
checkpoint_id = get_checkpoint_id(config)
|
||||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
with self._cursor() as cur:
|
||||||
|
if checkpoint_id is None:
|
||||||
|
cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = %s AND checkpoint_ns = %s "
|
||||||
|
"ORDER BY checkpoint_id DESC LIMIT 1",
|
||||||
|
(thread_id, checkpoint_ns),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = %s AND checkpoint_ns = %s "
|
||||||
|
"AND checkpoint_id = %s",
|
||||||
|
(thread_id, checkpoint_ns, checkpoint_id),
|
||||||
|
)
|
||||||
|
target_row = cur.fetchone()
|
||||||
|
if target_row is None:
|
||||||
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
target_id = cast(str, target_row["checkpoint_id"])
|
||||||
|
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
target_row["metadata"] or {}, channels
|
||||||
|
)
|
||||||
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
|
||||||
|
# WALK: page the parent chain, bounded by the deepest supersteps.
|
||||||
|
walk_sql = _build_delta_walk_sql(channels)
|
||||||
parent_of: dict[str, str | None] = {}
|
parent_of: dict[str, str | None] = {}
|
||||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||||
hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
shared_cpid_chain: list[str] = []
|
||||||
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
has_reached_root = False
|
||||||
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
|
||||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
|
||||||
seeded: set[str] = set()
|
|
||||||
cursor: str | None = None
|
cursor: str | None = None
|
||||||
|
|
||||||
with self._cursor() as cur:
|
while max_supersteps > 0:
|
||||||
while True:
|
walk_params: list[Any] = [
|
||||||
stage1_params: list[Any] = []
|
*channels,
|
||||||
for ch in channels:
|
thread_id,
|
||||||
stage1_params.extend([ch, ch])
|
checkpoint_ns,
|
||||||
stage1_params.extend(
|
cursor,
|
||||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
cursor,
|
||||||
)
|
_DELTA_PAGE_SIZE,
|
||||||
cur.execute(stage1_sql, stage1_params)
|
]
|
||||||
|
with self._cursor() as cur:
|
||||||
|
cur.execute(walk_sql, walk_params)
|
||||||
page = cur.fetchall()
|
page = cur.fetchall()
|
||||||
if not page:
|
if not page:
|
||||||
break
|
break
|
||||||
oldest = self._ingest_stage1_page(
|
oldest = _ingest_walk_page(
|
||||||
cast("list[Mapping[str, Any]]", page),
|
cast("list[Mapping[str, Any]]", page),
|
||||||
channels,
|
channels,
|
||||||
parent_of,
|
parent_of,
|
||||||
ver_by_i_by_cid,
|
ver_by_i_by_cid,
|
||||||
hs_by_i_by_cid,
|
)
|
||||||
)
|
has_reached_root = _advance_shared_chain(
|
||||||
self._try_advance_walks(
|
target_id, parent_of, shared_cpid_chain, max_supersteps
|
||||||
checkpoint_id,
|
)
|
||||||
channels,
|
if (
|
||||||
parent_of,
|
has_reached_root
|
||||||
ver_by_i_by_cid,
|
or len(shared_cpid_chain) >= max_supersteps
|
||||||
hs_by_i_by_cid,
|
or len(page) < _DELTA_PAGE_SIZE
|
||||||
chain_by_ch,
|
):
|
||||||
seed_ver_by_ch,
|
break
|
||||||
walk_cursor_by_ch,
|
cursor = oldest
|
||||||
seeded,
|
|
||||||
)
|
|
||||||
# Stop if every channel is seeded, or the page was short
|
|
||||||
# (chain exhausted — no more rows to fetch).
|
|
||||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
|
||||||
break
|
|
||||||
cursor = oldest
|
|
||||||
|
|
||||||
# Stage 2: per-channel UNION ALL — one writes branch per channel
|
chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch = (
|
||||||
# with non-empty chain, plus one blob branch per seeded channel.
|
self._resolve_delta_chains(
|
||||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
channels,
|
||||||
|
supersteps_by_ch,
|
||||||
|
shared_cpid_chain,
|
||||||
|
ver_by_i_by_cid,
|
||||||
|
has_reached_root,
|
||||||
|
thread_id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# FETCH: per-channel UNION ALL — writes branch per chained channel,
|
||||||
|
# blob branch per seeded channel.
|
||||||
|
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
|
||||||
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
||||||
stage2_sql = _build_delta_stage2_sql(
|
fetch_sql = _build_delta_fetch_sql(
|
||||||
channels_with_chain=channels_with_chain,
|
channels_with_chain=channels_with_chain,
|
||||||
channels_with_seed=channels_with_seed,
|
channels_with_seed=channels_with_seed,
|
||||||
)
|
)
|
||||||
|
if fetch_sql:
|
||||||
if stage2_sql:
|
fetch_params: list[Any] = []
|
||||||
stage2_params: list[Any] = []
|
|
||||||
for ch in channels_with_chain:
|
for ch in channels_with_chain:
|
||||||
stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]])
|
fetch_params.extend(
|
||||||
|
[thread_id, checkpoint_ns, ch, chained_cpid_by_ch[ch]]
|
||||||
|
)
|
||||||
for ch in channels_with_seed:
|
for ch in channels_with_seed:
|
||||||
stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
|
fetch_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
|
||||||
with self._cursor() as cur:
|
with self._cursor() as cur:
|
||||||
cur.execute(stage2_sql, stage2_params)
|
cur.execute(fetch_sql, fetch_params)
|
||||||
stage2_rows = cur.fetchall()
|
fetch_rows = cur.fetchall()
|
||||||
else:
|
else:
|
||||||
stage2_rows = []
|
fetch_rows = []
|
||||||
|
|
||||||
return self._build_delta_channels_writes_history(
|
return self._assemble_delta_history(
|
||||||
channels=channels,
|
channels=channels,
|
||||||
chain_by_ch=chain_by_ch,
|
chained_cpid_by_ch=chained_cpid_by_ch,
|
||||||
|
seed_cpid_by_ch=seed_cpid_by_ch,
|
||||||
seed_ver_by_ch=seed_ver_by_ch,
|
seed_ver_by_ch=seed_ver_by_ch,
|
||||||
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
|
fetch_rows=cast("list[_DeltaStage2Row]", fetch_rows),
|
||||||
)
|
)
|
||||||
|
|
||||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
|
|||||||
CheckpointMetadata,
|
CheckpointMetadata,
|
||||||
CheckpointTuple,
|
CheckpointTuple,
|
||||||
DeltaChannelHistory,
|
DeltaChannelHistory,
|
||||||
|
_parse_supersteps_since_last_snapshot_by_channel,
|
||||||
get_checkpoint_id,
|
get_checkpoint_id,
|
||||||
get_serializable_checkpoint_metadata,
|
get_serializable_checkpoint_metadata,
|
||||||
)
|
)
|
||||||
@@ -28,9 +29,11 @@ from langgraph.checkpoint.postgres import _ainternal
|
|||||||
from langgraph.checkpoint.postgres.base import (
|
from langgraph.checkpoint.postgres.base import (
|
||||||
_DELTA_PAGE_SIZE,
|
_DELTA_PAGE_SIZE,
|
||||||
BasePostgresSaver,
|
BasePostgresSaver,
|
||||||
_build_delta_stage1_sql,
|
_advance_shared_chain,
|
||||||
_build_delta_stage2_sql,
|
_build_delta_fetch_sql,
|
||||||
|
_build_delta_walk_sql,
|
||||||
_DeltaStage2Row,
|
_DeltaStage2Row,
|
||||||
|
_ingest_walk_page,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||||
|
|
||||||
@@ -408,89 +411,122 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|||||||
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
|
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
|
||||||
|
|
||||||
See `PostgresSaver.get_delta_channel_history` for design notes; this is
|
See `PostgresSaver.get_delta_channel_history` for design notes; this is
|
||||||
the async equivalent with internal stage-1 paging and per-channel
|
the async equivalent: paged supersteps-bounded WALK + per-channel
|
||||||
UNION ALL stage-2.
|
UNION ALL FETCH.
|
||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
channels = list(channels)
|
channels = list(channels)
|
||||||
thread_id = config["configurable"]["thread_id"]
|
thread_id = config["configurable"]["thread_id"]
|
||||||
|
if not thread_id:
|
||||||
|
raise ValueError("empty thread ID")
|
||||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
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"]
|
|
||||||
|
|
||||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
# Resolve the target checkpoint id + its metadata (for supersteps).
|
||||||
|
checkpoint_id = get_checkpoint_id(config)
|
||||||
|
async with self._cursor() as cur:
|
||||||
|
if checkpoint_id is None:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = %s AND checkpoint_ns = %s "
|
||||||
|
"ORDER BY checkpoint_id DESC LIMIT 1",
|
||||||
|
(thread_id, checkpoint_ns),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = %s AND checkpoint_ns = %s "
|
||||||
|
"AND checkpoint_id = %s",
|
||||||
|
(thread_id, checkpoint_ns, checkpoint_id),
|
||||||
|
)
|
||||||
|
target_row = await cur.fetchone()
|
||||||
|
if target_row is None:
|
||||||
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
target_id = cast(str, target_row["checkpoint_id"])
|
||||||
|
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
target_row["metadata"] or {}, channels
|
||||||
|
)
|
||||||
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
|
||||||
|
# WALK: page the parent chain, bounded by the deepest supersteps.
|
||||||
|
walk_sql = _build_delta_walk_sql(channels)
|
||||||
parent_of: dict[str, str | None] = {}
|
parent_of: dict[str, str | None] = {}
|
||||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||||
hs_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
shared_cpid_chain: list[str] = []
|
||||||
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
has_reached_root = False
|
||||||
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
|
||||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
|
||||||
seeded: set[str] = set()
|
|
||||||
cursor: str | None = None
|
cursor: str | None = None
|
||||||
|
|
||||||
async with self._cursor() as cur:
|
while max_supersteps > 0:
|
||||||
while True:
|
walk_params: list[Any] = [
|
||||||
stage1_params: list[Any] = []
|
*channels,
|
||||||
for ch in channels:
|
thread_id,
|
||||||
stage1_params.extend([ch, ch])
|
checkpoint_ns,
|
||||||
stage1_params.extend(
|
cursor,
|
||||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
cursor,
|
||||||
)
|
_DELTA_PAGE_SIZE,
|
||||||
await cur.execute(stage1_sql, stage1_params)
|
]
|
||||||
|
async with self._cursor() as cur:
|
||||||
|
await cur.execute(walk_sql, walk_params)
|
||||||
page = await cur.fetchall()
|
page = await cur.fetchall()
|
||||||
if not page:
|
if not page:
|
||||||
break
|
break
|
||||||
oldest = self._ingest_stage1_page(
|
oldest = _ingest_walk_page(
|
||||||
cast("list[Mapping[str, Any]]", page),
|
cast("list[Mapping[str, Any]]", page),
|
||||||
channels,
|
channels,
|
||||||
parent_of,
|
parent_of,
|
||||||
ver_by_i_by_cid,
|
ver_by_i_by_cid,
|
||||||
hs_by_i_by_cid,
|
)
|
||||||
)
|
has_reached_root = _advance_shared_chain(
|
||||||
self._try_advance_walks(
|
target_id, parent_of, shared_cpid_chain, max_supersteps
|
||||||
checkpoint_id,
|
)
|
||||||
channels,
|
if (
|
||||||
parent_of,
|
has_reached_root
|
||||||
ver_by_i_by_cid,
|
or len(shared_cpid_chain) >= max_supersteps
|
||||||
hs_by_i_by_cid,
|
or len(page) < _DELTA_PAGE_SIZE
|
||||||
chain_by_ch,
|
):
|
||||||
seed_ver_by_ch,
|
break
|
||||||
walk_cursor_by_ch,
|
cursor = oldest
|
||||||
seeded,
|
|
||||||
)
|
|
||||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
|
||||||
break
|
|
||||||
cursor = oldest
|
|
||||||
|
|
||||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch = (
|
||||||
|
self._resolve_delta_chains(
|
||||||
|
channels,
|
||||||
|
supersteps_by_ch,
|
||||||
|
shared_cpid_chain,
|
||||||
|
ver_by_i_by_cid,
|
||||||
|
has_reached_root,
|
||||||
|
thread_id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# FETCH: per-channel UNION ALL — writes branch per chained channel,
|
||||||
|
# blob branch per seeded channel.
|
||||||
|
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
|
||||||
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
||||||
stage2_sql = _build_delta_stage2_sql(
|
fetch_sql = _build_delta_fetch_sql(
|
||||||
channels_with_chain=channels_with_chain,
|
channels_with_chain=channels_with_chain,
|
||||||
channels_with_seed=channels_with_seed,
|
channels_with_seed=channels_with_seed,
|
||||||
)
|
)
|
||||||
|
|
||||||
if stage2_sql:
|
if fetch_sql:
|
||||||
stage2_params: list[Any] = []
|
fetch_params: list[Any] = []
|
||||||
for ch in channels_with_chain:
|
for ch in channels_with_chain:
|
||||||
stage2_params.extend([thread_id, checkpoint_ns, ch, chain_by_ch[ch]])
|
fetch_params.extend(
|
||||||
|
[thread_id, checkpoint_ns, ch, chained_cpid_by_ch[ch]]
|
||||||
|
)
|
||||||
for ch in channels_with_seed:
|
for ch in channels_with_seed:
|
||||||
stage2_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
|
fetch_params.extend([thread_id, checkpoint_ns, ch, seed_ver_by_ch[ch]])
|
||||||
async with self._cursor() as cur:
|
async with self._cursor() as cur:
|
||||||
await cur.execute(stage2_sql, stage2_params)
|
await cur.execute(fetch_sql, fetch_params)
|
||||||
stage2_rows = await cur.fetchall()
|
fetch_rows = await cur.fetchall()
|
||||||
else:
|
else:
|
||||||
stage2_rows = []
|
fetch_rows = []
|
||||||
|
|
||||||
return self._build_delta_channels_writes_history(
|
return self._assemble_delta_history(
|
||||||
channels=channels,
|
channels=channels,
|
||||||
chain_by_ch=chain_by_ch,
|
chained_cpid_by_ch=chained_cpid_by_ch,
|
||||||
|
seed_cpid_by_ch=seed_cpid_by_ch,
|
||||||
seed_ver_by_ch=seed_ver_by_ch,
|
seed_ver_by_ch=seed_ver_by_ch,
|
||||||
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
|
fetch_rows=cast("list[_DeltaStage2Row]", fetch_rows),
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
import random
|
import random
|
||||||
import warnings
|
import warnings
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Mapping, Sequence
|
||||||
@@ -15,10 +16,12 @@ from langgraph.checkpoint.base import (
|
|||||||
PendingWrite,
|
PendingWrite,
|
||||||
get_checkpoint_id,
|
get_checkpoint_id,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.serde.types import TASKS
|
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
|
||||||
from psycopg.types.json import Jsonb
|
from psycopg.types.json import Jsonb
|
||||||
|
|
||||||
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Page size for the paged WALK scan in `get_delta_channel_history`. Internal
|
||||||
# constant — exposing this as a kwarg is left as a follow-up.
|
# constant — exposing this as a kwarg is left as a follow-up.
|
||||||
_DELTA_PAGE_SIZE = 1024
|
_DELTA_PAGE_SIZE = 1024
|
||||||
|
|
||||||
@@ -160,7 +163,7 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
|||||||
|
|
||||||
|
|
||||||
class _DeltaStage2Row(TypedDict, total=False):
|
class _DeltaStage2Row(TypedDict, total=False):
|
||||||
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
|
"""One row from `_build_delta_fetch_sql` (a UNION ALL of writes and blobs)."""
|
||||||
|
|
||||||
_kind: str # "w" or "b"
|
_kind: str # "w" or "b"
|
||||||
checkpoint_id: str | None # "w" rows only
|
checkpoint_id: str | None # "w" rows only
|
||||||
@@ -172,42 +175,58 @@ class _DeltaStage2Row(TypedDict, total=False):
|
|||||||
version: str | None # "b" rows only
|
version: str | None # "b" rows only
|
||||||
|
|
||||||
|
|
||||||
# Multi-channel two-stage DeltaChannel reconstruction.
|
# Multi-channel two-pass DeltaChannel reconstruction.
|
||||||
#
|
#
|
||||||
# Stage 1 scans checkpoint metadata (no blob bytes) and emits one row per
|
# A `DeltaChannel` does not store its full value at every checkpoint — it
|
||||||
# checkpoint with K parallel JSONB key lookups (one column pair per
|
# stores periodic full-value *snapshots* and accumulates intermediate
|
||||||
# requested delta channel: ver_i / hs_i). No subqueries, no aggregation.
|
# *writes* between snapshots. To rebuild a channel's value at a target
|
||||||
# Python walks the parent chain once across all channels.
|
# checkpoint we need:
|
||||||
#
|
#
|
||||||
# Stage 2 fetches all writes and the seed blobs for ALL channels in a
|
# - the **seed** — the most recent snapshot at-or-before the target
|
||||||
# single roundtrip via `channel = ANY(%s)` and chain/seed-version
|
# (a single blob row in `checkpoint_blobs`); and
|
||||||
# filtering.
|
# - the **chain writes** — every write for this channel committed
|
||||||
|
# between that snapshot and the target, in order
|
||||||
|
# (rows in `checkpoint_writes`).
|
||||||
#
|
#
|
||||||
# Empirical comparison vs an alternative "ship full channel_versions /
|
# Two passes, in order:
|
||||||
# channel_values JSONB and let Python pick" form (1000 checkpoints,
|
|
||||||
# 8 total channels in graph, 3 delta channels requested):
|
|
||||||
#
|
#
|
||||||
# Postgres execution: A=0.24ms vs B=0.38ms (both negligible)
|
# 1. WALK — scan checkpoint metadata only (no blob bytes). For each
|
||||||
# End-to-end latency: A=6.83ms vs B=2.28ms (B is 3.0x faster)
|
# requested channel we read `channel_versions[ch]` (the seed's blob
|
||||||
# Wire payload: A=836KB vs B=330KB (61% smaller)
|
# version pointer). Then follow `parent_checkpoint_id` from the
|
||||||
# Buffer hits: identical (167 blocks)
|
# target backwards in pages of `_DELTA_PAGE_SIZE` rows.
|
||||||
#
|
#
|
||||||
# B (this dynamic-columns design) wins because it avoids JSONB
|
# Walk depth is driven by the *supersteps since last snapshot*
|
||||||
# serialization on the wire and JSONB-to-dict deserialization in
|
# counter — `metadata.counters_since_delta_snapshot[ch][1]` — read
|
||||||
# psycopg. Even at K=8 (8 delta channels = 16 dynamic columns), B
|
# from the target checkpoint. A channel's seed snapshot sits exactly
|
||||||
# still beats A end-to-end (4.2ms vs 6.8ms).
|
# `supersteps` hops back along the parent chain; walking by the
|
||||||
|
# counter (rather than scanning `channel_values` for the snapshot
|
||||||
|
# marker) is the only reliable way to locate seeds that aren't a
|
||||||
|
# `_DeltaSnapshot` sentinel — e.g. legacy plain-value blobs left by a
|
||||||
|
# thread that migrated from a non-delta channel, which `put` stores
|
||||||
|
# out of the inline `channel_values` map.
|
||||||
|
#
|
||||||
|
# 2. FETCH — given each channel's chained checkpoint ids and seed
|
||||||
|
# version from WALK, pull only the rows we need: writes for those
|
||||||
|
# exact checkpoint_ids and the seed blob at that exact version. One
|
||||||
|
# roundtrip, per-channel UNION ALL — no over-fetch.
|
||||||
|
#
|
||||||
|
# Walking the *parent chain* (not `list(before=...)`) matters: forked
|
||||||
|
# threads have multiple branches, and only on-path ancestors contribute.
|
||||||
|
|
||||||
|
|
||||||
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
def _build_delta_walk_sql(channels: Sequence[str]) -> str:
|
||||||
"""Build stage 1 SQL with 2K parallel JSONB key lookups.
|
"""Build the paged WALK SQL — scans checkpoint metadata only.
|
||||||
|
|
||||||
For channels=["messages", "files"] (with `paged=True`) the result is::
|
Emits one row per checkpoint with K parallel JSONB key lookups (one
|
||||||
|
`ver_i` column per requested channel — the channel's blob version, the
|
||||||
|
pointer we dereference in FETCH if this checkpoint is the seed). No
|
||||||
|
blob bytes; the result set fits a paged `LIMIT` cleanly.
|
||||||
|
|
||||||
|
For channels=["messages", "files"] the result is::
|
||||||
|
|
||||||
SELECT checkpoint_id, parent_checkpoint_id,
|
SELECT checkpoint_id, parent_checkpoint_id,
|
||||||
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
||||||
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0,
|
checkpoint -> 'channel_versions' ->> %s AS ver_1
|
||||||
checkpoint -> 'channel_versions' ->> %s AS ver_1,
|
|
||||||
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1
|
|
||||||
FROM checkpoints
|
FROM checkpoints
|
||||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||||
AND (%s::text IS NULL OR checkpoint_id < %s)
|
AND (%s::text IS NULL OR checkpoint_id < %s)
|
||||||
@@ -215,45 +234,38 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
|||||||
LIMIT %s
|
LIMIT %s
|
||||||
|
|
||||||
Channel names are passed as `%s` parameters (safe from SQL injection).
|
Channel names are passed as `%s` parameters (safe from SQL injection).
|
||||||
Only the column aliases `ver_i` / `hs_i` are interpolated into the
|
Only the column aliases `ver_i` are interpolated into the SQL string
|
||||||
SQL string (i is bounded by len(channels) and uses safe identifiers).
|
(i is bounded by len(channels) and uses safe identifiers).
|
||||||
|
|
||||||
Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ...,
|
Caller must extend params with `[ch_0, ch_1, ..., thread_id, ns,
|
||||||
thread_id, ns, cursor, cursor, page_size]` when `paged=True`.
|
cursor, cursor, page_size]`. The `cursor` is the smallest
|
||||||
|
`checkpoint_id` from the previous page (or `None` on the first page);
|
||||||
When `paged=False`, the WHERE has no cursor predicate and there's no
|
`(%s::text IS NULL OR ...)` makes the first-page `WHERE` a no-op.
|
||||||
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
|
||||||
"""
|
"""
|
||||||
cols = []
|
cols = [
|
||||||
for i in range(len(channels)):
|
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}"
|
||||||
cols.append(
|
for i in range(len(channels))
|
||||||
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, "
|
]
|
||||||
f"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}"
|
return (
|
||||||
)
|
|
||||||
sql = (
|
|
||||||
"SELECT checkpoint_id, parent_checkpoint_id, "
|
"SELECT checkpoint_id, parent_checkpoint_id, "
|
||||||
+ ", ".join(cols)
|
+ ", ".join(cols)
|
||||||
+ " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
|
+ " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = %s"
|
||||||
|
" AND (%s::text IS NULL OR checkpoint_id < %s)"
|
||||||
|
" ORDER BY checkpoint_id DESC LIMIT %s"
|
||||||
)
|
)
|
||||||
if paged:
|
|
||||||
sql += (
|
|
||||||
" AND (%s::text IS NULL OR checkpoint_id < %s)"
|
|
||||||
" ORDER BY checkpoint_id DESC LIMIT %s"
|
|
||||||
)
|
|
||||||
return sql
|
|
||||||
|
|
||||||
|
|
||||||
def _build_delta_stage2_sql(
|
def _build_delta_fetch_sql(
|
||||||
*,
|
*,
|
||||||
channels_with_chain: Sequence[str],
|
channels_with_chain: Sequence[str],
|
||||||
channels_with_seed: Sequence[str],
|
channels_with_seed: Sequence[str],
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Build stage 2 SQL as a per-channel UNION ALL.
|
"""Build the FETCH SQL as a per-channel UNION ALL.
|
||||||
|
|
||||||
For each channel with a non-empty chain, emit one branch reading
|
For each channel with a non-empty chain, emit one branch reading
|
||||||
`checkpoint_writes` for that specific channel + chain_cids. For each
|
`checkpoint_writes` for that specific channel + chain_cids. For each
|
||||||
channel with a seed_version, emit one branch reading `checkpoint_blobs`
|
channel with a seed_version, emit one branch reading `checkpoint_blobs`
|
||||||
for that channel + version. This avoids the over-fetch of the prior
|
for that channel + version. This avoids the over-fetch of a single
|
||||||
`channel = ANY(channels) AND checkpoint_id = ANY(union)` form when
|
`channel = ANY(channels) AND checkpoint_id = ANY(union)` form when
|
||||||
channels have different chain depths.
|
channels have different chain depths.
|
||||||
|
|
||||||
@@ -269,6 +281,7 @@ def _build_delta_stage2_sql(
|
|||||||
"""
|
"""
|
||||||
branches: list[str] = []
|
branches: list[str] = []
|
||||||
for _ in channels_with_chain:
|
for _ in channels_with_chain:
|
||||||
|
# NOTE: no ORDER BY on this branch — writes are sorted in assembly.
|
||||||
branches.append(
|
branches.append(
|
||||||
"SELECT 'w'::text AS _kind, "
|
"SELECT 'w'::text AS _kind, "
|
||||||
"checkpoint_id, channel, "
|
"checkpoint_id, channel, "
|
||||||
@@ -288,10 +301,59 @@ def _build_delta_stage2_sql(
|
|||||||
return " UNION ALL ".join(branches)
|
return " UNION ALL ".join(branches)
|
||||||
|
|
||||||
|
|
||||||
# Stage 1 rows are dynamic-shape dicts: {checkpoint_id, parent_checkpoint_id,
|
def _ingest_walk_page(
|
||||||
# ver_0, hs_0, ver_1, hs_1, ...}. Walking is parameterized by the channel
|
page_rows: Sequence[Mapping[str, Any]],
|
||||||
# list to map indices back to channel names — no static TypedDict here.
|
channels: Sequence[str],
|
||||||
# `dict[str, Any]` is the practical signature.
|
parent_of: dict[str, str | None],
|
||||||
|
ver_by_i_by_cid: list[dict[str, str | None]],
|
||||||
|
) -> str | None:
|
||||||
|
"""Fold one WALK page into `parent_of` + per-channel `ver_by_cid`.
|
||||||
|
|
||||||
|
Returns the oldest checkpoint_id seen on this page (smallest, since
|
||||||
|
pages come back DESC). Caller uses it as the cursor for the next page
|
||||||
|
(`AND checkpoint_id < cursor`).
|
||||||
|
"""
|
||||||
|
oldest: str | None = None
|
||||||
|
for r in page_rows:
|
||||||
|
cid = cast(str, r["checkpoint_id"])
|
||||||
|
parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
|
||||||
|
for i in range(len(channels)):
|
||||||
|
ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
|
||||||
|
# Rows are DESC; the last one is the smallest cid in the page.
|
||||||
|
oldest = cid
|
||||||
|
return oldest
|
||||||
|
|
||||||
|
|
||||||
|
def _advance_shared_chain(
|
||||||
|
target_id: str,
|
||||||
|
parent_of: Mapping[str, str | None],
|
||||||
|
shared_cpid_chain: list[str],
|
||||||
|
max_supersteps: int,
|
||||||
|
) -> bool:
|
||||||
|
"""Extend the shared parent chain as far as ingested pages allow.
|
||||||
|
|
||||||
|
The chain holds ancestors of the target, newest first: `chain[0]` is
|
||||||
|
the target's parent, `chain[1]` its grandparent, and so on. A single
|
||||||
|
chain is shared across all channels and grown to the maximum requested
|
||||||
|
depth; each channel later slices `chain[:supersteps]`.
|
||||||
|
|
||||||
|
Stops when:
|
||||||
|
- the chain reaches `max_supersteps` hops, OR
|
||||||
|
- the root is reached (parent is None) — returns True, OR
|
||||||
|
- the next ancestor's cid isn't in `parent_of` yet (waits for the
|
||||||
|
next page).
|
||||||
|
|
||||||
|
Returns True iff the root was reached.
|
||||||
|
"""
|
||||||
|
while len(shared_cpid_chain) < max_supersteps:
|
||||||
|
top_of_chain = shared_cpid_chain[-1] if shared_cpid_chain else target_id
|
||||||
|
if top_of_chain not in parent_of:
|
||||||
|
return False # wait for the next page
|
||||||
|
parent = parent_of[top_of_chain]
|
||||||
|
if parent is None:
|
||||||
|
return True # hit the root
|
||||||
|
shared_cpid_chain.append(parent)
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||||
@@ -336,107 +398,92 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|||||||
if t.decode() != "empty"
|
if t.decode() != "empty"
|
||||||
}
|
}
|
||||||
|
|
||||||
@staticmethod
|
def _resolve_delta_chains(
|
||||||
def _ingest_stage1_page(
|
self,
|
||||||
stage1_rows: Sequence[Mapping[str, Any]],
|
|
||||||
channels: Sequence[str],
|
channels: Sequence[str],
|
||||||
parent_of: dict[str, str | None],
|
supersteps_by_ch: Mapping[str, int],
|
||||||
ver_by_i_by_cid: list[dict[str, str | None]],
|
shared_cpid_chain: Sequence[str],
|
||||||
hs_by_i_by_cid: list[dict[str, bool]],
|
|
||||||
) -> str | None:
|
|
||||||
"""Fold one stage-1 page into the running walk-state mappings.
|
|
||||||
|
|
||||||
Returns the oldest checkpoint_id seen on this page (smallest, since
|
|
||||||
pages come back DESC). Caller uses it as the cursor for the next
|
|
||||||
page (`AND checkpoint_id < cursor`).
|
|
||||||
"""
|
|
||||||
oldest: str | None = None
|
|
||||||
for r in stage1_rows:
|
|
||||||
cid = cast(str, r["checkpoint_id"])
|
|
||||||
parent_of[cid] = cast("str | None", r["parent_checkpoint_id"])
|
|
||||||
for i in range(len(channels)):
|
|
||||||
ver_by_i_by_cid[i][cid] = cast("str | None", r.get(f"ver_{i}"))
|
|
||||||
hs_by_i_by_cid[i][cid] = bool(r.get(f"hs_{i}"))
|
|
||||||
# Rows are DESC; the last one is the smallest cid in the page.
|
|
||||||
oldest = cid
|
|
||||||
return oldest
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _try_advance_walks(
|
|
||||||
target_id: str,
|
|
||||||
channels: Sequence[str],
|
|
||||||
parent_of: Mapping[str, str | None],
|
|
||||||
ver_by_i_by_cid: Sequence[Mapping[str, str | None]],
|
ver_by_i_by_cid: Sequence[Mapping[str, str | None]],
|
||||||
hs_by_i_by_cid: Sequence[Mapping[str, bool]],
|
has_reached_root: bool,
|
||||||
chain_by_ch: dict[str, list[str]],
|
thread_id: str,
|
||||||
seed_ver_by_ch: dict[str, str | None],
|
) -> tuple[
|
||||||
walk_cursor_by_ch: dict[str, str | None],
|
dict[str, list[str]],
|
||||||
seeded: set[str],
|
dict[str, str | None],
|
||||||
) -> None:
|
dict[str, str | None],
|
||||||
"""Advance each not-yet-seeded channel's walk as far as possible.
|
]:
|
||||||
|
"""Slice the shared parent chain into per-channel chain/seed mappings.
|
||||||
|
|
||||||
Uses the partial `parent_of` map accumulated so far. A walk stops
|
For each channel the seed snapshot sits `supersteps` hops back, so
|
||||||
either because:
|
the seed checkpoint is `shared_cpid_chain[supersteps - 1]` and the
|
||||||
(a) it found a snapshot for its channel (channel becomes seeded),
|
chain is `shared_cpid_chain[:supersteps]` (newest first).
|
||||||
(b) it reached a real root (parent_of[cid] is None — fully
|
|
||||||
materialized at this point), or
|
|
||||||
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
|
||||||
a later page; the cursor stays put).
|
|
||||||
|
|
||||||
Mutates `chain_by_ch`, `seed_ver_by_ch`, `walk_cursor_by_ch`, and
|
When the chain is shorter than `supersteps` but the walk reached the
|
||||||
`seeded` in place.
|
root, the persisted chain is "compressed" relative to the logical
|
||||||
|
superstep count — either because intermediate supersteps were never
|
||||||
|
persisted (`durability="exit"`) or because the thread never produced
|
||||||
|
a snapshot at all. In both cases the seed candidate is the oldest
|
||||||
|
persisted checkpoint (`shared_cpid_chain[-1]`): FETCH loads its blob,
|
||||||
|
and assembly keeps it only if non-empty (a real snapshot or migrated
|
||||||
|
value) — otherwise it omits `seed` and replays the full chain on an
|
||||||
|
empty baseline.
|
||||||
"""
|
"""
|
||||||
|
chained_cpid_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
||||||
|
seed_cpid_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
||||||
|
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
||||||
for i, ch in enumerate(channels):
|
for i, ch in enumerate(channels):
|
||||||
if ch in seeded:
|
bound = supersteps_by_ch.get(ch, 0)
|
||||||
|
if bound <= 0:
|
||||||
continue
|
continue
|
||||||
# First-time entry: cursor starts at the target's parent.
|
if len(shared_cpid_chain) >= bound:
|
||||||
if ch not in walk_cursor_by_ch:
|
seed_depth = bound
|
||||||
walk_cursor_by_ch[ch] = parent_of.get(target_id)
|
elif has_reached_root:
|
||||||
cur_cid = walk_cursor_by_ch[ch]
|
seed_depth = len(shared_cpid_chain)
|
||||||
ch_chain = chain_by_ch[ch]
|
else:
|
||||||
hs_i = hs_by_i_by_cid[i]
|
logger.warning(
|
||||||
ver_i = ver_by_i_by_cid[i]
|
"cannot find seed snapshot for delta channel "
|
||||||
while cur_cid is not None:
|
"(thread_id=%s, channel=%s)",
|
||||||
if cur_cid not in parent_of:
|
thread_id,
|
||||||
# Need more pages to continue this walk.
|
ch,
|
||||||
break
|
)
|
||||||
ch_chain.append(cur_cid)
|
continue
|
||||||
if hs_i.get(cur_cid, False):
|
if seed_depth <= 0:
|
||||||
seed_ver_by_ch[ch] = ver_i.get(cur_cid)
|
continue
|
||||||
seeded.add(ch)
|
chained_cpid_by_ch[ch] = list(shared_cpid_chain[:seed_depth])
|
||||||
cur_cid = None
|
seed_cpid_by_ch[ch] = shared_cpid_chain[seed_depth - 1]
|
||||||
break
|
seed_ver_by_ch[ch] = ver_by_i_by_cid[i].get(seed_cpid_by_ch[ch])
|
||||||
cur_cid = parent_of[cur_cid]
|
return chained_cpid_by_ch, seed_cpid_by_ch, seed_ver_by_ch
|
||||||
walk_cursor_by_ch[ch] = cur_cid
|
|
||||||
|
|
||||||
def _build_delta_channels_writes_history(
|
def _assemble_delta_history(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
channels: Sequence[str],
|
channels: Sequence[str],
|
||||||
chain_by_ch: Mapping[str, list[str]],
|
chained_cpid_by_ch: Mapping[str, Sequence[str]],
|
||||||
|
seed_cpid_by_ch: Mapping[str, str | None],
|
||||||
seed_ver_by_ch: Mapping[str, str | None],
|
seed_ver_by_ch: Mapping[str, str | None],
|
||||||
stage2_rows: Sequence[_DeltaStage2Row],
|
fetch_rows: Sequence[_DeltaStage2Row],
|
||||||
) -> dict[str, DeltaChannelHistory]:
|
) -> dict[str, DeltaChannelHistory]:
|
||||||
"""Demux stage 2 rows per channel; produce per-channel histories.
|
"""Demux FETCH rows per channel and produce per-channel histories.
|
||||||
|
|
||||||
stage2_rows carry `channel` on every row. We build per-channel
|
`fetch_rows` carry `channel` on every row. Write rows (`_kind = 'w'`)
|
||||||
`writes_by_cid` and per-channel `seed_blob` dicts, then assemble
|
are bucketed per channel per checkpoint; seed-blob rows
|
||||||
a `DeltaChannelHistory` per requested channel. The `seed` key is omitted
|
(`_kind = 'b'`) give each channel its snapshot value.
|
||||||
when the walk reached root with no snapshot found, or when the
|
|
||||||
seed blob is sentinel "empty" — in both cases the consumer treats
|
The seed checkpoint's own writes are replayed on top of a
|
||||||
absence as "start empty".
|
`_DeltaSnapshot` seed (the snapshot is the value *prior* to its own
|
||||||
|
writes), but skipped for a migrated plain-value seed (a legacy
|
||||||
|
non-delta blob already incorporates those writes). The `seed` key is
|
||||||
|
omitted when no seed was located or the blob is the "empty"
|
||||||
|
tombstone — the consumer treats absence as "start empty".
|
||||||
"""
|
"""
|
||||||
# writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx)
|
# writes_by_ch_by_cid[channel][cid] = list of (type, blob, task_id, idx)
|
||||||
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
|
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
|
||||||
ch: {} for ch in channels
|
ch: {} for ch in channels
|
||||||
}
|
}
|
||||||
# seed_blob_by_ver[(channel, version)] = (type, blob)
|
seed_blob_by_ch: dict[str, tuple[str, bytes]] = {}
|
||||||
seed_blob_by_ver: dict[tuple[str, str], tuple[str, bytes]] = {}
|
|
||||||
|
|
||||||
for r in stage2_rows:
|
for r in fetch_rows:
|
||||||
ch = cast(str, r["channel"])
|
ch = cast(str, r["channel"])
|
||||||
kind = r["_kind"]
|
if r["_kind"] == "w":
|
||||||
if kind == "w":
|
|
||||||
cid = cast(str, r["checkpoint_id"])
|
cid = cast(str, r["checkpoint_id"])
|
||||||
writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
|
writes_by_ch_by_cid.setdefault(ch, {}).setdefault(cid, []).append(
|
||||||
cast(
|
cast(
|
||||||
@@ -444,35 +491,39 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|||||||
(r["type"], r["blob"], r["task_id"], r["idx"]),
|
(r["type"], r["blob"], r["task_id"], r["idx"]),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
else: # kind == "b"
|
else: # _kind == "b" — the seed blob for this channel.
|
||||||
ver = cast(str, r["version"])
|
seed_blob_by_ch[ch] = cast("tuple[str, bytes]", (r["type"], r["blob"]))
|
||||||
seed_blob_by_ver[(ch, ver)] = cast(
|
|
||||||
"tuple[str, bytes]", (r["type"], r["blob"])
|
|
||||||
)
|
|
||||||
|
|
||||||
# Sort writes per (channel, cid) newest-first by (task_id, idx)
|
# Within a checkpoint, writes apply oldest→newest by (task_id, idx).
|
||||||
for cid_map in writes_by_ch_by_cid.values():
|
for cid_map in writes_by_ch_by_cid.values():
|
||||||
for ws in cid_map.values():
|
for ws in cid_map.values():
|
||||||
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
|
ws.sort(key=lambda w: (w[2], w[3]))
|
||||||
|
|
||||||
result: dict[str, DeltaChannelHistory] = {}
|
result: dict[str, DeltaChannelHistory] = {}
|
||||||
for ch in channels:
|
for ch in channels:
|
||||||
chain_cids = chain_by_ch.get(ch, [])
|
entry: DeltaChannelHistory = {"writes": []}
|
||||||
seed_version = seed_ver_by_ch.get(ch)
|
|
||||||
|
skip_seed_checkpoint_writes = False
|
||||||
|
seed_blob = seed_blob_by_ch.get(ch)
|
||||||
|
if seed_blob is not None and seed_blob[0] != "empty":
|
||||||
|
seed_value = self.serde.loads_typed(seed_blob)
|
||||||
|
entry["seed"] = seed_value
|
||||||
|
# A migrated (non-delta) seed already includes the writes on
|
||||||
|
# its own checkpoint; a `_DeltaSnapshot` does not.
|
||||||
|
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
|
||||||
|
|
||||||
collected: list[PendingWrite] = []
|
|
||||||
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
||||||
for cid in chain_cids:
|
if cid_writes:
|
||||||
for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []):
|
collected: list[PendingWrite] = []
|
||||||
val = self.serde.loads_typed((type_tag, write_blob))
|
seed_cpid = seed_cpid_by_ch.get(ch)
|
||||||
collected.append((task_id, ch, val))
|
# Chain is newest→oldest; replay oldest→newest.
|
||||||
collected.reverse()
|
for cid in reversed(chained_cpid_by_ch.get(ch, [])):
|
||||||
|
if skip_seed_checkpoint_writes and cid == seed_cpid:
|
||||||
entry: DeltaChannelHistory = {"writes": collected}
|
continue
|
||||||
if seed_version is not None:
|
for type_tag, write_blob, task_id, _idx in cid_writes.get(cid, []):
|
||||||
blob = seed_blob_by_ver.get((ch, seed_version))
|
val = self.serde.loads_typed((type_tag, write_blob))
|
||||||
if blob is not None and blob[0] != "empty":
|
collected.append((task_id, ch, val))
|
||||||
entry["seed"] = self.serde.loads_typed(blob)
|
entry["writes"] = collected
|
||||||
result[ch] = entry
|
result[ch] = entry
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
@@ -24,10 +24,12 @@ from langgraph.checkpoint.base import (
|
|||||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||||
|
|
||||||
from langgraph.checkpoint.sqlite._delta import (
|
from langgraph.checkpoint.sqlite._delta import (
|
||||||
DELTA_STAGE1_SQL,
|
DELTA_WALK_SQL,
|
||||||
build_delta_channels_writes_history,
|
build_delta_channels_writes_history,
|
||||||
build_delta_stage2_sql,
|
build_delta_writes_fetch_sql,
|
||||||
step_walk_with_row,
|
parse_supersteps_since_last_snapshot_by_channel,
|
||||||
|
resolve_delta_chains,
|
||||||
|
step_walk_supersteps,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.sqlite.utils import search_where
|
from langgraph.checkpoint.sqlite.utils import search_where
|
||||||
|
|
||||||
@@ -505,68 +507,106 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|||||||
) -> Mapping[str, DeltaChannelHistory]:
|
) -> Mapping[str, DeltaChannelHistory]:
|
||||||
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
|
"""Fast-path override of `BaseCheckpointSaver.get_delta_channel_history`.
|
||||||
|
|
||||||
Two-stage query:
|
Walks the parent chain `supersteps` hops (read from the target's
|
||||||
|
`counters_since_delta_snapshot` metadata) to each channel's seed
|
||||||
|
snapshot, then collects the writes between the seed and the target.
|
||||||
|
|
||||||
* Stage 1 (paged): newest-first slice of `checkpoints` returning
|
* WALK: stream a newest-first slice of `checkpoints` returning
|
||||||
`(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per
|
`(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per
|
||||||
ancestor. Sqlite has no JSONB, so we ship the full serialized
|
ancestor. Sqlite has no JSONB, so we deserialize only the seed
|
||||||
checkpoint blob and inspect `channel_values` in Python. Pages
|
checkpoints to read their inline `channel_values`. Stops once the
|
||||||
newest-first by `checkpoint_id` with a `< cursor` predicate;
|
shared chain reaches the deepest requested `supersteps` or the
|
||||||
page size is `DELTA_PAGE_SIZE`. Stops paging when every channel
|
root is reached.
|
||||||
has found its seed or the chain is exhausted.
|
|
||||||
|
|
||||||
* Stage 2 (per-channel UNION ALL): one branch per channel reading
|
* FETCH (per-channel UNION ALL): one branch per channel reading
|
||||||
`writes` filtered to that channel's specific `chain_cids`. No
|
`writes` filtered to that channel's `chain_cids`. No separate
|
||||||
separate seed-blob fetch — sqlite stores `channel_values` inline
|
seed-blob fetch — sqlite stores `channel_values` inline.
|
||||||
in the checkpoint blob, so seeds come back from stage 1.
|
|
||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
channels = list(channels)
|
channels = list(channels)
|
||||||
thread_id = str(config["configurable"]["thread_id"])
|
thread_id = str(config["configurable"]["thread_id"])
|
||||||
|
if not thread_id:
|
||||||
|
raise ValueError("empty thread ID")
|
||||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
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}
|
# Resolve the target checkpoint id + its metadata (for supersteps).
|
||||||
seed_val_by_ch: dict[str, Any] = {}
|
checkpoint_id = get_checkpoint_id(config)
|
||||||
|
with self.cursor(transaction=False) as cur:
|
||||||
|
if checkpoint_id is None:
|
||||||
|
cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = ? AND checkpoint_ns = ? "
|
||||||
|
"ORDER BY checkpoint_id DESC LIMIT 1",
|
||||||
|
(thread_id, checkpoint_ns),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = ? AND checkpoint_ns = ? "
|
||||||
|
"AND checkpoint_id = ?",
|
||||||
|
(thread_id, checkpoint_ns, checkpoint_id),
|
||||||
|
)
|
||||||
|
target_row = cur.fetchone()
|
||||||
|
if target_row is None:
|
||||||
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
target_id = str(target_row[0])
|
||||||
|
metadata = json.loads(target_row[1]) if target_row[1] is not None else {}
|
||||||
|
supersteps_by_ch = parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
metadata, channels
|
||||||
|
)
|
||||||
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
needed_depths = set(supersteps_by_ch.values())
|
||||||
|
|
||||||
|
shared_cpid_chain: list[str] = []
|
||||||
walk_state: dict[str, Any] = {}
|
walk_state: dict[str, Any] = {}
|
||||||
seeded: set[str] = set()
|
seed_values_by_depth: dict[int, dict[str, Any]] = {}
|
||||||
|
|
||||||
with self.cursor(transaction=False) as cur:
|
with self.cursor(transaction=False) as cur:
|
||||||
cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id))
|
if max_supersteps > 0:
|
||||||
for row in cur:
|
cur.execute(DELTA_WALK_SQL, (thread_id, checkpoint_ns, target_id))
|
||||||
cid, parent_cid, type_tag, blob = row
|
for row in cur:
|
||||||
if step_walk_with_row(
|
cid, parent_cid, type_tag, blob = row
|
||||||
cid=cid,
|
if step_walk_supersteps(
|
||||||
parent_cid=parent_cid,
|
cid=cid,
|
||||||
type_tag=type_tag,
|
parent_cid=parent_cid,
|
||||||
blob=blob,
|
type_tag=type_tag,
|
||||||
target_id=checkpoint_id,
|
blob=blob,
|
||||||
serde=self.serde,
|
target_id=target_id,
|
||||||
chain_by_ch=chain_by_ch,
|
serde=self.serde,
|
||||||
seed_val_by_ch=seed_val_by_ch,
|
shared_cpid_chain=shared_cpid_chain,
|
||||||
walk_state=walk_state,
|
walk_state=walk_state,
|
||||||
seeded=seeded,
|
max_supersteps=max_supersteps,
|
||||||
channels=channels,
|
needed_depths=needed_depths,
|
||||||
):
|
seed_values_by_depth=seed_values_by_depth,
|
||||||
break
|
channels=channels,
|
||||||
|
):
|
||||||
|
break
|
||||||
|
|
||||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
(
|
||||||
stage2_sql = build_delta_stage2_sql(
|
chained_cpid_by_ch,
|
||||||
chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain],
|
seed_cpid_by_ch,
|
||||||
|
seed_value_by_ch,
|
||||||
|
) = resolve_delta_chains(
|
||||||
|
channels=channels,
|
||||||
|
supersteps_by_ch=supersteps_by_ch,
|
||||||
|
shared_cpid_chain=shared_cpid_chain,
|
||||||
|
seed_values_by_depth=seed_values_by_depth,
|
||||||
|
has_reached_root=bool(walk_state.get("reached_root")),
|
||||||
|
thread_id=thread_id,
|
||||||
)
|
)
|
||||||
if stage2_sql:
|
|
||||||
stage2_params: list[Any] = []
|
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
|
||||||
|
fetch_sql = build_delta_writes_fetch_sql(
|
||||||
|
chain_lens=[len(chained_cpid_by_ch[ch]) for ch in channels_with_chain],
|
||||||
|
)
|
||||||
|
if fetch_sql:
|
||||||
|
fetch_params: list[Any] = []
|
||||||
for ch in channels_with_chain:
|
for ch in channels_with_chain:
|
||||||
stage2_params.extend(
|
fetch_params.extend(
|
||||||
[thread_id, checkpoint_ns, ch, *chain_by_ch[ch]]
|
[thread_id, checkpoint_ns, ch, *chained_cpid_by_ch[ch]]
|
||||||
)
|
)
|
||||||
cur.execute(stage2_sql, stage2_params)
|
cur.execute(fetch_sql, fetch_params)
|
||||||
stage2_rows = cast(
|
stage2_rows = cast(
|
||||||
"list[tuple[str, str, str, int, str, bytes]]", cur.fetchall()
|
"list[tuple[str, str, str, int, str, bytes]]", cur.fetchall()
|
||||||
)
|
)
|
||||||
@@ -575,9 +615,9 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|||||||
|
|
||||||
return build_delta_channels_writes_history(
|
return build_delta_channels_writes_history(
|
||||||
channels=channels,
|
channels=channels,
|
||||||
chain_by_ch=chain_by_ch,
|
chained_cpid_by_ch=chained_cpid_by_ch,
|
||||||
seed_val_by_ch=seed_val_by_ch,
|
seed_cpid_by_ch=seed_cpid_by_ch,
|
||||||
seeded=seeded,
|
seed_value_by_ch=seed_value_by_ch,
|
||||||
stage2_rows=stage2_rows,
|
stage2_rows=stage2_rows,
|
||||||
serde=self.serde,
|
serde=self.serde,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,37 +1,50 @@
|
|||||||
"""Shared helpers for `get_delta_channel_history` on sqlite savers.
|
"""Shared helpers for `get_delta_channel_history` on sqlite savers.
|
||||||
|
|
||||||
Mirrors the two-stage shape of `BasePostgresSaver` (ancestor walk +
|
Mirrors the supersteps-based two-pass shape of `BasePostgresSaver`
|
||||||
per-channel UNION ALL writes fetch), but adapted for sqlite's
|
(ancestor walk bounded by `counters_since_delta_snapshot` + per-channel
|
||||||
constraints. The structural differences:
|
UNION ALL writes fetch), adapted for sqlite's constraints:
|
||||||
|
|
||||||
* No JSONB — to inspect `channel_values` for a checkpoint we must
|
* No JSONB — to inspect `channel_values` for a checkpoint we must
|
||||||
deserialize the full blob. Stage 1 streams the cursor row-by-row and
|
deserialize the full blob. The WALK streams the cursor row-by-row and
|
||||||
deserializes only the rows the merged walk visits, freeing each blob
|
deserializes only the seed checkpoints (the ones at a channel's
|
||||||
before advancing.
|
`supersteps` depth), freeing each blob before advancing.
|
||||||
* No separate blob table — `channel_values` lives inline in the
|
* No separate blob table — `channel_values` lives inline in the
|
||||||
checkpoint, so seeds come back from stage 1 with no second fetch.
|
checkpoint, so seeds come back from the WALK with no second fetch.
|
||||||
* Single merged walk (not K independent walks): each visited cid is
|
* Single shared parent-chain walk: each requested channel slices the
|
||||||
deserialized exactly once, regardless of how many channels are still
|
same chain to its own `supersteps` depth.
|
||||||
seeking their seed.
|
|
||||||
|
|
||||||
The streaming design keeps peak in-flight memory at roughly one
|
Walk depth is driven by the *supersteps since last snapshot* counter,
|
||||||
deserialized checkpoint at a time, instead of holding the entire
|
not by scanning `channel_values` for the snapshot marker — the only
|
||||||
ancestor chain's worth of raw blobs as a `fetchall()`-materialized list.
|
reliable way to locate seeds that aren't a `_DeltaSnapshot` sentinel
|
||||||
|
(e.g. legacy plain-value blobs from a thread migrated off a non-delta
|
||||||
|
channel).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Mapping, Sequence
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite
|
from langgraph.checkpoint.base import (
|
||||||
|
DeltaChannelHistory,
|
||||||
|
PendingWrite,
|
||||||
|
_parse_supersteps_since_last_snapshot_by_channel,
|
||||||
|
)
|
||||||
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||||
|
|
||||||
# Stage 1 streams ancestors of `target_cid` newest-first. The `<=`
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Re-exported under the package-local name used by the sqlite savers.
|
||||||
|
parse_supersteps_since_last_snapshot_by_channel = (
|
||||||
|
_parse_supersteps_since_last_snapshot_by_channel
|
||||||
|
)
|
||||||
|
|
||||||
|
# The WALK streams ancestors of `target_id` newest-first. The `<=`
|
||||||
# predicate keeps target itself in the stream so we can read its
|
# predicate keeps target itself in the stream so we can read its
|
||||||
# `parent_checkpoint_id` from the first row without a separate lookup;
|
# `parent_checkpoint_id` from the first matching row without a separate
|
||||||
# the caller skips target's own writes/seed (matches the
|
# lookup; target's own writes/seed are not part of the contract.
|
||||||
# `BaseCheckpointSaver` contract).
|
DELTA_WALK_SQL = (
|
||||||
DELTA_STAGE1_SQL = (
|
|
||||||
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
|
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
|
||||||
"FROM checkpoints "
|
"FROM checkpoints "
|
||||||
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
|
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
|
||||||
@@ -39,12 +52,12 @@ DELTA_STAGE1_SQL = (
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str:
|
def build_delta_writes_fetch_sql(*, chain_lens: Sequence[int]) -> str:
|
||||||
"""Stage-2 per-channel UNION ALL fetching writes from `writes`.
|
"""Per-channel UNION ALL fetching writes from `writes`.
|
||||||
|
|
||||||
One branch per channel with a non-empty chain. Each branch inlines its
|
One branch per channel with a non-empty chain. Each branch inlines its
|
||||||
own `IN (?, ?, ...)` placeholder list because sqlite has no array-bind
|
own `IN (?, ?, ...)` placeholder list because sqlite has no array-bind
|
||||||
equivalent of postgres's `= ANY(%s)`. Caller passes parameters in
|
equivalent of postgres's `= ANY(?)`. Caller passes parameters in
|
||||||
matching order: `[thread_id, checkpoint_ns, channel, *chain_cids]` per
|
matching order: `[thread_id, checkpoint_ns, channel, *chain_cids]` per
|
||||||
branch.
|
branch.
|
||||||
|
|
||||||
@@ -65,7 +78,7 @@ def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str:
|
|||||||
return " UNION ALL ".join(branches)
|
return " UNION ALL ".join(branches)
|
||||||
|
|
||||||
|
|
||||||
def step_walk_with_row(
|
def step_walk_supersteps(
|
||||||
*,
|
*,
|
||||||
cid: str,
|
cid: str,
|
||||||
parent_cid: str | None,
|
parent_cid: str | None,
|
||||||
@@ -73,75 +86,142 @@ def step_walk_with_row(
|
|||||||
blob: bytes,
|
blob: bytes,
|
||||||
target_id: str,
|
target_id: str,
|
||||||
serde: Any,
|
serde: Any,
|
||||||
chain_by_ch: dict[str, list[str]],
|
shared_cpid_chain: list[str],
|
||||||
seed_val_by_ch: dict[str, Any],
|
|
||||||
walk_state: dict[str, Any],
|
walk_state: dict[str, Any],
|
||||||
seeded: set[str],
|
max_supersteps: int,
|
||||||
|
needed_depths: set[int],
|
||||||
|
seed_values_by_depth: dict[int, dict[str, Any]],
|
||||||
channels: Sequence[str],
|
channels: Sequence[str],
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""Process one streamed stage-1 row in the merged ancestor walk.
|
"""Process one streamed WALK row, extending the shared parent chain.
|
||||||
|
|
||||||
The cursor returns (cid, parent_cid, type, blob) rows in
|
The cursor returns `(cid, parent_cid, type, checkpoint)` rows in
|
||||||
`checkpoint_id` DESC order starting at target. The first row is
|
`checkpoint_id` DESC order starting at target. The first row is target
|
||||||
target itself; we read its parent_cid to seed the walk and otherwise
|
itself; we read its `parent_cid` to seed the walk and skip it (target's
|
||||||
skip it (target's own writes/seed are not part of the contract).
|
own writes/seed are not part of the contract). Off-path rows (a sibling
|
||||||
|
branch on the same thread) advance the cursor without doing work.
|
||||||
|
|
||||||
For each subsequent row, if `cid` matches the walk's current
|
For each on-path ancestor we append its cid to `shared_cpid_chain`
|
||||||
position, we deserialize the blob, append the cid to every
|
(newest first). When the ancestor sits at a depth some channel needs as
|
||||||
not-yet-seeded channel's chain, and check `channel_values` for
|
its seed (`len(chain)` ∈ `needed_depths`), we deserialize it once and
|
||||||
seeds. The deserialized checkpoint is dropped before advancing — no
|
record its `channel_values` for the requested channels. The
|
||||||
cross-row cache, so peak in-flight is one deserialized checkpoint.
|
deserialized checkpoint is dropped immediately — peak in-flight is one
|
||||||
|
deserialized checkpoint.
|
||||||
|
|
||||||
Off-path rows (different branch on the same thread) advance the
|
Sets `walk_state["reached_root"]` when an ancestor has no parent.
|
||||||
cursor without doing any work.
|
Returns True when the walk can stop: chain reached `max_supersteps`, or
|
||||||
|
the root was reached.
|
||||||
Returns True when every requested channel is seeded — the caller
|
|
||||||
can stop iterating and close the cursor.
|
|
||||||
"""
|
"""
|
||||||
if "started" not in walk_state:
|
if "started" not in walk_state:
|
||||||
if cid == target_id:
|
if cid == target_id:
|
||||||
walk_state["started"] = True
|
walk_state["started"] = True
|
||||||
walk_state["cur_cid"] = parent_cid
|
walk_state["cur_cid"] = parent_cid
|
||||||
walk_state["active"] = {ch for ch in channels if ch not in seeded}
|
if parent_cid is None:
|
||||||
|
walk_state["reached_root"] = True
|
||||||
|
return True
|
||||||
# Not target yet (or target not present): keep streaming.
|
# Not target yet (or target not present): keep streaming.
|
||||||
return False
|
return False
|
||||||
active: set[str] = walk_state["active"]
|
if len(shared_cpid_chain) >= max_supersteps:
|
||||||
if not active:
|
|
||||||
return True
|
return True
|
||||||
if cid != walk_state["cur_cid"]:
|
if cid != walk_state["cur_cid"]:
|
||||||
# Off-path row from a sibling branch — skip without deserializing.
|
# Off-path row from a sibling branch — skip without deserializing.
|
||||||
return False
|
return False
|
||||||
for ch in active:
|
shared_cpid_chain.append(cid)
|
||||||
chain_by_ch[ch].append(cid)
|
depth = len(shared_cpid_chain) # 1-indexed position along the chain
|
||||||
ckpt = serde.loads_typed((type_tag, blob))
|
# Capture channel_values at any depth a channel may use as its seed: the
|
||||||
channel_values: Mapping[str, Any] = ckpt.get("channel_values") or {}
|
# exact `supersteps` depths, plus the root-most checkpoint (the seed
|
||||||
for ch in [ch for ch in active if ch in channel_values]:
|
# candidate when the chain is shorter than `supersteps`).
|
||||||
seed_val_by_ch[ch] = channel_values[ch]
|
if depth in needed_depths or parent_cid is None:
|
||||||
seeded.add(ch)
|
ckpt = serde.loads_typed((type_tag, blob))
|
||||||
active.discard(ch)
|
channel_values: Mapping[str, Any] = ckpt.get("channel_values") or {}
|
||||||
del ckpt, channel_values
|
seed_values_by_depth[depth] = {
|
||||||
|
ch: channel_values[ch] for ch in channels if ch in channel_values
|
||||||
|
}
|
||||||
|
del ckpt, channel_values
|
||||||
|
if parent_cid is None:
|
||||||
|
walk_state["reached_root"] = True
|
||||||
|
return True
|
||||||
walk_state["cur_cid"] = parent_cid
|
walk_state["cur_cid"] = parent_cid
|
||||||
return not active
|
return len(shared_cpid_chain) >= max_supersteps
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_delta_chains(
|
||||||
|
*,
|
||||||
|
channels: Sequence[str],
|
||||||
|
supersteps_by_ch: Mapping[str, int],
|
||||||
|
shared_cpid_chain: Sequence[str],
|
||||||
|
seed_values_by_depth: Mapping[int, Mapping[str, Any]],
|
||||||
|
has_reached_root: bool,
|
||||||
|
thread_id: str,
|
||||||
|
) -> tuple[
|
||||||
|
dict[str, list[str]],
|
||||||
|
dict[str, str | None],
|
||||||
|
dict[str, Any],
|
||||||
|
]:
|
||||||
|
"""Slice the shared parent chain into per-channel chain/seed mappings.
|
||||||
|
|
||||||
|
For each channel the seed snapshot sits `supersteps` hops back: the seed
|
||||||
|
checkpoint is `shared_cpid_chain[supersteps - 1]` and the chain is
|
||||||
|
`shared_cpid_chain[:supersteps]` (newest first), with the seed's inline
|
||||||
|
`channel_values[ch]` captured during the walk.
|
||||||
|
|
||||||
|
When the chain is shorter than `supersteps` but the walk reached the
|
||||||
|
root, the persisted chain is "compressed" relative to the logical
|
||||||
|
superstep count (`durability="exit"`, or a thread that never
|
||||||
|
snapshotted). The seed candidate is then the oldest persisted checkpoint
|
||||||
|
(`shared_cpid_chain[-1]`); if its `channel_values[ch]` is absent the
|
||||||
|
channel has no seed and replays the full chain on an empty baseline.
|
||||||
|
"""
|
||||||
|
chained_cpid_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
||||||
|
seed_cpid_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
||||||
|
seed_value_by_ch: dict[str, Any] = {}
|
||||||
|
for ch in channels:
|
||||||
|
bound = supersteps_by_ch.get(ch, 0)
|
||||||
|
if bound <= 0:
|
||||||
|
continue
|
||||||
|
if len(shared_cpid_chain) >= bound:
|
||||||
|
seed_depth = bound
|
||||||
|
elif has_reached_root:
|
||||||
|
seed_depth = len(shared_cpid_chain)
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"cannot find seed snapshot for delta channel "
|
||||||
|
"(thread_id=%s, channel=%s)",
|
||||||
|
thread_id,
|
||||||
|
ch,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
if seed_depth <= 0:
|
||||||
|
continue
|
||||||
|
chained_cpid_by_ch[ch] = list(shared_cpid_chain[:seed_depth])
|
||||||
|
seed_cpid_by_ch[ch] = shared_cpid_chain[seed_depth - 1]
|
||||||
|
seed_vals = seed_values_by_depth.get(seed_depth, {})
|
||||||
|
if ch in seed_vals:
|
||||||
|
seed_value_by_ch[ch] = seed_vals[ch]
|
||||||
|
return chained_cpid_by_ch, seed_cpid_by_ch, seed_value_by_ch
|
||||||
|
|
||||||
|
|
||||||
def build_delta_channels_writes_history(
|
def build_delta_channels_writes_history(
|
||||||
*,
|
*,
|
||||||
channels: Sequence[str],
|
channels: Sequence[str],
|
||||||
chain_by_ch: Mapping[str, list[str]],
|
chained_cpid_by_ch: Mapping[str, Sequence[str]],
|
||||||
seed_val_by_ch: Mapping[str, Any],
|
seed_cpid_by_ch: Mapping[str, str | None],
|
||||||
seeded: set[str],
|
seed_value_by_ch: Mapping[str, Any],
|
||||||
stage2_rows: Sequence[tuple[str, str, str, int, str, bytes]],
|
stage2_rows: Sequence[tuple[str, str, str, int, str, bytes]],
|
||||||
serde: Any,
|
serde: Any,
|
||||||
) -> dict[str, DeltaChannelHistory]:
|
) -> dict[str, DeltaChannelHistory]:
|
||||||
"""Demux stage-2 rows per channel; produce per-channel histories.
|
"""Demux writes rows per channel; produce per-channel histories.
|
||||||
|
|
||||||
Stage-2 rows are `(checkpoint_id, channel, task_id, idx, type, value)`.
|
`stage2_rows` are `(checkpoint_id, channel, task_id, idx, type, value)`.
|
||||||
Final write order is oldest→newest globally and `(task_id, idx)` within
|
Final write order is oldest→newest globally and `(task_id, idx)` within
|
||||||
a checkpoint, matching the contract on `DeltaChannelHistory.writes`.
|
a checkpoint, matching the contract on `DeltaChannelHistory.writes`.
|
||||||
|
|
||||||
`seed` is omitted when the walk reached a true root with no snapshot
|
The seed checkpoint's own writes are replayed on top of a
|
||||||
found (channel never entered `seeded`); consumers treat absence as
|
`_DeltaSnapshot` seed (the snapshot is the value *prior* to its own
|
||||||
"start empty".
|
writes), but skipped for a migrated plain-value seed (a legacy
|
||||||
|
non-delta blob already incorporates those writes). `seed` is omitted
|
||||||
|
when no seed was located (implicit empty baseline); consumers treat
|
||||||
|
absence as "start empty".
|
||||||
"""
|
"""
|
||||||
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
|
writes_by_ch_by_cid: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
|
||||||
ch: {} for ch in channels
|
ch: {} for ch in channels
|
||||||
@@ -156,17 +236,26 @@ def build_delta_channels_writes_history(
|
|||||||
|
|
||||||
result: dict[str, DeltaChannelHistory] = {}
|
result: dict[str, DeltaChannelHistory] = {}
|
||||||
for ch in channels:
|
for ch in channels:
|
||||||
chain_cids = chain_by_ch.get(ch, [])
|
entry: DeltaChannelHistory = {"writes": []}
|
||||||
|
|
||||||
|
skip_seed_checkpoint_writes = False
|
||||||
|
if ch in seed_value_by_ch:
|
||||||
|
seed_value = seed_value_by_ch[ch]
|
||||||
|
entry["seed"] = seed_value
|
||||||
|
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
|
||||||
|
|
||||||
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
cid_writes = writes_by_ch_by_cid.get(ch, {})
|
||||||
collected: list[PendingWrite] = []
|
if cid_writes:
|
||||||
# Chain is newest-first; iterate oldest-first for the public order.
|
collected: list[PendingWrite] = []
|
||||||
for cid in reversed(chain_cids):
|
seed_cpid = seed_cpid_by_ch.get(ch)
|
||||||
for type_tag, value_blob, task_id, _idx in cid_writes.get(cid, []):
|
# Chain is newest→oldest; replay oldest→newest.
|
||||||
collected.append(
|
for cid in reversed(list(chained_cpid_by_ch.get(ch, []))):
|
||||||
(task_id, ch, serde.loads_typed((type_tag, value_blob)))
|
if skip_seed_checkpoint_writes and cid == seed_cpid:
|
||||||
)
|
continue
|
||||||
entry: DeltaChannelHistory = {"writes": collected}
|
for type_tag, value_blob, task_id, _idx in cid_writes.get(cid, []):
|
||||||
if ch in seeded:
|
collected.append(
|
||||||
entry["seed"] = seed_val_by_ch[ch]
|
(task_id, ch, serde.loads_typed((type_tag, value_blob)))
|
||||||
|
)
|
||||||
|
entry["writes"] = collected
|
||||||
result[ch] = entry
|
result[ch] = entry
|
||||||
return result
|
return result
|
||||||
|
|||||||
@@ -25,10 +25,12 @@ from langgraph.checkpoint.base import (
|
|||||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||||
|
|
||||||
from langgraph.checkpoint.sqlite._delta import (
|
from langgraph.checkpoint.sqlite._delta import (
|
||||||
DELTA_STAGE1_SQL,
|
DELTA_WALK_SQL,
|
||||||
build_delta_channels_writes_history,
|
build_delta_channels_writes_history,
|
||||||
build_delta_stage2_sql,
|
build_delta_writes_fetch_sql,
|
||||||
step_walk_with_row,
|
parse_supersteps_since_last_snapshot_by_channel,
|
||||||
|
resolve_delta_chains,
|
||||||
|
step_walk_supersteps,
|
||||||
)
|
)
|
||||||
from langgraph.checkpoint.sqlite.utils import search_where
|
from langgraph.checkpoint.sqlite.utils import search_where
|
||||||
|
|
||||||
@@ -625,61 +627,98 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|||||||
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
|
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
|
||||||
|
|
||||||
See `SqliteSaver.get_delta_channel_history` for design notes; this
|
See `SqliteSaver.get_delta_channel_history` for design notes; this
|
||||||
is the async equivalent using `aiosqlite` cursors. Stage 1 pages
|
is the async equivalent using `aiosqlite` cursors. The WALK streams
|
||||||
the parent chain newest-first and Python-deserializes each
|
the parent chain newest-first, bounded by the target's
|
||||||
checkpoint blob to find per-channel snapshots; stage 2 fetches
|
`counters_since_delta_snapshot` supersteps, and deserializes only
|
||||||
only the relevant writes via per-channel UNION ALL.
|
seed checkpoints; FETCH pulls the relevant writes via per-channel
|
||||||
|
UNION ALL.
|
||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
channels = list(channels)
|
channels = list(channels)
|
||||||
await self.setup()
|
await self.setup()
|
||||||
thread_id = str(config["configurable"]["thread_id"])
|
thread_id = str(config["configurable"]["thread_id"])
|
||||||
|
if not thread_id:
|
||||||
|
raise ValueError("empty thread ID")
|
||||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
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}
|
# Resolve the target checkpoint id + its metadata (for supersteps).
|
||||||
seed_val_by_ch: dict[str, Any] = {}
|
checkpoint_id = get_checkpoint_id(config)
|
||||||
|
async with self.lock, self.conn.cursor() as cur:
|
||||||
|
if checkpoint_id is None:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = ? AND checkpoint_ns = ? "
|
||||||
|
"ORDER BY checkpoint_id DESC LIMIT 1",
|
||||||
|
(thread_id, checkpoint_ns),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT checkpoint_id, metadata FROM checkpoints "
|
||||||
|
"WHERE thread_id = ? AND checkpoint_ns = ? "
|
||||||
|
"AND checkpoint_id = ?",
|
||||||
|
(thread_id, checkpoint_ns, checkpoint_id),
|
||||||
|
)
|
||||||
|
target_row = await cur.fetchone()
|
||||||
|
if target_row is None:
|
||||||
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
target_id = str(target_row[0])
|
||||||
|
metadata = json.loads(target_row[1]) if target_row[1] is not None else {}
|
||||||
|
supersteps_by_ch = parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
metadata, channels
|
||||||
|
)
|
||||||
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
needed_depths = set(supersteps_by_ch.values())
|
||||||
|
|
||||||
|
shared_cpid_chain: list[str] = []
|
||||||
walk_state: dict[str, Any] = {}
|
walk_state: dict[str, Any] = {}
|
||||||
seeded: set[str] = set()
|
seed_values_by_depth: dict[int, dict[str, Any]] = {}
|
||||||
|
|
||||||
async with self.lock, self.conn.cursor() as cur:
|
async with self.lock, self.conn.cursor() as cur:
|
||||||
await cur.execute(
|
if max_supersteps > 0:
|
||||||
DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)
|
await cur.execute(DELTA_WALK_SQL, (thread_id, checkpoint_ns, target_id))
|
||||||
)
|
async for row in cur:
|
||||||
async for row in cur:
|
cid, parent_cid, type_tag, blob = row
|
||||||
cid, parent_cid, type_tag, blob = row
|
if step_walk_supersteps(
|
||||||
if step_walk_with_row(
|
cid=cid,
|
||||||
cid=cid,
|
parent_cid=parent_cid,
|
||||||
parent_cid=parent_cid,
|
type_tag=type_tag,
|
||||||
type_tag=type_tag,
|
blob=blob,
|
||||||
blob=blob,
|
target_id=target_id,
|
||||||
target_id=checkpoint_id,
|
serde=self.serde,
|
||||||
serde=self.serde,
|
shared_cpid_chain=shared_cpid_chain,
|
||||||
chain_by_ch=chain_by_ch,
|
walk_state=walk_state,
|
||||||
seed_val_by_ch=seed_val_by_ch,
|
max_supersteps=max_supersteps,
|
||||||
walk_state=walk_state,
|
needed_depths=needed_depths,
|
||||||
seeded=seeded,
|
seed_values_by_depth=seed_values_by_depth,
|
||||||
channels=channels,
|
channels=channels,
|
||||||
):
|
):
|
||||||
break
|
break
|
||||||
|
|
||||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
(
|
||||||
stage2_sql = build_delta_stage2_sql(
|
chained_cpid_by_ch,
|
||||||
chain_lens=[len(chain_by_ch[ch]) for ch in channels_with_chain],
|
seed_cpid_by_ch,
|
||||||
|
seed_value_by_ch,
|
||||||
|
) = resolve_delta_chains(
|
||||||
|
channels=channels,
|
||||||
|
supersteps_by_ch=supersteps_by_ch,
|
||||||
|
shared_cpid_chain=shared_cpid_chain,
|
||||||
|
seed_values_by_depth=seed_values_by_depth,
|
||||||
|
has_reached_root=bool(walk_state.get("reached_root")),
|
||||||
|
thread_id=thread_id,
|
||||||
)
|
)
|
||||||
if stage2_sql:
|
|
||||||
stage2_params: list[Any] = []
|
channels_with_chain = [ch for ch in channels if chained_cpid_by_ch[ch]]
|
||||||
|
fetch_sql = build_delta_writes_fetch_sql(
|
||||||
|
chain_lens=[len(chained_cpid_by_ch[ch]) for ch in channels_with_chain],
|
||||||
|
)
|
||||||
|
if fetch_sql:
|
||||||
|
fetch_params: list[Any] = []
|
||||||
for ch in channels_with_chain:
|
for ch in channels_with_chain:
|
||||||
stage2_params.extend(
|
fetch_params.extend(
|
||||||
[thread_id, checkpoint_ns, ch, *chain_by_ch[ch]]
|
[thread_id, checkpoint_ns, ch, *chained_cpid_by_ch[ch]]
|
||||||
)
|
)
|
||||||
await cur.execute(stage2_sql, stage2_params)
|
await cur.execute(fetch_sql, fetch_params)
|
||||||
stage2_rows = cast(
|
stage2_rows = cast(
|
||||||
"list[tuple[str, str, str, int, str, bytes]]",
|
"list[tuple[str, str, str, int, str, bytes]]",
|
||||||
await cur.fetchall(),
|
await cur.fetchall(),
|
||||||
@@ -689,9 +728,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|||||||
|
|
||||||
return build_delta_channels_writes_history(
|
return build_delta_channels_writes_history(
|
||||||
channels=channels,
|
channels=channels,
|
||||||
chain_by_ch=chain_by_ch,
|
chained_cpid_by_ch=chained_cpid_by_ch,
|
||||||
seed_val_by_ch=seed_val_by_ch,
|
seed_cpid_by_ch=seed_cpid_by_ch,
|
||||||
seeded=seeded,
|
seed_value_by_ch=seed_value_by_ch,
|
||||||
stage2_rows=stage2_rows,
|
stage2_rows=stage2_rows,
|
||||||
serde=self.serde,
|
serde=self.serde,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -34,6 +34,37 @@ PendingWrite = tuple[str, str, Any]
|
|||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
metadata: Mapping[str, Any],
|
||||||
|
channels: Sequence[str],
|
||||||
|
) -> dict[str, int]:
|
||||||
|
"""Per-channel supersteps-since-last-snapshot, parsed from metadata.
|
||||||
|
|
||||||
|
Reads `metadata.counters_since_delta_snapshot[ch]`, a `(updates,
|
||||||
|
supersteps)` pair, and returns `{ch: supersteps}` for channels whose
|
||||||
|
supersteps count is positive. Channels with no counter entry (just
|
||||||
|
snapshotted, or never written) are omitted — their seed, if any, is the
|
||||||
|
target checkpoint itself and needs no ancestor walk.
|
||||||
|
|
||||||
|
Used to drive the ancestor-walk depth in `get_delta_channel_history`:
|
||||||
|
a channel's seed snapshot sits exactly `supersteps` hops back along the
|
||||||
|
parent chain.
|
||||||
|
"""
|
||||||
|
counters = metadata.get("counters_since_delta_snapshot") or {}
|
||||||
|
result: dict[str, int] = {}
|
||||||
|
for ch in channels:
|
||||||
|
entry = counters.get(ch)
|
||||||
|
if not isinstance(entry, (list, tuple)) or len(entry) < 2:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
supersteps = int(entry[1])
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
continue
|
||||||
|
if supersteps > 0:
|
||||||
|
result[ch] = supersteps
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
# Marked as total=False to allow for future expansion.
|
# Marked as total=False to allow for future expansion.
|
||||||
class CheckpointMetadata(TypedDict, total=False):
|
class CheckpointMetadata(TypedDict, total=False):
|
||||||
"""Metadata associated with a checkpoint."""
|
"""Metadata associated with a checkpoint."""
|
||||||
@@ -619,34 +650,30 @@ class BaseCheckpointSaver(Generic[V]):
|
|||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
|
channels = list(channels)
|
||||||
seed_by_ch: dict[str, Any] = {}
|
|
||||||
remaining: set[str] = set(channels)
|
|
||||||
target_tuple = self.get_tuple(config)
|
target_tuple = self.get_tuple(config)
|
||||||
cursor_config: RunnableConfig | None = (
|
if target_tuple is None:
|
||||||
target_tuple.parent_config if target_tuple else None
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
target_tuple.metadata or {}, channels
|
||||||
)
|
)
|
||||||
while cursor_config is not None and remaining:
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
|
||||||
|
chain: list[CheckpointTuple] = []
|
||||||
|
has_reached_root = False
|
||||||
|
cursor_config: RunnableConfig | None = target_tuple.parent_config
|
||||||
|
while cursor_config is not None and len(chain) < max_supersteps:
|
||||||
tup = self.get_tuple(cursor_config)
|
tup = self.get_tuple(cursor_config)
|
||||||
if tup is None:
|
if tup is None:
|
||||||
break
|
break
|
||||||
if tup.pending_writes:
|
chain.append(tup)
|
||||||
for write in reversed(tup.pending_writes):
|
if tup.parent_config is None:
|
||||||
ch = write[1]
|
has_reached_root = True
|
||||||
if ch in remaining:
|
break
|
||||||
collected_by_ch[ch].append(write)
|
|
||||||
for ch in list(remaining):
|
|
||||||
if ch in tup.checkpoint["channel_values"]:
|
|
||||||
seed_by_ch[ch] = tup.checkpoint["channel_values"][ch]
|
|
||||||
remaining.discard(ch)
|
|
||||||
cursor_config = tup.parent_config
|
cursor_config = tup.parent_config
|
||||||
result: dict[str, DeltaChannelHistory] = {}
|
return self._assemble_default_delta_history(
|
||||||
for ch in channels:
|
channels, supersteps_by_ch, chain, has_reached_root, config
|
||||||
entry: DeltaChannelHistory = {"writes": list(reversed(collected_by_ch[ch]))}
|
)
|
||||||
if ch in seed_by_ch:
|
|
||||||
entry["seed"] = seed_by_ch[ch]
|
|
||||||
result[ch] = entry
|
|
||||||
return result
|
|
||||||
|
|
||||||
async def aget_delta_channel_history(
|
async def aget_delta_channel_history(
|
||||||
self, *, config: RunnableConfig, channels: Sequence[str]
|
self, *, config: RunnableConfig, channels: Sequence[str]
|
||||||
@@ -660,32 +687,106 @@ class BaseCheckpointSaver(Generic[V]):
|
|||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
|
channels = list(channels)
|
||||||
seed_by_ch: dict[str, Any] = {}
|
|
||||||
remaining: set[str] = set(channels)
|
|
||||||
target_tuple = await self.aget_tuple(config)
|
target_tuple = await self.aget_tuple(config)
|
||||||
cursor_config: RunnableConfig | None = (
|
if target_tuple is None:
|
||||||
target_tuple.parent_config if target_tuple else None
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
target_tuple.metadata or {}, channels
|
||||||
)
|
)
|
||||||
while cursor_config is not None and remaining:
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
|
||||||
|
chain: list[CheckpointTuple] = []
|
||||||
|
has_reached_root = False
|
||||||
|
cursor_config: RunnableConfig | None = target_tuple.parent_config
|
||||||
|
while cursor_config is not None and len(chain) < max_supersteps:
|
||||||
tup = await self.aget_tuple(cursor_config)
|
tup = await self.aget_tuple(cursor_config)
|
||||||
if tup is None:
|
if tup is None:
|
||||||
break
|
break
|
||||||
if tup.pending_writes:
|
chain.append(tup)
|
||||||
for write in reversed(tup.pending_writes):
|
if tup.parent_config is None:
|
||||||
ch = write[1]
|
has_reached_root = True
|
||||||
if ch in remaining:
|
break
|
||||||
collected_by_ch[ch].append(write)
|
|
||||||
for ch in list(remaining):
|
|
||||||
if ch in tup.checkpoint["channel_values"]:
|
|
||||||
seed_by_ch[ch] = tup.checkpoint["channel_values"][ch]
|
|
||||||
remaining.discard(ch)
|
|
||||||
cursor_config = tup.parent_config
|
cursor_config = tup.parent_config
|
||||||
|
return self._assemble_default_delta_history(
|
||||||
|
channels, supersteps_by_ch, chain, has_reached_root, config
|
||||||
|
)
|
||||||
|
|
||||||
|
def _assemble_default_delta_history(
|
||||||
|
self,
|
||||||
|
channels: Sequence[str],
|
||||||
|
supersteps_by_ch: Mapping[str, int],
|
||||||
|
chain: Sequence[CheckpointTuple],
|
||||||
|
has_reached_root: bool,
|
||||||
|
config: RunnableConfig,
|
||||||
|
) -> dict[str, DeltaChannelHistory]:
|
||||||
|
"""Slice the walked parent chain into per-channel histories.
|
||||||
|
|
||||||
|
For each channel the seed snapshot sits `supersteps` hops back along
|
||||||
|
`chain` (newest first): the seed checkpoint is `chain[supersteps - 1]`
|
||||||
|
and its `channel_values[ch]` is the seed value. The seed checkpoint's
|
||||||
|
own writes are replayed on top of a `_DeltaSnapshot` seed but skipped
|
||||||
|
for a migrated plain-value seed (the legacy non-delta blob already
|
||||||
|
incorporates them).
|
||||||
|
|
||||||
|
When the chain is shorter than `supersteps` but the root was reached,
|
||||||
|
the persisted chain is "compressed" relative to the logical superstep
|
||||||
|
count (`durability="exit"`, or a thread that never snapshotted). The
|
||||||
|
seed candidate is then the oldest checkpoint (`chain[-1]`); if its
|
||||||
|
`channel_values[ch]` is absent the channel has no seed and replays
|
||||||
|
the full chain on an empty baseline.
|
||||||
|
"""
|
||||||
|
# Imported lazily to avoid a hard checkpoint→serde-types coupling at
|
||||||
|
# module import; only the delta surface needs the runtime check.
|
||||||
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||||
|
|
||||||
|
thread_id = config["configurable"].get("thread_id")
|
||||||
result: dict[str, DeltaChannelHistory] = {}
|
result: dict[str, DeltaChannelHistory] = {}
|
||||||
for ch in channels:
|
for ch in channels:
|
||||||
entry: DeltaChannelHistory = {"writes": list(reversed(collected_by_ch[ch]))}
|
entry: DeltaChannelHistory = {"writes": []}
|
||||||
if ch in seed_by_ch:
|
bound = supersteps_by_ch.get(ch, 0)
|
||||||
entry["seed"] = seed_by_ch[ch]
|
if bound <= 0:
|
||||||
|
result[ch] = entry
|
||||||
|
continue
|
||||||
|
if len(chain) >= bound:
|
||||||
|
seed_depth = bound
|
||||||
|
elif has_reached_root:
|
||||||
|
seed_depth = len(chain)
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"cannot find seed snapshot for delta channel "
|
||||||
|
"(thread_id=%s, channel=%s)",
|
||||||
|
thread_id,
|
||||||
|
ch,
|
||||||
|
)
|
||||||
|
result[ch] = entry
|
||||||
|
continue
|
||||||
|
if seed_depth <= 0:
|
||||||
|
result[ch] = entry
|
||||||
|
continue
|
||||||
|
chain_slice: Sequence[CheckpointTuple] = chain[:seed_depth]
|
||||||
|
seed_tuple: CheckpointTuple | None = chain[seed_depth - 1]
|
||||||
|
|
||||||
|
skip_seed_checkpoint_writes = False
|
||||||
|
channel_values = seed_tuple.checkpoint["channel_values"]
|
||||||
|
if ch in channel_values:
|
||||||
|
seed_value = channel_values[ch]
|
||||||
|
entry["seed"] = seed_value
|
||||||
|
skip_seed_checkpoint_writes = not isinstance(seed_value, _DeltaSnapshot)
|
||||||
|
else:
|
||||||
|
# No stored value at the oldest checkpoint → empty baseline,
|
||||||
|
# so the chain's own writes are all replayed from empty.
|
||||||
|
seed_tuple = None
|
||||||
|
|
||||||
|
collected: list[PendingWrite] = []
|
||||||
|
# Chain is newest→oldest; replay oldest→newest.
|
||||||
|
for tup in reversed(list(chain_slice)):
|
||||||
|
if skip_seed_checkpoint_writes and tup is seed_tuple:
|
||||||
|
continue
|
||||||
|
for write in tup.pending_writes or []:
|
||||||
|
if write[1] == ch:
|
||||||
|
collected.append(write)
|
||||||
|
entry["writes"] = collected
|
||||||
result[ch] = entry
|
result[ch] = entry
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from langgraph.checkpoint.base import (
|
|||||||
DeltaChannelHistory,
|
DeltaChannelHistory,
|
||||||
PendingWrite,
|
PendingWrite,
|
||||||
SerializerProtocol,
|
SerializerProtocol,
|
||||||
|
_parse_supersteps_since_last_snapshot_by_channel,
|
||||||
get_checkpoint_id,
|
get_checkpoint_id,
|
||||||
get_checkpoint_metadata,
|
get_checkpoint_metadata,
|
||||||
)
|
)
|
||||||
@@ -144,14 +145,20 @@ class InMemorySaver(
|
|||||||
) -> Mapping[str, DeltaChannelHistory]:
|
) -> Mapping[str, DeltaChannelHistory]:
|
||||||
"""Override: walk the parent chain ONCE for all requested channels.
|
"""Override: walk the parent chain ONCE for all requested channels.
|
||||||
|
|
||||||
Each channel terminates independently at the nearest ancestor
|
Walk depth is driven by the target checkpoint's
|
||||||
whose stored blob is non-empty. Other channels keep walking until
|
`counters_since_delta_snapshot[ch]` supersteps counter: each
|
||||||
they find their own terminator or hit the root.
|
channel's seed snapshot sits exactly `supersteps` hops back along
|
||||||
|
the parent chain. This locates seeds reliably even when they aren't
|
||||||
|
a `_DeltaSnapshot` sentinel — e.g. a legacy plain-value blob from a
|
||||||
|
thread migrated off a non-delta channel.
|
||||||
|
|
||||||
Pre-delta plain-value blobs subsume their ancestor's pending
|
Pre-delta plain-value seeds subsume their own checkpoint's pending
|
||||||
writes (the value already includes them); `_DeltaSnapshot` blobs
|
writes (the value already includes them), so those are skipped;
|
||||||
do not (snapshot is the value AT that ancestor, prior to its own
|
`_DeltaSnapshot` seeds do not (the snapshot is the value AT that
|
||||||
pending writes that produce the child).
|
ancestor, prior to its own pending writes), so they are replayed.
|
||||||
|
When the chain reaches the root short of `supersteps`, the snapshot
|
||||||
|
was never persisted (implicit empty baseline) — replay the full
|
||||||
|
chain with no seed.
|
||||||
"""
|
"""
|
||||||
if not channels:
|
if not channels:
|
||||||
return {}
|
return {}
|
||||||
@@ -159,72 +166,98 @@ class InMemorySaver(
|
|||||||
# module import; only this override needs the runtime check.
|
# module import; only this override needs the runtime check.
|
||||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||||
|
|
||||||
|
channels = list(channels)
|
||||||
thread_id = config["configurable"]["thread_id"]
|
thread_id = config["configurable"]["thread_id"]
|
||||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||||
checkpoint_id = config["configurable"].get("checkpoint_id", "")
|
|
||||||
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
|
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
|
||||||
|
|
||||||
|
# Resolve the target checkpoint id + its metadata (for supersteps).
|
||||||
|
checkpoint_id = get_checkpoint_id(config)
|
||||||
|
if checkpoint_id is None:
|
||||||
|
checkpoint_id = next(reversed(ns_storage), None)
|
||||||
|
target_entry = ns_storage.get(checkpoint_id) if checkpoint_id else None
|
||||||
|
if target_entry is None:
|
||||||
|
return {ch: {"writes": []} for ch in channels}
|
||||||
|
metadata = self.serde.loads_typed(target_entry[1])
|
||||||
|
supersteps_by_ch = _parse_supersteps_since_last_snapshot_by_channel(
|
||||||
|
metadata or {}, channels
|
||||||
|
)
|
||||||
|
max_supersteps = max(supersteps_by_ch.values(), default=0)
|
||||||
|
|
||||||
|
# Walk the parent chain (newest first) up to the deepest supersteps.
|
||||||
chain: list[str] = []
|
chain: list[str] = []
|
||||||
target_entry = ns_storage.get(checkpoint_id)
|
has_reached_root = False
|
||||||
current: str | None = target_entry[2] if target_entry is not None else None
|
current: str | None = target_entry[2]
|
||||||
while current is not None:
|
while current is not None and len(chain) < max_supersteps:
|
||||||
entry = ns_storage.get(current)
|
entry = ns_storage.get(current)
|
||||||
if entry is None:
|
if entry is None:
|
||||||
break
|
break
|
||||||
chain.append(current)
|
chain.append(current)
|
||||||
_, _, parent = entry
|
parent = entry[2]
|
||||||
current = parent
|
if parent is None:
|
||||||
|
has_reached_root = True
|
||||||
collected_by_ch: dict[str, list[PendingWrite]] = {c: [] for c in channels}
|
|
||||||
seed_by_ch: dict[str, Any] = {}
|
|
||||||
remaining: set[str] = set(channels)
|
|
||||||
|
|
||||||
for cp_id in chain:
|
|
||||||
if not remaining:
|
|
||||||
break
|
break
|
||||||
entry = ns_storage.get(cp_id)
|
current = parent
|
||||||
ckpt = self.serde.loads_typed(entry[0]) if entry is not None else None
|
|
||||||
|
|
||||||
terminated_here: set[str] = set()
|
|
||||||
blob_value_by_ch: dict[str, Any] = {}
|
|
||||||
if ckpt is not None:
|
|
||||||
versions = ckpt.get("channel_versions", {})
|
|
||||||
for ch in remaining:
|
|
||||||
ver = versions.get(ch)
|
|
||||||
if ver is None:
|
|
||||||
continue
|
|
||||||
blob_entry = self.blobs.get((thread_id, checkpoint_ns, ch, ver))
|
|
||||||
if blob_entry is None or blob_entry[0] == "empty":
|
|
||||||
continue
|
|
||||||
blob_value_by_ch[ch] = self.serde.loads_typed(blob_entry)
|
|
||||||
terminated_here.add(ch)
|
|
||||||
|
|
||||||
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
|
|
||||||
for (_task_id, _idx), (tid, ch, serialized, _) in sorted(
|
|
||||||
step_writes.items(), reverse=True
|
|
||||||
):
|
|
||||||
if ch not in remaining:
|
|
||||||
continue
|
|
||||||
blob_value = blob_value_by_ch.get(ch)
|
|
||||||
if blob_value is not None and not isinstance(
|
|
||||||
blob_value, _DeltaSnapshot
|
|
||||||
):
|
|
||||||
continue
|
|
||||||
collected_by_ch[ch].append(
|
|
||||||
(tid, ch, self.serde.loads_typed(serialized))
|
|
||||||
)
|
|
||||||
|
|
||||||
for ch in terminated_here:
|
|
||||||
seed_by_ch[ch] = blob_value_by_ch[ch]
|
|
||||||
remaining.discard(ch)
|
|
||||||
|
|
||||||
result: dict[str, DeltaChannelHistory] = {}
|
result: dict[str, DeltaChannelHistory] = {}
|
||||||
for ch in channels:
|
for ch in channels:
|
||||||
entry_h: DeltaChannelHistory = {
|
entry_h: DeltaChannelHistory = {"writes": []}
|
||||||
"writes": list(reversed(collected_by_ch[ch]))
|
bound = supersteps_by_ch.get(ch, 0)
|
||||||
}
|
if bound <= 0:
|
||||||
if ch in seed_by_ch:
|
result[ch] = entry_h
|
||||||
entry_h["seed"] = seed_by_ch[ch]
|
continue
|
||||||
|
if len(chain) >= bound:
|
||||||
|
seed_depth = bound
|
||||||
|
elif has_reached_root:
|
||||||
|
seed_depth = len(chain)
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
"cannot find seed snapshot for delta channel "
|
||||||
|
"(thread_id=%s, channel=%s)",
|
||||||
|
thread_id,
|
||||||
|
ch,
|
||||||
|
)
|
||||||
|
result[ch] = entry_h
|
||||||
|
continue
|
||||||
|
if seed_depth <= 0:
|
||||||
|
result[ch] = entry_h
|
||||||
|
continue
|
||||||
|
chain_slice: list[str] = chain[:seed_depth]
|
||||||
|
seed_cpid: str | None = chain[seed_depth - 1]
|
||||||
|
|
||||||
|
skip_seed_checkpoint_writes = False
|
||||||
|
seed_entry = ns_storage.get(seed_cpid)
|
||||||
|
if seed_entry is not None:
|
||||||
|
ckpt = self.serde.loads_typed(seed_entry[0])
|
||||||
|
ver = (ckpt.get("channel_versions") or {}).get(ch)
|
||||||
|
blob_entry = (
|
||||||
|
self.blobs.get((thread_id, checkpoint_ns, ch, ver))
|
||||||
|
if ver is not None
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
if blob_entry is not None and blob_entry[0] != "empty":
|
||||||
|
seed_value = self.serde.loads_typed(blob_entry)
|
||||||
|
entry_h["seed"] = seed_value
|
||||||
|
skip_seed_checkpoint_writes = not isinstance(
|
||||||
|
seed_value, _DeltaSnapshot
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# No stored value at the oldest checkpoint → empty
|
||||||
|
# baseline; the chain's own writes replay from empty.
|
||||||
|
seed_cpid = None
|
||||||
|
|
||||||
|
collected: list[PendingWrite] = []
|
||||||
|
# Chain is newest→oldest; replay oldest→newest.
|
||||||
|
for cp_id in reversed(chain_slice):
|
||||||
|
if skip_seed_checkpoint_writes and cp_id == seed_cpid:
|
||||||
|
continue
|
||||||
|
step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {})
|
||||||
|
for (_task_id, _idx), (tid, w_ch, serialized, _) in sorted(
|
||||||
|
step_writes.items()
|
||||||
|
):
|
||||||
|
if w_ch == ch:
|
||||||
|
collected.append((tid, ch, self.serde.loads_typed(serialized)))
|
||||||
|
entry_h["writes"] = collected
|
||||||
result[ch] = entry_h
|
result[ch] = entry_h
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
@@ -351,9 +351,12 @@ class TestInMemorySaverDeltaChannel:
|
|||||||
cp1["id"] = "cp1"
|
cp1["id"] = "cp1"
|
||||||
cp2 = empty_checkpoint()
|
cp2 = empty_checkpoint()
|
||||||
cp2["id"] = "cp2"
|
cp2["id"] = "cp2"
|
||||||
|
# Target (cp2) carries the supersteps counter; no snapshot was ever
|
||||||
|
# taken, so the walk runs back to the root → no seed, empty baseline.
|
||||||
|
cp2_md = {"counters_since_delta_snapshot": {channel: [1, 2]}}
|
||||||
saver.storage[thread_id][ns] = {
|
saver.storage[thread_id][ns] = {
|
||||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed(cp2_md), "cp1"),
|
||||||
}
|
}
|
||||||
# Writes stored at cp1 produced the cp1 snapshot; part of history.
|
# Writes stored at cp1 produced the cp1 snapshot; part of history.
|
||||||
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = (
|
||||||
@@ -456,10 +459,15 @@ class TestBaseFallbackGetChannelWrites:
|
|||||||
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
|
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
|
||||||
cp2 = empty_checkpoint()
|
cp2 = empty_checkpoint()
|
||||||
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
|
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
|
||||||
|
# The target (cp2) carries `counters_since_delta_snapshot` — the
|
||||||
|
# supersteps-since-last-snapshot driving the ancestor-walk depth. No
|
||||||
|
# snapshot was ever taken, so the count runs back past the oldest
|
||||||
|
# persisted checkpoint (root) → implicit empty baseline, no seed.
|
||||||
|
cp2_md = {"counters_since_delta_snapshot": {channel: [2, 3]}}
|
||||||
saver.storage[thread_id][ns] = {
|
saver.storage[thread_id][ns] = {
|
||||||
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
|
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
|
||||||
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
|
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
|
||||||
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]),
|
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed(cp2_md), cp1["id"]),
|
||||||
}
|
}
|
||||||
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
|
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
|
||||||
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
|
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
|
||||||
@@ -607,10 +615,14 @@ class TestPreDeltaBlobTerminator:
|
|||||||
cp3["id"] = "cp3"
|
cp3["id"] = "cp3"
|
||||||
cp3["channel_versions"][channel] = v3
|
cp3["channel_versions"][channel] = v3
|
||||||
|
|
||||||
|
# cp1 is the pre-delta (migration) seed; the channel became a
|
||||||
|
# DeltaChannel at cp2, so its supersteps counter starts there: cp2 -> 1,
|
||||||
|
# cp3 -> 2. The target (cp3) walks 2 hops back to the cp1 seed blob.
|
||||||
|
cp3_md = {"counters_since_delta_snapshot": {channel: [2, 2]}}
|
||||||
saver.storage[thread_id][ns] = {
|
saver.storage[thread_id][ns] = {
|
||||||
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
|
||||||
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
"cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"),
|
||||||
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"),
|
"cp3": (serde.dumps_typed(cp3), serde.dumps_typed(cp3_md), "cp2"),
|
||||||
}
|
}
|
||||||
# Write under cp1 would be from the pre-delta era and MUST be ignored
|
# Write under cp1 would be from the pre-delta era and MUST be ignored
|
||||||
# (the blob already captures it). We add one and assert it is not
|
# (the blob already captures it). We add one and assert it is not
|
||||||
|
|||||||
@@ -70,6 +70,42 @@ def delta_channels_to_snapshot(
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def advance_delta_counters(
|
||||||
|
channels: Mapping[str, BaseChannel],
|
||||||
|
prev_metadata: Mapping[str, Any] | None,
|
||||||
|
updated_channels: set[str] | None,
|
||||||
|
) -> tuple[dict[str, tuple[int, int]], set[str]]:
|
||||||
|
"""Advance per-delta-channel `(updates, supersteps)` counters one step.
|
||||||
|
|
||||||
|
Mirrors the counter bookkeeping in `PregelLoop._put_checkpoint`: every
|
||||||
|
superstep bumps `supersteps` for all delta channels and `updates` for
|
||||||
|
those written this step, then channels that hit their snapshot cadence
|
||||||
|
are reset to `(0, 0)`.
|
||||||
|
|
||||||
|
Used by `update_state` (which creates a delta checkpoint outside the
|
||||||
|
main loop) so the written checkpoint carries a correct
|
||||||
|
`counters_since_delta_snapshot` — without it, the supersteps-based
|
||||||
|
`get_delta_channel_history` walk can't locate the seed. Returns
|
||||||
|
`(new_counters, channels_to_snapshot)`; callers drop `(0, 0)` entries
|
||||||
|
before writing the metadata field.
|
||||||
|
"""
|
||||||
|
prev = dict((prev_metadata or {}).get("counters_since_delta_snapshot") or {})
|
||||||
|
updated = updated_channels or set()
|
||||||
|
new_counters: dict[str, tuple[int, int]] = {}
|
||||||
|
for name, ch in channels.items():
|
||||||
|
if not isinstance(ch, DeltaChannel):
|
||||||
|
continue
|
||||||
|
u, s = prev.get(name, (0, 0))
|
||||||
|
s += 1
|
||||||
|
if name in updated:
|
||||||
|
u += 1
|
||||||
|
new_counters[name] = (u, s)
|
||||||
|
channels_to_snapshot = delta_channels_to_snapshot(channels, new_counters)
|
||||||
|
for k in channels_to_snapshot:
|
||||||
|
new_counters[k] = (0, 0)
|
||||||
|
return new_counters, channels_to_snapshot
|
||||||
|
|
||||||
|
|
||||||
def create_checkpoint(
|
def create_checkpoint(
|
||||||
checkpoint: Checkpoint,
|
checkpoint: Checkpoint,
|
||||||
channels: Mapping[str, BaseChannel] | None,
|
channels: Mapping[str, BaseChannel] | None,
|
||||||
|
|||||||
@@ -130,6 +130,7 @@ from langgraph.pregel._algo import (
|
|||||||
from langgraph.pregel._call import identifier
|
from langgraph.pregel._call import identifier
|
||||||
from langgraph.pregel._checkpoint import (
|
from langgraph.pregel._checkpoint import (
|
||||||
achannels_from_checkpoint,
|
achannels_from_checkpoint,
|
||||||
|
advance_delta_counters,
|
||||||
channels_from_checkpoint,
|
channels_from_checkpoint,
|
||||||
copy_checkpoint,
|
copy_checkpoint,
|
||||||
create_checkpoint,
|
create_checkpoint,
|
||||||
@@ -2032,15 +2033,40 @@ class Pregel(
|
|||||||
checkpointer.get_next_version,
|
checkpointer.get_next_version,
|
||||||
self.trigger_to_nodes,
|
self.trigger_to_nodes,
|
||||||
)
|
)
|
||||||
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
|
# Advance delta-channel snapshot counters for this manual super-
|
||||||
|
# step so the written checkpoint carries a correct
|
||||||
|
# `counters_since_delta_snapshot` (mirrors the main loop). Without
|
||||||
|
# it, the supersteps-based delta history walk can't find the seed.
|
||||||
|
updated_channel_names = {
|
||||||
|
c for task in run_tasks for c, _ in task.writes if c != PUSH
|
||||||
|
}
|
||||||
|
new_counters, channels_to_snapshot = advance_delta_counters(
|
||||||
|
channels,
|
||||||
|
saved.metadata if saved else None,
|
||||||
|
updated_channel_names,
|
||||||
|
)
|
||||||
|
# Keep `updated_channels` at its default (None) — matching the
|
||||||
|
# historical checkpoint shape here — so resume/trigger semantics
|
||||||
|
# (e.g. deferred nodes) are unaffected. `get_next_version` /
|
||||||
|
# `channels_to_snapshot` are no-ops unless a delta channel snapshots.
|
||||||
|
checkpoint = create_checkpoint(
|
||||||
|
checkpoint,
|
||||||
|
channels,
|
||||||
|
step + 1,
|
||||||
|
get_next_version=checkpointer.get_next_version,
|
||||||
|
channels_to_snapshot=channels_to_snapshot,
|
||||||
|
)
|
||||||
|
update_metadata: dict[str, Any] = {
|
||||||
|
"source": "update",
|
||||||
|
"step": step + 1,
|
||||||
|
"parents": saved.metadata.get("parents", {}) if saved else {},
|
||||||
|
}
|
||||||
|
if non_zero := {k: v for k, v in new_counters.items() if v != (0, 0)}:
|
||||||
|
update_metadata["counters_since_delta_snapshot"] = non_zero
|
||||||
next_config = checkpointer.put(
|
next_config = checkpointer.put(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
checkpoint,
|
checkpoint,
|
||||||
{
|
update_metadata,
|
||||||
"source": "update",
|
|
||||||
"step": step + 1,
|
|
||||||
"parents": saved.metadata.get("parents", {}) if saved else {},
|
|
||||||
},
|
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions, checkpoint["channel_versions"]
|
checkpoint_previous_versions, checkpoint["channel_versions"]
|
||||||
),
|
),
|
||||||
@@ -2500,16 +2526,41 @@ class Pregel(
|
|||||||
checkpointer.get_next_version,
|
checkpointer.get_next_version,
|
||||||
self.trigger_to_nodes,
|
self.trigger_to_nodes,
|
||||||
)
|
)
|
||||||
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
|
# Advance delta-channel snapshot counters for this manual super-
|
||||||
|
# step so the written checkpoint carries a correct
|
||||||
|
# `counters_since_delta_snapshot` (mirrors the main loop). Without
|
||||||
|
# it, the supersteps-based delta history walk can't find the seed.
|
||||||
|
updated_channel_names = {
|
||||||
|
c for task in run_tasks for c, _ in task.writes if c != PUSH
|
||||||
|
}
|
||||||
|
new_counters, channels_to_snapshot = advance_delta_counters(
|
||||||
|
channels,
|
||||||
|
saved.metadata if saved else None,
|
||||||
|
updated_channel_names,
|
||||||
|
)
|
||||||
|
# Keep `updated_channels` at its default (None) — matching the
|
||||||
|
# historical checkpoint shape here — so resume/trigger semantics
|
||||||
|
# (e.g. deferred nodes) are unaffected. `get_next_version` /
|
||||||
|
# `channels_to_snapshot` are no-ops unless a delta channel snapshots.
|
||||||
|
checkpoint = create_checkpoint(
|
||||||
|
checkpoint,
|
||||||
|
channels,
|
||||||
|
step + 1,
|
||||||
|
get_next_version=checkpointer.get_next_version,
|
||||||
|
channels_to_snapshot=channels_to_snapshot,
|
||||||
|
)
|
||||||
|
update_metadata: dict[str, Any] = {
|
||||||
|
"source": "update",
|
||||||
|
"step": step + 1,
|
||||||
|
"parents": saved.metadata.get("parents", {}) if saved else {},
|
||||||
|
}
|
||||||
|
if non_zero := {k: v for k, v in new_counters.items() if v != (0, 0)}:
|
||||||
|
update_metadata["counters_since_delta_snapshot"] = non_zero
|
||||||
# save checkpoint, after applying writes
|
# save checkpoint, after applying writes
|
||||||
next_config = await checkpointer.aput(
|
next_config = await checkpointer.aput(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
checkpoint,
|
checkpoint,
|
||||||
{
|
update_metadata,
|
||||||
"source": "update",
|
|
||||||
"step": step + 1,
|
|
||||||
"parents": saved.metadata.get("parents", {}) if saved else {},
|
|
||||||
},
|
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions, checkpoint["channel_versions"]
|
checkpoint_previous_versions, checkpoint["channel_versions"]
|
||||||
),
|
),
|
||||||
|
|||||||
Reference in New Issue
Block a user