diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 6ca2448d0..c939bb2e5 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -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( diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py index 1fe617ff7..46f45cb47 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/_delta.py @@ -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)) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 368428c68..f506ea031 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -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 diff --git a/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py b/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py new file mode 100644 index 000000000..b7718d801 --- /dev/null +++ b/libs/checkpoint-sqlite/tests/test_delta_parent_walk.py @@ -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]