mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 10:49:56 +02:00
fix(checkpoint-sqlite): walk delta ancestors by parent pointer
Stage 1 of the sqlite delta history filtered `checkpoint_id <= target` and streamed `ORDER BY checkpoint_id DESC`. Both predicates encode the same extra assumption: that every child's checkpoint id sorts above its parent's. Ancestry is defined by `parent_checkpoint_id`, and nothing in the contract requires ids to be monotonic. A parent whose id sorted above its child's was dropped from the stream, so its stored value and its writes were lost with no error raised. Removing the range filter alone would not help: in DESC order that parent arrives before the target, so the walk passes it before it has started. A single-pass ordered stream cannot express this walk. Replace it with a recursive CTE anchored at the target that follows `parent_checkpoint_id`. Rows arrive in walk order, so `step_walk_with_row` keeps its existing shape, and the query now reads only true ancestors instead of every row at or below the target, which is strictly less IO than before. Following pointers can loop where a bounded id scan could not, and a loop is reachable through `put` alone rather than only by corruption: it writes with `INSERT OR REPLACE`, so re-putting an existing checkpoint id under a descendant's config repoints that checkpoint at its own descendant. The walk therefore stops on a repeated checkpoint id. sqlite yields recursive rows lazily, so abandoning the cursor ends the recursion instead of waiting on it. Postgres needs no equivalent change. It pages the whole thread and follows parent pointers already, so it returns the correct history for this scenario. Fixes #8550 Co-authored-by: lylelllll <59271327+lylelllll@users.noreply.github.com>
This commit is contained in:
co-authored by
lylelllll
parent
658541c496
commit
ecc420e7fe
@@ -538,7 +538,10 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
seeded: set[str] = set()
|
||||
|
||||
with self.cursor(transaction=False) as cur:
|
||||
cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id))
|
||||
cur.execute(
|
||||
DELTA_STAGE1_SQL,
|
||||
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
|
||||
)
|
||||
for row in cur:
|
||||
cid, parent_cid, type_tag, blob = row
|
||||
if step_walk_with_row(
|
||||
|
||||
@@ -26,16 +26,38 @@ from typing import Any
|
||||
|
||||
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite
|
||||
|
||||
# Stage 1 streams ancestors of `target_cid` newest-first. The `<=`
|
||||
# predicate keeps target itself in the stream so we can read its
|
||||
# `parent_checkpoint_id` from the first row without a separate lookup;
|
||||
# the caller skips target's own writes/seed (matches the
|
||||
# `BaseCheckpointSaver` contract).
|
||||
# Stage 1 streams target followed by its ancestors, oldest step last, by
|
||||
# following `parent_checkpoint_id` in a recursive CTE. Target is the anchor
|
||||
# row, so the caller reads its `parent_checkpoint_id` without a separate
|
||||
# lookup and skips its own writes/seed (matches the `BaseCheckpointSaver`
|
||||
# contract).
|
||||
#
|
||||
# Ancestry is defined by `parent_checkpoint_id` alone. An earlier form
|
||||
# filtered `checkpoint_id <= target` and ordered by `checkpoint_id DESC`,
|
||||
# which additionally required every child's id to sort above its parent's.
|
||||
# Nothing in the contract promises that, and a parent sorting above its
|
||||
# child was dropped from the stream entirely, silently costing that
|
||||
# parent's seed and writes. See #8550.
|
||||
#
|
||||
# A parent chain can cycle: `put` writes with `INSERT OR REPLACE`, so
|
||||
# re-putting an existing checkpoint id under a descendant's config rewrites
|
||||
# its parent. The old id-range scan read a finite row set and could not
|
||||
# loop; this one can, so `step_walk_with_row` stops on a repeated
|
||||
# checkpoint_id. sqlite yields recursive rows lazily, so abandoning the
|
||||
# cursor ends the recursion rather than waiting on it.
|
||||
DELTA_STAGE1_SQL = (
|
||||
"WITH RECURSIVE ancestors(checkpoint_id, parent_checkpoint_id, type, "
|
||||
"checkpoint) AS ("
|
||||
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
|
||||
"FROM checkpoints "
|
||||
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
|
||||
"ORDER BY checkpoint_id DESC"
|
||||
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? "
|
||||
"UNION ALL "
|
||||
"SELECT c.checkpoint_id, c.parent_checkpoint_id, c.type, c.checkpoint "
|
||||
"FROM checkpoints c JOIN ancestors a "
|
||||
"ON c.checkpoint_id = a.parent_checkpoint_id "
|
||||
"WHERE c.thread_id = ? AND c.checkpoint_ns = ?"
|
||||
") "
|
||||
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint FROM ancestors"
|
||||
)
|
||||
|
||||
|
||||
@@ -95,14 +117,20 @@ def step_walk_with_row(
|
||||
Off-path rows (different branch on the same thread) advance the
|
||||
cursor without doing any work.
|
||||
|
||||
Returns True when every requested channel is seeded — the caller
|
||||
can stop iterating and close the cursor.
|
||||
Returns True when the caller can stop iterating and close the cursor,
|
||||
either because every requested channel is seeded or because the parent
|
||||
chain revisited a checkpoint it had already walked. `put` writes with
|
||||
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
|
||||
descendant's config points it at its own descendant and makes the chain
|
||||
a loop; without this check the recursive stage-1 query would feed rows
|
||||
forever.
|
||||
"""
|
||||
if "started" not in walk_state:
|
||||
if cid == target_id:
|
||||
walk_state["started"] = True
|
||||
walk_state["cur_cid"] = parent_cid
|
||||
walk_state["active"] = {ch for ch in channels if ch not in seeded}
|
||||
walk_state["walked"] = {cid}
|
||||
# Not target yet (or target not present): keep streaming.
|
||||
return False
|
||||
active: set[str] = walk_state["active"]
|
||||
@@ -111,6 +139,10 @@ def step_walk_with_row(
|
||||
if cid != walk_state["cur_cid"]:
|
||||
# Off-path row from a sibling branch — skip without deserializing.
|
||||
return False
|
||||
walked: set[str] = walk_state["walked"]
|
||||
if cid in walked:
|
||||
return True
|
||||
walked.add(cid)
|
||||
for ch in active:
|
||||
chain_by_ch[ch].append(cid)
|
||||
ckpt = serde.loads_typed((type_tag, blob))
|
||||
|
||||
@@ -650,7 +650,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
async with self.lock, self.conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)
|
||||
DELTA_STAGE1_SQL,
|
||||
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
|
||||
)
|
||||
async for row in cur:
|
||||
cid, parent_cid, type_tag, blob = row
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""The stage-1 ancestor walk follows `parent_checkpoint_id`, not id order.
|
||||
|
||||
Ancestry is defined by the `parent_checkpoint_id` column. An earlier stage-1
|
||||
query filtered `checkpoint_id <= target` and streamed `ORDER BY checkpoint_id
|
||||
DESC`, so it also required every child's id to sort above its parent's. Real
|
||||
ids are `uuid6` and happen to satisfy that, but the contract never promised it,
|
||||
and a parent sorting above its child was dropped from the stream: its seed and
|
||||
its writes vanished with no error. See #8550.
|
||||
|
||||
Following parent pointers in a recursive CTE removes both assumptions, at the
|
||||
cost of needing a cycle guard, which the bounded id scan got for free.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
DeltaChannelHistory,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
CHANNEL = "ch"
|
||||
CONFIG: dict[str, Any] = {"configurable": {"thread_id": "t", "checkpoint_ns": ""}}
|
||||
EXPECTED: DeltaChannelHistory = {
|
||||
"writes": [("task", CHANNEL, "write-root")],
|
||||
"seed": "seed",
|
||||
}
|
||||
|
||||
|
||||
def _checkpoint(checkpoint_id: str, values: dict[str, Any]) -> Checkpoint:
|
||||
value = empty_checkpoint()
|
||||
value["id"] = checkpoint_id
|
||||
value["channel_values"] = values
|
||||
return value
|
||||
|
||||
|
||||
# `z` sorts above `a`, so the parent's id sorts above its child's. Real uuid6
|
||||
# ids never do this; nothing in the contract stops a caller supplying ids that
|
||||
# do, and clock skew between two processes writing one thread produces it.
|
||||
PARENT_ID_ORDERS = [
|
||||
pytest.param("z-older", "a-newer", id="parent_id_sorts_above_child"),
|
||||
pytest.param("a-older", "z-newer", id="parent_id_sorts_below_child"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
|
||||
def test_sync_walk_reaches_parent_whatever_the_id_order(
|
||||
root_id: str, child_id: str
|
||||
) -> None:
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = saver.put(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
|
||||
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
child = saver.put(root, _checkpoint(child_id, {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=child, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
# The unoptimised implementation is the contract; the fast path must
|
||||
# agree with it on the same rows.
|
||||
assert (
|
||||
got[CHANNEL]
|
||||
== BaseCheckpointSaver.get_delta_channel_history(
|
||||
saver, config=child, channels=[CHANNEL]
|
||||
)[CHANNEL]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
|
||||
async def test_async_walk_reaches_parent_whatever_the_id_order(
|
||||
root_id: str, child_id: str
|
||||
) -> None:
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = await saver.aput(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
|
||||
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
child = await saver.aput(root, _checkpoint(child_id, {}), {}, {})
|
||||
|
||||
got = await saver.aget_delta_channel_history(config=child, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_long_chain_walks_the_whole_thread() -> None:
|
||||
"""A migrated thread can hold its only stored value at the root.
|
||||
|
||||
The walk then legitimately runs the length of the thread, so nothing in
|
||||
the cycle check may cut it short. Ids descend as the chain grows here, so
|
||||
id order fights the walk at every step.
|
||||
"""
|
||||
steps = 40
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
parent = saver.put(
|
||||
CONFIG, _checkpoint(f"id-{steps:03d}", {CHANNEL: "seed"}), {}, {}
|
||||
)
|
||||
saver.put_writes(parent, [(CHANNEL, "write-root")], "task")
|
||||
for step in range(steps - 1, 0, -1):
|
||||
parent = saver.put(parent, _checkpoint(f"id-{step:03d}", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=parent, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_cyclic_parent_chain_terminates() -> None:
|
||||
"""A parent chain that loops must not feed the walk forever.
|
||||
|
||||
Reachable through `put` alone, no corruption needed: it writes with
|
||||
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
|
||||
descendant's config repoints that checkpoint at its own descendant. The
|
||||
old id-range scan read a finite row set and could not loop; a recursive
|
||||
parent-pointer query can, so the walk stops on a repeated id.
|
||||
|
||||
Note this one fails by hanging, not by asserting, since a regression
|
||||
means the row stream never ends. The package has no timeout plugin, so
|
||||
the CI job timeout is the backstop.
|
||||
"""
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
a = saver.put(CONFIG, _checkpoint("cid-a", {}), {}, {})
|
||||
b = saver.put(a, _checkpoint("cid-b", {}), {}, {})
|
||||
# Re-put "cid-a" with "cid-b" as its parent: a -> b -> a.
|
||||
saver.put(b, _checkpoint("cid-a", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=b, channels=[CHANNEL])
|
||||
# Nothing on the cycle stores a value, so no seed. The point of the
|
||||
# test is that the call returns at all.
|
||||
assert "seed" not in got[CHANNEL]
|
||||
Reference in New Issue
Block a user