mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 18:59:42 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e857fede04 | ||
|
|
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