mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-31 12:19:58 +02:00
fix(checkpoint-postgres): find plain-value seeds when walking delta history (#8535)
Fixes langchain-ai/langgraph#8534 `put` splits stored values in two: primitives stay inline in the checkpoint's `channel_values`, everything else moves to `checkpoint_blobs`, and only `_DeltaSnapshot` leaves an inline marker behind when it moves. Stage-1 seed detection tested for that marker, so a plain value — what a thread migrated from `BinaryOperatorAggregate` leaves behind — was invisible to the walk. ### Effect Migrated threads found no seed, walked to the root, and replayed every write on every read. Values still came out correct, because replaying an additive reducer from empty rebuilds the same list, which is why nothing looked wrong. What was lost is early termination — the entire point of `DeltaChannel`: <!-- linear:table-colwidths:266,266,266 --> | thread length | writes replayed, before | after | | -- | -- | -- | | 2 turns | 3 | 1 | | 6 turns | 7 | 1 | | 20 turns | 21 | 1 | Read latency is flat at \~0.6ms across all three after the change. ### Approach Stage 1 now checks both places a value can live rather than trusting the marker. It probes `checkpoint_blobs`: ```sql EXISTS (SELECT 1 FROM checkpoint_blobs b0 WHERE b0.thread_id = checkpoints.thread_id AND b0.checkpoint_ns = checkpoints.checkpoint_ns AND b0.channel = %s AND b0.version = checkpoint -> 'channel_versions' ->> %s AND b0.type <> 'empty') AS hb_0 ``` and selects the inline value alongside it, since `None`, `str`, `int`, `float` and `bool` stay in `channel_values` with no blob row: ```sql checkpoint -> 'channel_values' -> %s AS inline_0 ``` The blob predicate matches `checkpoint_blobs`' primary key `(thread_id, checkpoint_ns, channel, version)` exactly, so it is one index lookup per row per channel, bounded by the 1024-row page. I picked reading storage over the cheaper alternative — also writing the marker for plain values — because **that would not fix any thread already on disk.** Existing checkpoints have no marker and there is nowhere to add one retroactively. The seed resolves to the blob when one exists and the inline value otherwise. That ordering is also what keeps a genuine inline `true` — a `bool` channel holding `True` — distinguishable from the literal `true` marker `put` inlines for a `_DeltaSnapshot`: only the snapshot has a blob. `None` is deliberately not treated as a seed; a JSON null is indistinguishable from "nothing stored" at this layer, so the walk continues and replay from empty is correct. Params go from two to four per channel; both callers updated. The inline half came out of review on this PR — a blob-only probe would have left scalar-aggregate migrations (an integer sum, say) still replaying their full history. ### On the `type <> 'empty'` predicate Being upfront since it isn't demonstrable with a test: `put` does not currently produce `empty` rows on this path — `blob_versions` is filtered to keys present in `channel_values`, so `_dump_blobs`' empty branch is unreachable from it. I confirmed there are no `empty` rows in a populated test database. I kept it because stage 2 already applies the same check when resolving the seed blob. Without it the two stages could disagree: stage 1 terminates the walk on a row stage 2 then discards, producing no seed *and* a truncated write chain — the same failure shape this function exists to avoid. Rationale is in the docstring so the next reader doesn't have to ask. Happy to drop it if you'd rather not carry an unexercised predicate. ### Tests `libs/checkpoint-postgres/tests/test_delta_plain_value_seed.py` — blob-stored plain-value seed, `_DeltaSnapshot` seed, a version bump with nothing stored (which must not stop the walk short of an older real value), inline primitives (`int`, `str`, `float`, `None`), and inline `True` versus the snapshot marker. Each fails against the behaviour it fixes. Verified: postgres suite 269 passed on PG 15 and 16; delta-channel conformance against `AsyncPostgresSaver` went from 6 of 8 to 8 of 8, including the pre-existing `test_history_migration_plain_value_as_seed` failure this was causing; `make lint` clean. ### Not included I wanted a Postgres conformance runner alongside `checkpoint-sqlite`'s, but it needs `langgraph-checkpoint-conformance` as a dev dependency and the contributing guide asks for maintainer sign-off before adding one. The direct tests above cover the same ground without it. Worth flagging separately: **conformance effectively runs against** `InMemorySaver` **only today.** `libs/checkpoint-conformance/tests/` contains just `test_validate_memory.py`, and `checkpoint-sqlite`'s `test_conformance_delta.py` silently skips because the package isn't installed in its test environment (`importorskip`). Wiring it up for sqlite and postgres is what would have caught this bug, and langchain-ai/langgraph#8534 notes it. Sqlite is unaffected by the bug itself — it stores `channel_values` inline and inspects them directly. `langgraph-api` already resolves seeds by version rather than by marker.
This commit is contained in:
@@ -67,24 +67,12 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
"v": 4,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
"my_key": "meow",
|
||||
"node": "node"
|
||||
},
|
||||
"channel_versions": {
|
||||
"__start__": 2,
|
||||
"my_key": 3,
|
||||
"start:node": 3,
|
||||
"node": 3
|
||||
},
|
||||
"channel_values": {"my_key": "meow", "node": "node"},
|
||||
"channel_versions": {"__start__": 2, "my_key": 3, "start:node": 3, "node": 3},
|
||||
"versions_seen": {
|
||||
"__input__": {},
|
||||
"__start__": {
|
||||
"__start__": 1
|
||||
},
|
||||
"node": {
|
||||
"start:node": 2
|
||||
}
|
||||
"__start__": {"__start__": 1},
|
||||
"node": {"start:node": 2},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -108,24 +96,12 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
|
||||
"v": 4,
|
||||
"ts": "2024-07-31T20:14:19.804150+00:00",
|
||||
"id": "1ef4f797-8335-6428-8001-8a1503f9b875",
|
||||
"channel_values": {
|
||||
"my_key": "meow",
|
||||
"node": "node"
|
||||
},
|
||||
"channel_versions": {
|
||||
"__start__": 2,
|
||||
"my_key": 3,
|
||||
"start:node": 3,
|
||||
"node": 3
|
||||
},
|
||||
"channel_values": {"my_key": "meow", "node": "node"},
|
||||
"channel_versions": {"__start__": 2, "my_key": 3, "start:node": 3, "node": 3},
|
||||
"versions_seen": {
|
||||
"__input__": {},
|
||||
"__start__": {
|
||||
"__start__": 1
|
||||
},
|
||||
"node": {
|
||||
"start:node": 2
|
||||
}
|
||||
"__start__": {"__start__": 1},
|
||||
"node": {"start:node": 2},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -478,9 +478,11 @@ class PostgresSaver(BasePostgresSaver):
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
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]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
inline_by_i_by_cid: list[dict[str, Any]] = [{} for _ in channels]
|
||||
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
||||
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = None
|
||||
@@ -489,7 +491,8 @@ class PostgresSaver(BasePostgresSaver):
|
||||
while True:
|
||||
stage1_params: list[Any] = []
|
||||
for ch in channels:
|
||||
stage1_params.extend([ch, ch])
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
@@ -502,16 +505,19 @@ class PostgresSaver(BasePostgresSaver):
|
||||
channels,
|
||||
parent_of,
|
||||
ver_by_i_by_cid,
|
||||
hs_by_i_by_cid,
|
||||
hb_by_i_by_cid,
|
||||
inline_by_i_by_cid,
|
||||
)
|
||||
self._try_advance_walks(
|
||||
checkpoint_id,
|
||||
channels,
|
||||
parent_of,
|
||||
ver_by_i_by_cid,
|
||||
hs_by_i_by_cid,
|
||||
hb_by_i_by_cid,
|
||||
inline_by_i_by_cid,
|
||||
chain_by_ch,
|
||||
seed_ver_by_ch,
|
||||
seed_inline_by_ch,
|
||||
walk_cursor_by_ch,
|
||||
seeded,
|
||||
)
|
||||
@@ -546,6 +552,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
channels=channels,
|
||||
chain_by_ch=chain_by_ch,
|
||||
seed_ver_by_ch=seed_ver_by_ch,
|
||||
seed_inline_by_ch=seed_inline_by_ch,
|
||||
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
|
||||
)
|
||||
|
||||
|
||||
@@ -426,9 +426,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
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]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
inline_by_i_by_cid: list[dict[str, Any]] = [{} for _ in channels]
|
||||
chain_by_ch: dict[str, list[str]] = {ch: [] for ch in channels}
|
||||
seed_ver_by_ch: dict[str, str | None] = {ch: None for ch in channels}
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = None
|
||||
@@ -437,7 +439,8 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
while True:
|
||||
stage1_params: list[Any] = []
|
||||
for ch in channels:
|
||||
stage1_params.extend([ch, ch])
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
@@ -450,16 +453,19 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
channels,
|
||||
parent_of,
|
||||
ver_by_i_by_cid,
|
||||
hs_by_i_by_cid,
|
||||
hb_by_i_by_cid,
|
||||
inline_by_i_by_cid,
|
||||
)
|
||||
self._try_advance_walks(
|
||||
checkpoint_id,
|
||||
channels,
|
||||
parent_of,
|
||||
ver_by_i_by_cid,
|
||||
hs_by_i_by_cid,
|
||||
hb_by_i_by_cid,
|
||||
inline_by_i_by_cid,
|
||||
chain_by_ch,
|
||||
seed_ver_by_ch,
|
||||
seed_inline_by_ch,
|
||||
walk_cursor_by_ch,
|
||||
seeded,
|
||||
)
|
||||
@@ -490,6 +496,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
channels=channels,
|
||||
chain_by_ch=chain_by_ch,
|
||||
seed_ver_by_ch=seed_ver_by_ch,
|
||||
seed_inline_by_ch=seed_inline_by_ch,
|
||||
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
|
||||
)
|
||||
|
||||
|
||||
@@ -199,27 +199,68 @@ class _DeltaStage2Row(TypedDict, total=False):
|
||||
|
||||
|
||||
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
"""Build stage 1 SQL with 2K parallel JSONB key lookups.
|
||||
"""Build stage 1 SQL with K parallel version lookups + seed probes.
|
||||
|
||||
For channels=["messages", "files"] (with `paged=True`) the result is::
|
||||
|
||||
SELECT checkpoint_id, parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
||||
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_0,
|
||||
EXISTS (SELECT 1 FROM checkpoint_blobs b0
|
||||
WHERE b0.thread_id = checkpoints.thread_id
|
||||
AND b0.checkpoint_ns = checkpoints.checkpoint_ns
|
||||
AND b0.channel = %s
|
||||
AND b0.version = checkpoint -> 'channel_versions' ->> %s
|
||||
AND b0.type <> 'empty') AS hb_0,
|
||||
checkpoint -> 'channel_values' -> %s AS inline_0,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver_1,
|
||||
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_1
|
||||
EXISTS (...) AS hb_1,
|
||||
checkpoint -> 'channel_values' -> %s AS inline_1
|
||||
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
|
||||
|
||||
Channel names are passed as `%s` parameters (safe from SQL injection).
|
||||
Only the column aliases `ver_i` / `hs_i` are interpolated into the
|
||||
SQL string (i is bounded by len(channels) and uses safe identifiers).
|
||||
A stored value for a channel lives in one of two places, because `put`
|
||||
splits them:
|
||||
|
||||
Caller must extend params with `[ch_0, ch_0, ch_1, ch_1, ...,
|
||||
thread_id, ns, cursor, cursor, page_size]` when `paged=True`.
|
||||
* **blob** — non-primitive values (and `_DeltaSnapshot`) are moved to
|
||||
`checkpoint_blobs`. `hb_i` ("has blob") probes for one. The probe hits
|
||||
that table's primary key `(thread_id, checkpoint_ns, channel, version)`
|
||||
exactly, so it is an index lookup per row per channel.
|
||||
* **inline** — `None`, `str`, `int`, `float` and `bool` stay in the
|
||||
checkpoint's own `channel_values` and get no blob row at all. `inline_i`
|
||||
returns that value.
|
||||
|
||||
Testing only for a key in `channel_values` (the previous approach) missed
|
||||
blob-stored plain values, since `put` leaves an inline marker there for
|
||||
`_DeltaSnapshot` but not for a plain value — which is what a thread
|
||||
migrated from a pre-delta channel type leaves behind. Probing only the
|
||||
blobs table would conversely miss inline primitives. Both are needed, and
|
||||
the caller treats "either present" as the seed.
|
||||
|
||||
`hb_i` also disambiguates the two: for a `_DeltaSnapshot`, `inline_i` is the
|
||||
literal `true` marker rather than the value, so a blob must win over an
|
||||
inline reading whenever one exists. That ordering is what makes a genuine
|
||||
inline `true` (a bool channel) distinguishable from the marker.
|
||||
|
||||
The `type <> 'empty'` predicate mirrors the check stage 2 already applies
|
||||
when resolving the seed blob. `put` does not currently produce `empty` rows
|
||||
on this path — `blob_versions` is filtered to keys present in
|
||||
`channel_values`, so `_dump_blobs`' empty branch is unreachable from it —
|
||||
but without the predicate the two stages could disagree: stage 1 would
|
||||
terminate the walk on a row stage 2 then discards, yielding no seed *and* a
|
||||
truncated write chain, which is the failure this function exists to avoid.
|
||||
|
||||
Channel names are passed as `%s` parameters (safe from SQL injection).
|
||||
Only the column aliases `ver_i` / `hb_i` / `inline_i` and the subquery alias
|
||||
`b{i}` are interpolated into the SQL string (i is bounded by len(channels)
|
||||
and uses safe identifiers).
|
||||
|
||||
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
|
||||
cursor, cursor, page_size]` when `paged=True` — four per channel: the
|
||||
version lookup, the blob's channel, the version the blob must match, and the
|
||||
inline lookup.
|
||||
|
||||
When `paged=False`, the WHERE has no cursor predicate and there's no
|
||||
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
||||
@@ -228,7 +269,13 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
for i in range(len(channels)):
|
||||
cols.append(
|
||||
f"checkpoint -> 'channel_versions' ->> %s AS ver_{i}, "
|
||||
f"(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS hs_{i}"
|
||||
f"EXISTS (SELECT 1 FROM checkpoint_blobs b{i} "
|
||||
f"WHERE b{i}.thread_id = checkpoints.thread_id "
|
||||
f"AND b{i}.checkpoint_ns = checkpoints.checkpoint_ns "
|
||||
f"AND b{i}.channel = %s "
|
||||
f"AND b{i}.version = checkpoint -> 'channel_versions' ->> %s "
|
||||
f"AND b{i}.type <> 'empty') AS hb_{i}, "
|
||||
f"checkpoint -> 'channel_values' -> %s AS inline_{i}"
|
||||
)
|
||||
sql = (
|
||||
"SELECT checkpoint_id, parent_checkpoint_id, "
|
||||
@@ -342,7 +389,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
channels: Sequence[str],
|
||||
parent_of: dict[str, str | None],
|
||||
ver_by_i_by_cid: list[dict[str, str | None]],
|
||||
hs_by_i_by_cid: list[dict[str, bool]],
|
||||
hb_by_i_by_cid: list[dict[str, bool]],
|
||||
inline_by_i_by_cid: list[dict[str, Any]],
|
||||
) -> str | None:
|
||||
"""Fold one stage-1 page into the running walk-state mappings.
|
||||
|
||||
@@ -356,7 +404,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
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}"))
|
||||
hb_by_i_by_cid[i][cid] = bool(r.get(f"hb_{i}"))
|
||||
inline_by_i_by_cid[i][cid] = r.get(f"inline_{i}")
|
||||
# Rows are DESC; the last one is the smallest cid in the page.
|
||||
oldest = cid
|
||||
return oldest
|
||||
@@ -367,9 +416,11 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
channels: Sequence[str],
|
||||
parent_of: Mapping[str, str | None],
|
||||
ver_by_i_by_cid: Sequence[Mapping[str, str | None]],
|
||||
hs_by_i_by_cid: Sequence[Mapping[str, bool]],
|
||||
hb_by_i_by_cid: Sequence[Mapping[str, bool]],
|
||||
inline_by_i_by_cid: Sequence[Mapping[str, Any]],
|
||||
chain_by_ch: dict[str, list[str]],
|
||||
seed_ver_by_ch: dict[str, str | None],
|
||||
seed_inline_by_ch: dict[str, Any],
|
||||
walk_cursor_by_ch: dict[str, str | None],
|
||||
seeded: set[str],
|
||||
) -> None:
|
||||
@@ -377,14 +428,15 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
Uses the partial `parent_of` map accumulated so far. A walk stops
|
||||
either because:
|
||||
(a) it found a snapshot for its channel (channel becomes seeded),
|
||||
(a) it found a stored value for its channel — a blob or an inline
|
||||
primitive (channel becomes seeded),
|
||||
(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
|
||||
`seeded` in place.
|
||||
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
||||
`walk_cursor_by_ch`, and `seeded` in place.
|
||||
"""
|
||||
for i, ch in enumerate(channels):
|
||||
if ch in seeded:
|
||||
@@ -394,15 +446,22 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
walk_cursor_by_ch[ch] = parent_of.get(target_id)
|
||||
cur_cid = walk_cursor_by_ch[ch]
|
||||
ch_chain = chain_by_ch[ch]
|
||||
hs_i = hs_by_i_by_cid[i]
|
||||
hb_i = hb_by_i_by_cid[i]
|
||||
inline_i = inline_by_i_by_cid[i]
|
||||
ver_i = ver_by_i_by_cid[i]
|
||||
while cur_cid is not None:
|
||||
if cur_cid not in parent_of:
|
||||
# Need more pages to continue this walk.
|
||||
break
|
||||
ch_chain.append(cur_cid)
|
||||
if hs_i.get(cur_cid, False):
|
||||
has_blob = hb_i.get(cur_cid, False)
|
||||
inline = inline_i.get(cur_cid)
|
||||
if has_blob or inline is not None:
|
||||
# A blob wins: for a `_DeltaSnapshot` the inline reading is
|
||||
# the `true` marker, not the value.
|
||||
seed_ver_by_ch[ch] = ver_i.get(cur_cid)
|
||||
if not has_blob:
|
||||
seed_inline_by_ch[ch] = inline
|
||||
seeded.add(ch)
|
||||
cur_cid = None
|
||||
break
|
||||
@@ -415,16 +474,23 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
channels: Sequence[str],
|
||||
chain_by_ch: Mapping[str, list[str]],
|
||||
seed_ver_by_ch: Mapping[str, str | None],
|
||||
seed_inline_by_ch: Mapping[str, Any],
|
||||
stage2_rows: Sequence[_DeltaStage2Row],
|
||||
) -> dict[str, DeltaChannelHistory]:
|
||||
"""Demux stage 2 rows per channel; produce per-channel histories.
|
||||
|
||||
stage2_rows carry `channel` on every row. We build per-channel
|
||||
`writes_by_cid` and per-channel `seed_blob` dicts, then assemble
|
||||
a `DeltaChannelHistory` per requested channel. The `seed` key is omitted
|
||||
when the walk reached root with no snapshot found, or when the
|
||||
seed blob is sentinel "empty" — in both cases the consumer treats
|
||||
absence as "start empty".
|
||||
a `DeltaChannelHistory` per requested channel.
|
||||
|
||||
A seed comes from the blobs table when the walk found one there, and
|
||||
otherwise from `seed_inline_by_ch` — `put` keeps `None`, `str`, `int`,
|
||||
`float` and `bool` values in the checkpoint's own `channel_values` with
|
||||
no blob row, so those never appear in `stage2_rows`.
|
||||
|
||||
The `seed` key is omitted when the walk reached root without finding a
|
||||
stored value, or when the seed blob is sentinel "empty" — in both cases
|
||||
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: dict[str, dict[str, list[tuple[str, bytes, str, int]]]] = {
|
||||
@@ -473,6 +539,10 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
blob = seed_blob_by_ver.get((ch, seed_version))
|
||||
if blob is not None and blob[0] != "empty":
|
||||
entry["seed"] = self.serde.loads_typed(blob)
|
||||
elif ch in seed_inline_by_ch:
|
||||
# Inline primitive: stored in the checkpoint, not the blobs
|
||||
# table, so stage 2 never returned a row for it.
|
||||
entry["seed"] = seed_inline_by_ch[ch]
|
||||
result[ch] = entry
|
||||
return result
|
||||
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
"""Seed detection for `DeltaChannel` histories on Postgres.
|
||||
|
||||
`put` splits stored values in two: primitives stay inline in the checkpoint's
|
||||
`channel_values`, everything else moves to `checkpoint_blobs`. Only
|
||||
`_DeltaSnapshot` leaves an inline marker behind when it moves, so the stage-1
|
||||
walk has to check both places — a blob probe alone misses inline primitives, and
|
||||
an inline-key check alone missed blob-stored plain values, which is what a thread
|
||||
migrated from a pre-delta channel type leaves behind. See #8534.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import Checkpoint, empty_checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
CHANNEL = "items"
|
||||
|
||||
|
||||
async def _build_chain(saver: AsyncPostgresSaver, seed_value: Any) -> tuple[str, dict]:
|
||||
"""Store `seed_value` at step 1, then two steps that store nothing.
|
||||
|
||||
Every step carries a write so the walk has something to collect.
|
||||
Returns `(thread_id, head_config)`.
|
||||
"""
|
||||
thread_id = str(uuid4())
|
||||
parent: dict | None = None
|
||||
for step in range(4):
|
||||
config: dict = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
if parent is not None:
|
||||
config["configurable"]["checkpoint_id"] = parent["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
cp: Checkpoint = empty_checkpoint()
|
||||
cp["id"] = str(uuid6(clock_seq=step))
|
||||
new_versions: dict[str, Any] = {}
|
||||
if step == 1:
|
||||
cp["channel_values"][CHANNEL] = seed_value
|
||||
cp["channel_versions"][CHANNEL] = "v1"
|
||||
new_versions[CHANNEL] = "v1"
|
||||
else:
|
||||
cp["channel_versions"][CHANNEL] = f"v{step}"
|
||||
parent = await saver.aput(
|
||||
config, cp, {"source": "loop", "step": step, "parents": {}}, new_versions
|
||||
)
|
||||
await saver.aput_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
|
||||
assert parent is not None
|
||||
return thread_id, parent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plain_value_seed_is_found() -> None:
|
||||
"""A pre-delta plain value must be located as the seed.
|
||||
|
||||
Before #8534 the walk ran to the root and returned no seed, which happens
|
||||
to reconstruct correctly for additive reducers while costing an
|
||||
O(thread length) replay on every read.
|
||||
"""
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
_, head = await _build_chain(saver, [10, 20])
|
||||
|
||||
result = await saver.aget_delta_channel_history(config=head, channels=[CHANNEL])
|
||||
entry = result[CHANNEL]
|
||||
|
||||
assert entry.get("seed") == [10, 20], (
|
||||
f"expected the plain value as seed, got {entry.get('seed', '<missing>')}"
|
||||
)
|
||||
# Only the writes between the seed and the head's parent replay: step 1
|
||||
# (the seed's own) and step 2. Step 0 is older than the seed, step 3 is
|
||||
# pending at the head.
|
||||
assert [w[2] for w in entry["writes"]] == ["w1", "w2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_snapshot_seed_is_found() -> None:
|
||||
"""The `_DeltaSnapshot` path keeps working, so both seed kinds agree."""
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
_, head = await _build_chain(saver, _DeltaSnapshot([10, 20]))
|
||||
|
||||
result = await saver.aget_delta_channel_history(config=head, channels=[CHANNEL])
|
||||
entry = result[CHANNEL]
|
||||
|
||||
seed = entry.get("seed")
|
||||
assert isinstance(seed, _DeltaSnapshot), f"expected a snapshot, got {seed!r}"
|
||||
assert seed.value == [10, 20]
|
||||
assert [w[2] for w in entry["writes"]] == ["w1", "w2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_version_bump_without_a_value_does_not_hide_an_older_seed() -> None:
|
||||
"""A delta-era step bumps `channel_versions` without storing a value, so no
|
||||
blob exists for that version. The probe must report no seed there and keep
|
||||
walking rather than stopping at a version it cannot resolve.
|
||||
|
||||
Step 0 holds the real value; step 1 bumps the version with nothing stored.
|
||||
Walking back from the head has to pass step 1 to reach step 0.
|
||||
"""
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
thread_id = str(uuid4())
|
||||
parent: dict | None = None
|
||||
for step in range(4):
|
||||
config: dict = {
|
||||
"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}
|
||||
}
|
||||
if parent is not None:
|
||||
config["configurable"]["checkpoint_id"] = parent["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
cp: Checkpoint = empty_checkpoint()
|
||||
cp["id"] = str(uuid6(clock_seq=step))
|
||||
new_versions: dict[str, Any] = {}
|
||||
if step == 0:
|
||||
cp["channel_values"][CHANNEL] = [10, 20]
|
||||
cp["channel_versions"][CHANNEL] = "v0"
|
||||
new_versions[CHANNEL] = "v0"
|
||||
elif step == 1:
|
||||
# Version bumped, value absent -> no blob row written.
|
||||
cp["channel_versions"][CHANNEL] = "v1"
|
||||
new_versions[CHANNEL] = "v1"
|
||||
else:
|
||||
cp["channel_versions"][CHANNEL] = "v1"
|
||||
parent = await saver.aput(
|
||||
config,
|
||||
cp,
|
||||
{"source": "loop", "step": step, "parents": {}},
|
||||
new_versions,
|
||||
)
|
||||
await saver.aput_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
|
||||
assert parent is not None
|
||||
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=parent, channels=[CHANNEL]
|
||||
)
|
||||
entry = result[CHANNEL]
|
||||
|
||||
assert entry.get("seed") == [10, 20], (
|
||||
"the walk stopped at the empty blob instead of reaching the real "
|
||||
f"value at step 0; got {entry.get('seed', '<missing>')}"
|
||||
)
|
||||
assert [w[2] for w in entry["writes"]] == ["w0", "w1", "w2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_primitive_seed_is_found() -> None:
|
||||
"""`put` keeps `None`, `str`, `int`, `float` and `bool` in the checkpoint's
|
||||
own `channel_values` with no blob row, so a blob probe alone cannot see
|
||||
them. Stage 1 reads the inline value too and uses it when there is no blob.
|
||||
"""
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
for seed_value in (42, "x", 3.5, None):
|
||||
_, head = await _build_chain(saver, seed_value)
|
||||
entry = (
|
||||
await saver.aget_delta_channel_history(config=head, channels=[CHANNEL])
|
||||
)[CHANNEL]
|
||||
if seed_value is None:
|
||||
# A JSON null is indistinguishable from "no value stored", so
|
||||
# the walk keeps going; replay from empty is the correct result.
|
||||
assert "seed" not in entry
|
||||
else:
|
||||
assert entry.get("seed") == seed_value, (
|
||||
f"inline {type(seed_value).__name__} seed not found: "
|
||||
f"{entry.get('seed', '<missing>')!r}"
|
||||
)
|
||||
assert [w[2] for w in entry["writes"]] == ["w1", "w2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_true_is_not_read_as_a_snapshot_marker() -> None:
|
||||
"""`put` inlines a literal `true` in `channel_values` as the marker for a
|
||||
`_DeltaSnapshot`, which is also what a genuine `bool` channel holding
|
||||
`True` looks like. A blob exists only in the snapshot case, so preferring
|
||||
the blob keeps the two apart.
|
||||
"""
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
|
||||
_, head = await _build_chain(saver, True)
|
||||
entry = (
|
||||
await saver.aget_delta_channel_history(config=head, channels=[CHANNEL])
|
||||
)[CHANNEL]
|
||||
assert entry.get("seed") is True, (
|
||||
f"a real inline True must survive, got {entry.get('seed', '<missing>')!r}"
|
||||
)
|
||||
|
||||
_, snap_head = await _build_chain(saver, _DeltaSnapshot(True))
|
||||
snap_entry = (
|
||||
await saver.aget_delta_channel_history(config=snap_head, channels=[CHANNEL])
|
||||
)[CHANNEL]
|
||||
seed = snap_entry.get("seed")
|
||||
assert isinstance(seed, _DeltaSnapshot), (
|
||||
f"the marker must resolve to the blob, not inline true; got {seed!r}"
|
||||
)
|
||||
assert seed.value is True
|
||||
Reference in New Issue
Block a user