From af82510cf9bf348d0507cd15feff9c6ac928186a Mon Sep 17 00:00:00 2001 From: Will Fu-Hinthorn Date: Wed, 29 Apr 2026 13:56:25 -0700 Subject: [PATCH] cleanup --- .../langgraph/checkpoint/postgres/__init__.py | 3 +- .../langgraph/checkpoint/postgres/aio.py | 5 +-- .../langgraph/checkpoint/postgres/base.py | 12 ++---- .../langgraph/checkpoint/base/__init__.py | 12 ++---- libs/langgraph/langgraph/pregel/_loop.py | 6 +-- .../tests/test_delta_channel_benchmark.py | 6 ++- libs/langgraph/tests/test_pregel.py | 38 ++----------------- 7 files changed, 21 insertions(+), 61 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index cd2fb33de..237fd8932 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -443,8 +443,7 @@ class PostgresSaver(BasePostgresSaver): One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a - single roundtrip. Rationale + benchmark in - `notes/delta_channel_query_bench.md`. + single roundtrip; the ancestor walk runs in Python. """ thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index bd3899f78..965c8f31b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -404,9 +404,8 @@ class AsyncPostgresSaver(BasePostgresSaver): One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a - single roundtrip; rows assembled by the shared pure helper on - `BasePostgresSaver`. Rationale + benchmark in - `notes/delta_channel_query_bench.md`. + single roundtrip; rows are assembled by the shared pure helper on + `BasePostgresSaver`. """ thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"].get("checkpoint_ns", "") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 50103b7c9..d5b963c87 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -178,12 +178,9 @@ class _DeltaCombinedRow(TypedDict, total=False): version: str | None -# DeltaChannel reconstruction: one combined CTE+UNION ALL query per channel. -# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE -# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are -# 3x-100x faster in the realistic depth range and the Python walk is O(n). -# The three plain SELECTs are further collapsed into one UNION ALL query so -# that only one roundtrip is needed per channel reconstruction. +# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints, +# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs +# in Python in `_build_delta_channel_writes_history`. # # Parameter order: (channel, thread_id, checkpoint_ns, # thread_id, checkpoint_ns, channel, @@ -301,8 +298,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): "tuple[str, bytes]", (r["type"], r["blob"]) ) - # Sort writes within each checkpoint (task_id DESC, idx DESC) to match - # the prior CTE ordering — newest write first per ancestor. + # newest write first per ancestor (task_id DESC, idx DESC) for ws in writes_by_cid.values(): ws.sort(key=lambda w: (w[2], w[3]), reverse=True) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 4cfa32f5a..891f7c95e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -490,15 +490,9 @@ class BaseCheckpointSaver(Generic[V]): """Pure storage read used by `_get_channel_writes_history`. Must return the same value as `get_tuple` but must NOT trigger channel - reconstruction (i.e., must not call `channels_from_checkpoint`). The - default implementation delegates to `get_tuple`, which is correct for - savers whose `get_tuple` is a pure storage query (the common case). - - Override this if your saver performs channel hydration inside `get_tuple`. - Doing so structurally breaks the otherwise-possible cycle: - _get_channel_writes_history -> _get_tuple_raw -> get_tuple - -> channels_from_checkpoint - -> _get_channel_writes_history (cycle!) + reconstruction; otherwise the channel-hydration path would re-enter + `_get_channel_writes_history`. Override only if `get_tuple` itself + performs channel hydration. """ return self.get_tuple(config) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index a3a3a4765..e6aea17d8 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -190,6 +190,8 @@ class PregelLoop: _migrate_checkpoint: Callable[[Checkpoint], None] | None submit: Submit channels: Mapping[str, BaseChannel] + # Only set on AsyncPregelLoop; sync loops keep this as None. + _delta_write_futs: list[Any] | None = None managed: ManagedValueMapping checkpoint: Checkpoint checkpoint_id_saved: str @@ -422,7 +424,7 @@ class PregelLoop: writes_to_save, task_id, ) - if hasattr(self, "_delta_write_futs") and any( + if self._delta_write_futs is not None and any( isinstance(self.specs.get(c), DeltaChannel) for c, _ in writes_to_save ): self._delta_write_futs.append(fut) @@ -1313,8 +1315,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): - _delta_write_futs: list[asyncio.Future[Any]] - def __init__( self, input: Any | None, diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 8136ed717..cf17d0f12 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -17,6 +17,7 @@ from __future__ import annotations import contextlib import math +import os import sys import time from typing import Annotated, Any @@ -34,7 +35,10 @@ try: from langgraph.checkpoint.postgres import PostgresSaver _POSTGRES_AVAILABLE = True - _POSTGRES_URI = "postgres://sydney_runkle@localhost:5441/postgres?sslmode=disable" + _POSTGRES_URI = os.environ.get( + "LANGGRAPH_BENCH_POSTGRES_URI", + "postgres://postgres@localhost:5432/postgres?sslmode=disable", + ) except ImportError: _POSTGRES_AVAILABLE = False diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 677fe7cf3..d010eecbe 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -16,7 +16,7 @@ from typing import Annotated, Any, Literal, get_type_hints import pytest from langchain_core.language_models import GenericFakeChatModel -from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, RemoveMessage from langchain_core.runnables import ( RunnableConfig, RunnableLambda, @@ -25,6 +25,7 @@ from langchain_core.runnables import ( from langchain_core.runnables.graph import Edge from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( + DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint, CheckpointMetadata, @@ -50,7 +51,7 @@ from langgraph.config import get_stream_writer from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand from langgraph.func import entrypoint, task from langgraph.graph import END, START, StateGraph -from langgraph.graph.message import MessagesState, add_messages +from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages from langgraph.pregel import ( NodeBuilder, Pregel, @@ -9405,11 +9406,6 @@ def test_fork_does_not_apply_pending_writes( async def test_delta_channel_end_to_end_inmemory() -> None: """Full graph run: DeltaChannel accumulates correctly across multiple turns.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9446,11 +9442,6 @@ async def test_delta_channel_end_to_end_inmemory() -> None: async def test_delta_channel_time_travel() -> None: """Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9503,11 +9494,6 @@ async def test_delta_channel_time_travel() -> None: async def test_delta_channel_remove_message_end_to_end() -> None: """RemoveMessage inside a DeltaChannel graph must persist and reload correctly.""" - from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9549,11 +9535,6 @@ async def test_delta_channel_remove_message_end_to_end() -> None: async def test_delta_channel_update_by_id_end_to_end() -> None: """Updating a message by ID via DeltaChannel must persist and reload correctly.""" - from langchain_core.messages import HumanMessage - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9588,11 +9569,6 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: async def test_delta_channel_durability_exit_stores_snapshot() -> None: """DeltaChannel must reload from a durability='exit' checkpoint.""" - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] @@ -9620,14 +9596,6 @@ async def test_delta_channel_durability_exit_stores_snapshot() -> None: async def test_delta_channel_async_write_ordering() -> None: """In async mode, DeltaChannel write futures are awaited before the checkpoint is committed, so aput_writes always precedes aput for sentinel checkpoints.""" - from typing import Annotated - - from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.checkpoint.memory import InMemorySaver - - from langgraph.graph import START, StateGraph - from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]