chore(langgraph): use two phase read to avoid unnecessary data transport (#7660)

## Summary

Replaces the single-roundtrip `UNION ALL` DeltaChannel read with a
two-stage query that avoids fetching unused snapshot blobs, then removes
the old combined path entirely.

### Problem

`_get_channel_writes_history` used a single `UNION ALL` query that
fetched **all** checkpoint metadata, writes, and blobs for a
`(thread_id, channel)` in one shot. With `snapshot_frequency=N`, this
pulled back O(N/freq) full-size snapshot blobs even though only the
nearest one is needed to seed reconstruction. At 500 turns with
`snapshot_frequency=10`, this meant fetching ~100 complete
message-history snapshots per read.

### Solution

Two-stage read:
- **Stage 1** — lightweight scan of `checkpoints` only (no blob bytes):
walks the parent chain from the target checkpoint and stops at the first
ancestor with a snapshot, returning `chain_cids` and `seed_version`
- **Stage 2** — targeted fetch: only the writes for `chain_cids` and the
single seed blob at `seed_version`

The two-stage path is now unconditional — the old combined query and
`LG_DELTA_TWO_STAGE_QUERY` env-var gate have been removed.

### Sentinel cleanup

`DELTA_SENTINEL` is now a pure in-memory signal and is never written to
storage:
- Postgres `put()` already stripped it from `channel_values` before
writing blobs
- Memory saver `put()` now stores `"empty"` instead of serializing the
sentinel
- `EXT_DELTA_SENTINEL` (msgpack ext code 8) removed from
`JsonPlusSerializer`
- `DELTA_SENTINEL` is kept as an in-memory marker:
`DeltaChannel.checkpoint()` returns it so savers know to skip it, and
`_ChannelWritesHistory.seed` uses it to mean "no snapshot found, start
from empty"

## Performance

Benchmarked at `snapshot_frequency=10` on Postgres (`~100 tok/msg`):

| turns | old combined query | two-stage |
|------:|-------------------:|----------:|
| 50    | 6.0ms              | 2.8ms  (2.1x faster) |
| 100   | 10.1ms             | 5.6ms  (1.8x faster) |
| 500   | **216.1ms**        | 15.3ms (**14x faster**) |

The old query's read time grew super-linearly with turn count because
each read fetched O(N/freq) full snapshot blobs. Two-stage keeps read
depth bounded by `snapshot_frequency` regardless of thread length.

## Test plan

- `make test` in `libs/checkpoint`, `libs/checkpoint-postgres`,
`libs/langgraph`
- Removed `test_delta_sentinel_serde_round_trip` (sentinel no longer
serializable)
- Updated `test_memory.py` — delta channel blobs stored as `"empty"`,
not serialized sentinel
- Updated `test_channels.py` — `channel_values` no longer contains
sentinel key for DeltaChannels
- Deleted `test_delta_channel_two_stage_benchmark.py` (one-stage vs
two-stage comparison; path no longer exists)

---------

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Quanzheng Long
2026-05-01 11:06:54 -04:00
committed by GitHub
co-authored by Sydney Runkle Claude Sonnet 4.6
parent c8c58a0768
commit d8b7800183
9 changed files with 170 additions and 160 deletions
@@ -19,6 +19,7 @@ from langgraph.checkpoint.base import (
get_serializable_checkpoint_metadata,
)
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
@@ -26,9 +27,11 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_COMBINED_SQL,
SELECT_DELTA_STAGE1_SQL,
SELECT_DELTA_STAGE2_SQL,
BasePostgresSaver,
_DeltaCombinedRow,
_DeltaStage1Row,
_DeltaStage2Row,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -308,7 +311,12 @@ class PostgresSaver(BasePostgresSaver):
# others are stored in blobs table
blob_values = {}
for k, v in checkpoint["channel_values"].items():
if v is None or isinstance(v, (str, int, float, bool)):
if v is DELTA_SENTINEL:
copy["channel_values"].pop(k)
elif isinstance(v, _DeltaSnapshot):
blob_values[k] = copy["channel_values"].pop(k)
copy["channel_values"][k] = True
elif v is None or isinstance(v, (str, int, float, bool)):
pass
else:
blob_values[k] = copy["channel_values"].pop(k)
@@ -441,41 +449,49 @@ class PostgresSaver(BasePostgresSaver):
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; the ancestor walk runs in Python.
Two-stage query: stage 1 scans checkpoint metadata to walk the parent
chain and locate the nearest snapshot; stage 2 fetches only the
chain-limited writes and single seed blob.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
# Caller didn't specify a target — resolve to the latest
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
# returns the newest; its config carries the resolved id.
target = self.get_tuple(config)
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
with self._cursor() as cur:
cur.execute(
SELECT_DELTA_COMBINED_SQL,
SELECT_DELTA_STAGE1_SQL,
(channel, channel, thread_id, checkpoint_ns),
)
stage1_rows = cur.fetchall()
chain_cids, seed_version = self._walk_stage1(
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
)
seed_versions = [seed_version] if seed_version else []
with self._cursor() as cur:
cur.execute(
SELECT_DELTA_STAGE2_SQL,
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
chain_cids,
thread_id,
checkpoint_ns,
channel,
seed_versions,
),
)
rows = cur.fetchall()
stage2_rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
target_id=checkpoint_id,
rows=cast("list[_DeltaCombinedRow]", rows),
chain_cids=chain_cids,
seed_version=seed_version,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
)
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: