diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index aa20f5339..78222205f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -451,7 +451,8 @@ class PostgresSaver(BasePostgresSaver): * Stage 1 (paged): dynamic SELECT over `checkpoints` with three columns per channel: its version, an `EXISTS` probe for a stored blob at that version, and its inline value. Pages newest-first by - `checkpoint_id` with a cursor; page size is `_DELTA_PAGE_SIZE`. + `checkpoint_id`, starting at the target; page size is + `_DELTA_PAGE_SIZE`. Stops paging when every channel has found its seed or a page comes back short. @@ -476,7 +477,7 @@ class PostgresSaver(BasePostgresSaver): # Stage 1: paged K-JSONB-lookup scan, walking the parent chain in # Python after each page. Stops as soon as every channel has its seed. - stage1_sql = _build_delta_stage1_sql(channels, paged=True) + stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True) parent_of: dict[str, str | None] = {} ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels] hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels] @@ -486,7 +487,7 @@ class PostgresSaver(BasePostgresSaver): seed_inline_by_ch: dict[str, Any] = {} walk_cursor_by_ch: dict[str, str | None] = {} seeded: set[str] = set() - cursor: str | None = None + cursor: str | None = checkpoint_id with self._cursor() as cur: while True: @@ -495,7 +496,7 @@ class PostgresSaver(BasePostgresSaver): # 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] + [thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE] ) cur.execute(stage1_sql, stage1_params) page = cur.fetchall() @@ -527,6 +528,7 @@ class PostgresSaver(BasePostgresSaver): if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE: break cursor = oldest + stage1_sql = _build_delta_stage1_sql(channels, paged=True) # Stage 2: per-channel UNION ALL — one writes branch per channel # with non-empty chain, plus one blob branch per seeded channel. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 246c47c0a..2b813b391 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -423,7 +423,7 @@ class AsyncPostgresSaver(BasePostgresSaver): return {ch: {"writes": []} for ch in channels} checkpoint_id = target.config["configurable"]["checkpoint_id"] - stage1_sql = _build_delta_stage1_sql(channels, paged=True) + stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True) parent_of: dict[str, str | None] = {} ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels] hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels] @@ -433,7 +433,7 @@ class AsyncPostgresSaver(BasePostgresSaver): seed_inline_by_ch: dict[str, Any] = {} walk_cursor_by_ch: dict[str, str | None] = {} seeded: set[str] = set() - cursor: str | None = None + cursor: str | None = checkpoint_id async with self._cursor() as cur: while True: @@ -442,7 +442,7 @@ class AsyncPostgresSaver(BasePostgresSaver): # 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] + [thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE] ) await cur.execute(stage1_sql, stage1_params) page = await cur.fetchall() @@ -472,6 +472,7 @@ class AsyncPostgresSaver(BasePostgresSaver): if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE: break cursor = oldest + stage1_sql = _build_delta_stage1_sql(channels, paged=True) channels_with_chain = [ch for ch in channels if chain_by_ch[ch]] channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 343107884..81b618c03 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -178,10 +178,13 @@ class _DeltaStage2Row(TypedDict, total=False): # `_build_delta_stage2_sql` document their shapes. -def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: +def _build_delta_stage1_sql( + channels: Sequence[str], *, paged: bool, include_cursor: bool = False +) -> str: """Build stage 1 SQL with K parallel version lookups + seed probes. - For channels=["messages", "files"] (with `paged=True`) the result is:: + For channels=["messages", "files"] (with `paged=True, include_cursor=True`) + the result is:: SELECT checkpoint_id, parent_checkpoint_id, checkpoint -> 'channel_versions' ->> %s AS ver_0, @@ -197,7 +200,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: 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) + AND checkpoint_id <= %s ORDER BY checkpoint_id DESC LIMIT %s @@ -238,9 +241,16 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: 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. + 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. + + Pages run newest-first from the target down. The first page passes the + target as the cursor with `include_cursor=True`, so it opens with the + target's own row, whose parent starts the walk; each later page continues + below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6 + is time-ordered, which `get_tuple` also relies on to find the latest + checkpoint), so no row newer than the target is part of its chain. 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. @@ -264,7 +274,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str: ) if paged: sql += ( - " AND (%s::text IS NULL OR checkpoint_id < %s)" + f" AND checkpoint_id {'<=' if include_cursor else '<'} %s" " ORDER BY checkpoint_id DESC LIMIT %s" ) return sql @@ -413,8 +423,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): materialized at this point), (c) the next ancestor cid isn't in `parent_of` yet (waiting for a later page; the cursor stays put), or - (d) the target's own row isn't in `parent_of` yet (the walk has - not started; no cursor is set, so a later page retries). + (d) the target's own row isn't in `parent_of` (the target doesn't + exist, so the walk never starts). Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`, `walk_cursor_by_ch`, and `seeded` in place. @@ -422,8 +432,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): for i, ch in enumerate(channels): if ch in seeded: continue - # Pages start at the thread head, so the target may not have - # loaded yet; a `None` cursor would read as "target is a root". + # The first page opens with the target's row, so it's missing only + # when the target doesn't exist; a `None` cursor would read as + # "target is a root". if ch not in walk_cursor_by_ch: if target_id not in parent_of: continue diff --git a/libs/checkpoint-postgres/tests/test_delta_pagination.py b/libs/checkpoint-postgres/tests/test_delta_pagination.py index 6f6cfeaff..0c7407068 100644 --- a/libs/checkpoint-postgres/tests/test_delta_pagination.py +++ b/libs/checkpoint-postgres/tests/test_delta_pagination.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import Any from uuid import uuid4 @@ -14,7 +15,7 @@ from langgraph.checkpoint.serde.types import _DeltaSnapshot from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver -from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE +from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE, BasePostgresSaver from tests.conftest import DEFAULT_URI CHANNEL = "items" @@ -23,8 +24,8 @@ SEED_STEP = 1 SEED_VALUE = [10, 20] TARGET_STEP = 4 -# The real page size is the control; the rest leave the target off the first -# page (three checkpoints are newer than it). +# The real page size is the control; the rest split the walk from the target +# to its seed across pages. PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1] @@ -92,7 +93,7 @@ def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None: @pytest.mark.parametrize("page_size", PAGE_SIZES) -async def test_async_target_older_than_the_first_page( +async def test_async_walk_continues_across_pages( page_size: int, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size) @@ -106,7 +107,7 @@ async def test_async_target_older_than_the_first_page( @pytest.mark.parametrize("page_size", PAGE_SIZES) -def test_sync_target_older_than_the_first_page( +def test_sync_walk_continues_across_pages( page_size: int, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size) @@ -119,6 +120,50 @@ def test_sync_target_older_than_the_first_page( _assert_history(result[CHANNEL], page_size) +def _record_rows_read(monkeypatch: pytest.MonkeyPatch) -> list[str]: + read: list[str] = [] + ingest = BasePostgresSaver._ingest_stage1_page + + def record(rows: Sequence[Mapping[str, Any]], *args: Any) -> str | None: + read.extend(row["checkpoint_id"] for row in rows) + return ingest(rows, *args) + + monkeypatch.setattr(BasePostgresSaver, "_ingest_stage1_page", staticmethod(record)) + return read + + +def _ids_from_target_down(configs: list[dict]) -> list[str]: + return [c["configurable"]["checkpoint_id"] for c in configs[TARGET_STEP::-1]] + + +async def test_async_walk_reads_nothing_newer_than_the_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + read = _record_rows_read(monkeypatch) + async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: + await saver.setup() + configs = await _abuild_chain(saver) + result = await saver.aget_delta_channel_history( + config=configs[TARGET_STEP], channels=[CHANNEL] + ) + _assert_history(result[CHANNEL], _DELTA_PAGE_SIZE) + assert read == _ids_from_target_down(configs) + + +def test_sync_walk_reads_nothing_newer_than_the_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + read = _record_rows_read(monkeypatch) + with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: + saver.setup() + configs = _build_chain(saver) + result = saver.get_delta_channel_history( + config=configs[TARGET_STEP], channels=[CHANNEL] + ) + _assert_history(result[CHANNEL], _DELTA_PAGE_SIZE) + assert read == _ids_from_target_down(configs) + + async def test_root_target_has_no_history_and_still_terminates( monkeypatch: pytest.MonkeyPatch, ) -> None: