From cd8335263e0bc257e4facd3ca6cb23c0168d72e0 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Wed, 5 Aug 2026 23:05:51 -0400 Subject: [PATCH] fix(checkpoint-postgres): derive the delta walk cursor once the target loads `get_delta_channel_history` pages the `checkpoints` table newest-first from the head of the thread, so the target checkpoint's own row does not land until paging reaches back to it. `_try_advance_walks` read the walk cursor out of the partial `parent_of` map on the first page regardless, and `parent_of.get(target_id)` returns `None` there for a target that is simply not loaded yet, which is the same value that means "the target is a root". The cursor is derived under a `ch not in walk_cursor_by_ch` guard, so it is never recomputed. A target older than one page (1024 checkpoints) kept that `None` for the rest of the run: the chain stayed empty, the channel never seeded, and the history came back with no seed and no writes. Downstream `channels_from_checkpoint` hydrates that as an empty channel, so `get_state`, `get_state_history` and `update_state` against an older checkpoint returned an empty `DeltaChannel` on a thread with real accumulated history, silently and without an error. Start the walk only once `target_id` is actually in `parent_of`, so "not loaded yet" and "is a root" stop sharing a representation. Both savers share this method, so both are covered. Fixes #8448 Co-authored-by: Navneet-Scaler <147032454+Navneet-Scaler@users.noreply.github.com> --- .../langgraph/checkpoint/postgres/base.py | 15 +- .../tests/test_delta_pagination.py | 172 ++++++++++++++++++ 2 files changed, 185 insertions(+), 2 deletions(-) create mode 100644 libs/checkpoint-postgres/tests/test_delta_pagination.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index d58d2cc38..1ad0a9b74 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -435,15 +435,26 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): (c) the next ancestor cid isn't in `parent_of` yet (waiting for a later page; the cursor stays put). + A walk that hasn't started yet is a fourth state, distinct from all + three: pages arrive newest-first from the head of the thread, so a + target deeper than one page isn't in `parent_of` on the early pages + and the walk can only begin once its own row lands. + 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: continue - # First-time entry: cursor starts at the target's parent. + # First-time entry: cursor starts at the target's parent, but + # only once the target's own row has loaded. Reading it earlier + # would record `None` for "target not seen yet" the same way it + # records `None` for "target is a root", and this guard fires + # only once, so the walk would stay stranded for good. if ch not in walk_cursor_by_ch: - walk_cursor_by_ch[ch] = parent_of.get(target_id) + if target_id not in parent_of: + continue + walk_cursor_by_ch[ch] = parent_of[target_id] cur_cid = walk_cursor_by_ch[ch] ch_chain = chain_by_ch[ch] hb_i = hb_by_i_by_cid[i] diff --git a/libs/checkpoint-postgres/tests/test_delta_pagination.py b/libs/checkpoint-postgres/tests/test_delta_pagination.py new file mode 100644 index 000000000..ee7c2b498 --- /dev/null +++ b/libs/checkpoint-postgres/tests/test_delta_pagination.py @@ -0,0 +1,172 @@ +"""Stage-1 pagination for `DeltaChannel` histories on Postgres. + +`get_delta_channel_history` pages the `checkpoints` table newest-first in +chunks of `_DELTA_PAGE_SIZE`, starting from the head of the thread rather than +from the target checkpoint. The target's own row therefore only lands once +paging has reached back to it, which for a long thread can be several pages in. + +Until then `parent_of` has no entry for it, and reading the walk cursor out of +that map records `None`, the same value that means "the target is a root". The +cursor is derived once, so a target older than the first page kept that `None` +forever: no seed, no writes, and the channel hydrated empty with no error. +See #8448. + +These tests shrink the page size instead of writing 1024+ real checkpoints per +case. The behaviour under test is "the target is not on the first page", and +the page the target lands on is the only thing that decides it. +""" + +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +import pytest +from langgraph.checkpoint.base import ( + Checkpoint, + DeltaChannelHistory, + empty_checkpoint, +) +from langgraph.checkpoint.base.id import uuid6 +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 tests.conftest import DEFAULT_URI + +CHANNEL = "items" +STEPS = 8 +SEED_STEP = 1 +SEED_VALUE = [10, 20] +# Deep enough that the smaller page sizes below have to page past it. +TARGET_STEP = 4 + +# The real page size holds this whole thread on one page, so it is the control. +# The rest each leave the target off the first page: there are three +# checkpoints newer than `TARGET_STEP`, so any page size at or below 3 does it. +PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1] + + +def _step_args( + thread_id: str, step: int, parent: dict | None +) -> tuple[dict, Checkpoint, dict[str, Any]]: + """Return the `(config, checkpoint, new_versions)` triple for one step. + + Step `SEED_STEP` stores a snapshot for `CHANNEL`; the rest only bump the + channel version, which is what a delta channel does between snapshots. + """ + config: dict = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + if parent is not None: + config["configurable"]["checkpoint_id"] = parent["configurable"][ + "checkpoint_id" + ] + checkpoint: Checkpoint = empty_checkpoint() + checkpoint["id"] = str(uuid6(clock_seq=step)) + checkpoint["channel_versions"][CHANNEL] = f"v{step}" + if step == SEED_STEP: + checkpoint["channel_values"][CHANNEL] = _DeltaSnapshot(list(SEED_VALUE)) + return config, checkpoint, {CHANNEL: f"v{step}"} + return config, checkpoint, {} + + +async def _abuild_chain(saver: AsyncPostgresSaver) -> list[dict]: + """Write `STEPS` linked checkpoints, one write each; return their configs.""" + thread_id = str(uuid4()) + parent: dict | None = None + configs: list[dict] = [] + for step in range(STEPS): + config, checkpoint, new_versions = _step_args(thread_id, step, parent) + parent = await saver.aput( + config, + checkpoint, + {"source": "loop", "step": step, "parents": {}}, + new_versions, + ) + await saver.aput_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4())) + configs.append(parent) + return configs + + +def _build_chain(saver: PostgresSaver) -> list[dict]: + """Sync twin of `_abuild_chain`.""" + thread_id = str(uuid4()) + parent: dict | None = None + configs: list[dict] = [] + for step in range(STEPS): + config, checkpoint, new_versions = _step_args(thread_id, step, parent) + parent = saver.put( + config, + checkpoint, + {"source": "loop", "step": step, "parents": {}}, + new_versions, + ) + saver.put_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4())) + configs.append(parent) + return configs + + +def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None: + """Check the walk from `TARGET_STEP` back to the snapshot at `SEED_STEP`. + + The chain is steps 3, 2 and 1: the target's own writes are pending for its + next super-step and excluded, and the walk stops at the snapshot. Writes + come back oldest first. + """ + seed = entry.get("seed") + assert isinstance(seed, _DeltaSnapshot), ( + f"page_size={page_size}: expected a snapshot seed, " + f"got {entry.get('seed', '')!r}" + ) + assert seed.value == SEED_VALUE + assert [w[2] for w in entry["writes"]] == ["w1", "w2", "w3"], ( + f"page_size={page_size}: got {[w[2] for w in entry['writes']]}" + ) + + +@pytest.mark.parametrize("page_size", PAGE_SIZES) +async def test_async_target_older_than_the_first_page( + page_size: int, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size) + 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], page_size) + + +@pytest.mark.parametrize("page_size", PAGE_SIZES) +def test_sync_target_older_than_the_first_page( + page_size: int, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size) + 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], page_size) + + +async def test_root_target_has_no_history_and_still_terminates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A root target is the case where a `None` cursor is the right answer. + + Its walk has nowhere to go, so it never seeds and collects none of the + thread's writes, and paging has to run out on a short page rather than + spin. Page size 1 gives one page per checkpoint, so the loop runs the + length of the thread before stopping. + """ + monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", 1) + 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[0], channels=[CHANNEL] + ) + assert result[CHANNEL] == {"writes": []}