diff --git a/.gitignore b/.gitignore index 145295a29..b154d3b56 100644 --- a/.gitignore +++ b/.gitignore @@ -100,3 +100,4 @@ dmypy.json .turbo .editorconfig .scratch +.worktrees/ diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 600127333..237fd8932 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -4,15 +4,17 @@ import threading from collections import defaultdict from collections.abc import Iterator, Sequence from contextlib import contextmanager -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, + _ChannelWritesHistory, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb from psycopg_pool import ConnectionPool from langgraph.checkpoint.postgres import _internal -from langgraph.checkpoint.postgres.base import BasePostgresSaver +from langgraph.checkpoint.postgres.base import ( + SELECT_DELTA_COMBINED_SQL, + BasePostgresSaver, + _DeltaCombinedRow, +) from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver Conn = _internal.Conn # For backward compatibility @@ -430,6 +436,48 @@ class PostgresSaver(BasePostgresSaver): with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur + def _get_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _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. + """ + 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, + ( + channel, + thread_id, + checkpoint_ns, + thread_id, + checkpoint_ns, + channel, + thread_id, + checkpoint_ns, + channel, + ), + ) + rows = cur.fetchall() + return self._build_delta_channel_writes_history( + channel=channel, + target_id=checkpoint_id, + rows=cast("list[_DeltaCombinedRow]", rows), + ) + def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: """ Convert a database row into a CheckpointTuple object. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5cd533f8d..965c8f31b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -4,15 +4,17 @@ import asyncio from collections import defaultdict from collections.abc import AsyncIterator, Iterator, Sequence from contextlib import asynccontextmanager -from typing import Any +from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, + _ChannelWritesHistory, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -23,7 +25,11 @@ from psycopg.types.json import Jsonb from psycopg_pool import AsyncConnectionPool from langgraph.checkpoint.postgres import _ainternal -from langgraph.checkpoint.postgres.base import BasePostgresSaver +from langgraph.checkpoint.postgres.base import ( + SELECT_DELTA_COMBINED_SQL, + BasePostgresSaver, + _DeltaCombinedRow, +) from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver Conn = _ainternal.Conn # For backward compatibility @@ -391,6 +397,46 @@ class AsyncPostgresSaver(BasePostgresSaver): async with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur + async def _aget_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _ChannelWritesHistory: + """Fast-path override of `BaseCheckpointSaver._aget_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; rows are assembled by the shared pure helper on + `BasePostgresSaver`. + """ + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id is None: + target = await self.aget_tuple(config) + if target is None: + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + checkpoint_id = target.config["configurable"]["checkpoint_id"] + async with self._cursor() as cur: + await cur.execute( + SELECT_DELTA_COMBINED_SQL, + ( + channel, + thread_id, + checkpoint_ns, + thread_id, + checkpoint_ns, + channel, + thread_id, + checkpoint_ns, + channel, + ), + ) + rows = await cur.fetchall() + return self._build_delta_channel_writes_history( + channel=channel, + target_id=checkpoint_id, + rows=cast("list[_DeltaCombinedRow]", rows), + ) + async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: """ Convert a database row into a CheckpointTuple object. diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 8f27fee5b..d5b963c87 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -4,13 +4,16 @@ import random import warnings from collections.abc import Sequence from importlib.metadata import version as get_version -from typing import Any, cast +from typing import Any, TypedDict, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, + PendingWrite, + _ChannelWritesHistory, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS @@ -153,6 +156,62 @@ INSERT_CHECKPOINT_WRITES_SQL = """ """ +class _DeltaCombinedRow(TypedDict, total=False): + """One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables). + + Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are + relevant for that kind; irrelevant columns are NULL and typed as `None`. + """ + + _kind: str # always present: "p", "w", or "b" + # checkpoint row ("p") + checkpoint_id: str | None + parent_checkpoint_id: str | None + ver: str | None + # write / blob rows ("w", "b") + type: str | None + blob: bytes | None + # write row only ("w") + task_id: str | None + idx: int | None + # blob row only ("b") + version: str | None + + +# 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, +# thread_id, checkpoint_ns, channel) +SELECT_DELTA_COMBINED_SQL = """ + SELECT 'p'::text AS _kind, + checkpoint_id, + parent_checkpoint_id, + checkpoint -> 'channel_versions' ->> %s AS ver, + NULL::text AS type, + NULL::bytea AS blob, + NULL::text AS task_id, + NULL::int AS idx, + NULL::text AS version + FROM checkpoints + WHERE thread_id = %s AND checkpoint_ns = %s + UNION ALL + SELECT 'w', + checkpoint_id, NULL, NULL, + type, blob, task_id, idx, NULL + FROM checkpoint_writes + WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s + UNION ALL + SELECT 'b', + NULL, NULL, NULL, + type, blob, NULL, NULL, version + FROM checkpoint_blobs + WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s +""" + + class BasePostgresSaver(BaseCheckpointSaver[str]): SELECT_SQL = SELECT_SQL SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL @@ -195,6 +254,83 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): if t.decode() != "empty" } + def _build_delta_channel_writes_history( + self, + *, + channel: str, + target_id: str, + rows: Sequence[_DeltaCombinedRow], + ) -> _ChannelWritesHistory: + """Reconstruct one delta channel's history from the combined UNION ALL rows. + + Pure data transform shared by sync (`PostgresSaver`) and async + (`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL` + and feed the tagged rows here. + + Walk is newest → oldest from the target's parent. A non-sentinel + blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the + walk and is returned as the seed so replay starts from it. + + Writes stored at `target_id` itself are pending writes for the next + step and are excluded — the walk begins at the target's parent. + """ + parent_of: dict[str, str | None] = {} + ver_of: dict[str, str | None] = {} + writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {} + blob_by_ver: dict[str, tuple[str, bytes]] = {} + + for r in rows: + kind = r["_kind"] + if kind == "p": + cid = cast(str, r["checkpoint_id"]) + parent_of[cid] = r["parent_checkpoint_id"] + ver_of[cid] = r["ver"] + elif kind == "w": + cid = cast(str, r["checkpoint_id"]) + writes_by_cid.setdefault(cid, []).append( + cast( + "tuple[str, bytes, str, int]", + (r["type"], r["blob"], r["task_id"], r["idx"]), + ) + ) + else: # kind == "b" + blob_by_ver[cast(str, r["version"])] = cast( + "tuple[str, bytes]", (r["type"], r["blob"]) + ) + + # 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) + + ancestors: list[str] = [] + cur_cid: str | None = parent_of.get(target_id) + while cur_cid is not None: + ancestors.append(cur_cid) + cur_cid = parent_of.get(cur_cid) + if not ancestors: + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[]) + + collected: list[PendingWrite] = [] # newest first; reversed at the end + for cid in ancestors: + # Collect writes first — they encode the transition FROM this + # ancestor's state to its child's and must be included even if + # this ancestor is also the seed checkpoint. + for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []): + val = self.serde.loads_typed((type_tag, write_blob)) + collected.append((task_id, channel, val)) + # Then check seed terminator. + ver = ver_of.get(cid) + if ver is not None: + seed_blob = blob_by_ver.get(ver) + if seed_blob is not None and seed_blob[0] != "empty": + blob_value = self.serde.loads_typed(seed_blob) + if blob_value is not DELTA_SENTINEL: + collected.reverse() + return _ChannelWritesHistory(seed=blob_value, writes=collected) + + collected.reverse() # oldest → newest + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + def _dump_blobs( self, thread_id: str, diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index f084f8b22..72044e179 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -12,7 +12,7 @@ readme = "README.md" license = "MIT" license-files = ['LICENSE'] dependencies = [ - "langgraph-checkpoint>=2.1.2,<5.0.0", + "langgraph-checkpoint>=4.0.3,<5.0.0", "orjson>=3.11.5", "psycopg>=3.2.0", "psycopg-pool>=3.2.0", diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index a46163395..fd42146cf 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -361,9 +361,9 @@ async def test_get_checkpoint_no_channel_values( load_checkpoint_tuple = saver._load_checkpoint_tuple - def patched_load_checkpoint_tuple(value): + async def patched_load_checkpoint_tuple(value): value["checkpoint"].pop("channel_values", None) - return load_checkpoint_tuple(value) + return await load_checkpoint_tuple(value) monkeypatch.setattr( saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple @@ -371,3 +371,47 @@ async def test_get_checkpoint_no_channel_values( checkpoint = await saver.aget_tuple(config) assert checkpoint.checkpoint["channel_values"] == {} + + +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +async def test_delta_channel_chain_reconstruction(saver_name: str) -> None: + """AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal.""" + pytest.importorskip( + "langgraph.channels.delta", reason="langgraph core not installed" + ) + + from typing import Annotated + + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.channels.delta import DeltaChannel + from langgraph.graph import START, StateGraph + from langgraph.graph.message import _messages_delta_reducer + from typing_extensions import TypedDict + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def respond(state: State) -> dict: + n = len(state["messages"]) + return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + + async with _saver(saver_name) as saver: + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "diff-channel-test-1"}} + + await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config) + await graph.ainvoke( + {"messages": [HumanMessage(content="there", id="h2")]}, config + ) + + state = await graph.aget_state(config) + msgs = state.values["messages"] + assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}" + assert msgs[0].content == "hi" + assert msgs[1].content == "reply-1" + assert msgs[2].content == "there" + assert msgs[3].content == "reply-3" diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 33495d5e3..891f7c95e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -3,7 +3,7 @@ from __future__ import annotations import copy import logging from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence -from typing import ( # noqa: UP035 +from typing import ( Any, Generic, Literal, @@ -18,6 +18,9 @@ from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.checkpoint.serde.types import ( + DELTA_SENTINEL as DELTA_SENTINEL, +) from langgraph.checkpoint.serde.types import ( ERROR, INTERRUPT, @@ -28,6 +31,8 @@ from langgraph.checkpoint.serde.types import ( V = TypeVar("V", int, float, str) PendingWrite = tuple[str, str, Any] + + logger = logging.getLogger(__name__) @@ -119,6 +124,30 @@ class CheckpointTuple(NamedTuple): pending_writes: list[PendingWrite] | None = None +class _ChannelWritesHistory(NamedTuple): + """Result of `BaseCheckpointSaver._get_channel_writes_history`. + + Storage-level view of what one channel wrote across the ancestor chain + of a target checkpoint: + + * `seed` — the nearest ancestor's stored blob value for this channel, + or `DELTA_SENTINEL` if the walk reached the root without finding a + stored value. A non-sentinel seed typically indicates a pre-delta + snapshot preserved across a channel-type migration (e.g. + `BinaryOperatorAggregate` storage extended under `DeltaChannel`). + * `writes` — on-path deltas oldest→newest, one `PendingWrite` per + step that wrote to this channel. Writes stored at the target + checkpoint itself are pending for the next super-step and are + excluded. + + Experimental: method surface may change; the NamedTuple shape is the + contract. + """ + + seed: Any + writes: list[PendingWrite] + + class BaseCheckpointSaver(Generic[V]): """Base class for creating a graph checkpointer. @@ -457,6 +486,104 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + def _get_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: + """Pure storage read used by `_get_channel_writes_history`. + + Must return the same value as `get_tuple` but must NOT trigger channel + 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) + + async def _aget_tuple_raw(self, config: RunnableConfig) -> CheckpointTuple | None: + """Async version of `_get_tuple_raw`. See docstring there.""" + return await self.aget_tuple(config) + + def _get_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _ChannelWritesHistory: + """**Experimental.** Query one channel's writes along the parent chain. + + Storage-level query, not channel semantics: returns `(seed, writes)` + reflecting what storage knows about a single channel across the + ancestor chain of the target checkpoint identified by `config`. + + * `writes` — on-path deltas oldest→newest as `PendingWrite` tuples. + Writes stored at the target `checkpoint_id` itself are pending + for the next super-step and are excluded. + * `seed` — the nearest ancestor's stored blob value for this + channel; `DELTA_SENTINEL` if the walk reached the root without + finding a stored value. A non-sentinel seed typically indicates + a pre-delta snapshot preserved across a channel-type migration. + + Walks the **parent chain** (not `list(before=...)`): for forked + threads, only on-path ancestors contribute. + + Reference implementation walks `get_tuple` + `parent_config`, + inspecting each ancestor's `channel_values[channel]` for the seed + terminator. Savers with direct storage access (`InMemorySaver`, + `PostgresSaver`) override for performance; the return contract is + fixed here. + + Underscore-prefixed because the method surface is experimental. + """ + collected: list[PendingWrite] = [] # newest first; reversed at the end + target_tuple = self._get_tuple_raw(config) + cursor_config: RunnableConfig | None = ( + target_tuple.parent_config if target_tuple else None + ) + while cursor_config is not None: + tup = self._get_tuple_raw(cursor_config) + if tup is None: + break + # Collect this ancestor's writes FIRST — they encode the + # transition from this ancestor's state to its child's, so + # they must be included whether or not this ancestor is the + # seed terminator. + if tup.pending_writes: + # Within a superstep, pending_writes are oldest→newest; + # reverse to scan newest-first. + for write in reversed(tup.pending_writes): + if write[1] != channel: + continue + collected.append(write) + # Seed terminator: any non-sentinel blob on an ancestor + # establishes the reconstruction base. Stop here. + ancestor_value = tup.checkpoint["channel_values"].get(channel) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + collected.reverse() + return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + cursor_config = tup.parent_config + collected.reverse() + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + + async def _aget_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _ChannelWritesHistory: + """Async version of `_get_channel_writes_history`. See docstring there.""" + collected: list[PendingWrite] = [] + target_tuple = await self._aget_tuple_raw(config) + cursor_config: RunnableConfig | None = ( + target_tuple.parent_config if target_tuple else None + ) + while cursor_config is not None: + tup = await self._aget_tuple_raw(cursor_config) + if tup is None: + break + if tup.pending_writes: + for write in reversed(tup.pending_writes): + if write[1] != channel: + continue + collected.append(write) + ancestor_value = tup.checkpoint["channel_values"].get(channel) + if ancestor_value is not None and ancestor_value is not DELTA_SENTINEL: + collected.reverse() + return _ChannelWritesHistory(seed=ancestor_value, writes=collected) + cursor_config = tup.parent_config + collected.reverse() + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + def get_next_version(self, current: V | None, channel: None) -> V: """Generate the next version ID for a channel. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index da353c2b4..4cf1f80fc 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -14,16 +14,20 @@ from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, + PendingWrite, SerializerProtocol, + _ChannelWritesHistory, get_checkpoint_id, get_checkpoint_metadata, ) +from langgraph.checkpoint.serde.types import _DeltaSnapshot logger = logging.getLogger(__name__) @@ -121,16 +125,114 @@ class InMemorySaver( return self.stack.__exit__(__exc_type, __exc_value, __traceback) def _load_blobs( - self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions + self, + thread_id: str, + checkpoint_ns: str, + versions: ChannelVersions, ) -> dict[str, Any]: - channel_values: dict[str, Any] = {} - for k, v in versions.items(): - kk = (thread_id, checkpoint_ns, k, v) - if kk in self.blobs: - vv = self.blobs[kk] - if vv[0] != "empty": - channel_values[k] = self.serde.loads_typed(vv) - return channel_values + result: dict[str, Any] = {} + for k, ver in versions.items(): + kk = (thread_id, checkpoint_ns, k, ver) + if kk not in self.blobs: + continue + vv = self.blobs[kk] + if vv[0] == "empty": + continue + result[k] = self.serde.loads_typed(vv) + return result + + def _get_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _ChannelWritesHistory: + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id", "") + ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {}) + # Walk the parent chain newest→oldest. Skip the target itself — + # writes stored AT `checkpoint_id` are pending for the next step + # (pregel applies them via `apply_writes`; they aren't part of the + # snapshot value AT `checkpoint_id`). + chain: list[str] = [] + target_entry = ns_storage.get(checkpoint_id) + current: str | None = target_entry[2] if target_entry is not None else None + while current is not None: + entry = ns_storage.get(current) + if entry is None: + break + chain.append(current) + _, _, parent = entry + current = parent + # Scan newest→oldest. A pre-delta blob on an ancestor terminates the + # walk and is bound as `seed`; without this, a thread migrated from + # pre-delta storage would replay ancestor writes all the way to the + # root AND miss any value that lived only in the old blob (e.g. from + # `update_state`). + # + # At each ancestor, check the blob BEFORE processing its pending + # writes: a pre-delta blob represents the state AT that ancestor, + # which already subsumes any writes stored under it. Processing + # those writes first would fold them into the reconstructed value + # twice (once via the blob, once via replay). + collected: list[PendingWrite] = [] # newest first + for cp_id in chain: # newest → oldest + entry = ns_storage.get(cp_id) + if entry is not None: + ckpt = self.serde.loads_typed(entry[0]) + ver = ckpt.get("channel_versions", {}).get(channel) + if ver is not None: + blob_entry = self.blobs.get( + (thread_id, checkpoint_ns, channel, ver) + ) + if blob_entry is not None and blob_entry[0] != "empty": + blob_value = self.serde.loads_typed(blob_entry) + if blob_value is not DELTA_SENTINEL: + if isinstance(blob_value, _DeltaSnapshot): + # Step-based snapshot: the blob is state AT this + # ancestor, but the ancestor's pending_writes + # encode the NEXT step's transition and are NOT + # subsumed by the snapshot — collect them first. + step_writes = self.writes.get( + (thread_id, checkpoint_ns, cp_id), {} + ) + for (_task_id, _idx), ( + tid, + ch, + serialized, + _, + ) in sorted(step_writes.items(), reverse=True): + if ch != channel: + continue + collected.append( + (tid, ch, self.serde.loads_typed(serialized)) + ) + collected.reverse() + return _ChannelWritesHistory( + seed=blob_value, writes=collected + ) + # Pre-delta blob: state AT this ancestor already + # subsumes its pending_writes — skip them. + collected.reverse() + return _ChannelWritesHistory( + seed=blob_value, writes=collected + ) + + step_writes = self.writes.get((thread_id, checkpoint_ns, cp_id), {}) + # Within a superstep, sorted by (task_id, idx) = oldest → newest; + # reverse for newest-first scan. + for (_task_id, _idx), (tid, ch, serialized, _) in sorted( + step_writes.items(), reverse=True + ): + if ch != channel: + continue + val = self.serde.loads_typed(serialized) + collected.append((tid, ch, val)) + collected.reverse() + return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected) + + async def _aget_channel_writes_history( + self, config: RunnableConfig, channel: str + ) -> _ChannelWritesHistory: + return self._get_channel_writes_history(config, channel) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the in-memory storage. diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 86b72759c..281db5406 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -33,14 +33,18 @@ from langchain_core.load.load import Reviver from langgraph.checkpoint.serde import _msgpack as _lg_msgpack from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.event_hooks import emit_serde_event -from langgraph.checkpoint.serde.types import SendProtocol +from langgraph.checkpoint.serde.types import ( + DELTA_SENTINEL, + SendProtocol, + _DeltaSentinel, + _DeltaSnapshot, +) from langgraph.store.base import Item if TYPE_CHECKING: from langgraph.checkpoint.serde._msgpack import ( AllowedMsgpackModules, ) - from langgraph.checkpoint.serde.types import SendProtocol LC_REVIVER = Reviver() EMPTY_BYTES = b"" @@ -317,10 +321,16 @@ EXT_METHOD_SINGLE_ARG = 3 EXT_PYDANTIC_V1 = 4 EXT_PYDANTIC_V2 = 5 EXT_NUMPY_ARRAY = 6 +EXT_DELTA_SNAPSHOT = 7 +EXT_DELTA_SENTINEL = 8 def _msgpack_default(obj: Any) -> str | ormsgpack.Ext: - if hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2 + if isinstance(obj, _DeltaSnapshot): + return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value)) + elif isinstance(obj, _DeltaSentinel): + return ormsgpack.Ext(EXT_DELTA_SENTINEL, b"") + elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2 return ormsgpack.Ext( EXT_PYDANTIC_V2, _msgpack_enc( @@ -646,7 +656,15 @@ def _create_msgpack_ext_hook( return False def ext_hook(code: int, data: bytes) -> Any: - if code == EXT_CONSTRUCTOR_SINGLE_ARG: + if code == EXT_DELTA_SENTINEL: + return DELTA_SENTINEL + elif code == EXT_DELTA_SNAPSHOT: + return _DeltaSnapshot( + ormsgpack.unpackb( + data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS + ) + ) + elif code == EXT_CONSTRUCTOR_SINGLE_ARG: try: tup = ormsgpack.unpackb( data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 65a2b0c8e..d68ec413a 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -1,6 +1,7 @@ from collections.abc import Sequence from typing import ( Any, + NamedTuple, Protocol, TypeVar, runtime_checkable, @@ -14,6 +15,39 @@ INTERRUPT = "__interrupt__" RESUME = "__resume__" TASKS = "__pregel_tasks" + +class _DeltaSentinel: + """Singleton marker stored (as zero bytes) in checkpoint_blobs for a + DeltaChannel field. The actual per-step writes live in checkpoint_writes + and are replayed through the reducer at load time. + + Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same + module-level instance. + """ + + __slots__ = () + + def __repr__(self) -> str: + return "DELTA_SENTINEL" + + +DELTA_SENTINEL = _DeltaSentinel() + + +class _DeltaSnapshot(NamedTuple): + """Snapshot blob for a DeltaChannel with finite snapshot_frequency. + + Stored in checkpoint_blobs via the `EXT_DELTA_SNAPSHOT` msgpack ext code. + The ancestor walk in `_get_channel_writes_history` terminates when it + encounters this type (any non-sentinel blob stops the walk). + + `from_checkpoint` reconstructs the channel value directly from `.value` + without replaying writes — the snapshot IS the accumulated state. + """ + + value: Any + + Value = TypeVar("Value", covariant=True) Update = TypeVar("Update", contravariant=True) C = TypeVar("C") diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index 15c54e6c0..ca4f2333f 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -1048,3 +1048,15 @@ def test_msgpack_nested_pydantic_serializes_as_dict( # No blocking should occur - inner is serialized as dict, not ext assert "blocked" not in caplog.text.lower() assert result == obj + + +def test_delta_sentinel_serde_round_trip() -> None: + from langgraph.checkpoint.base import DELTA_SENTINEL + from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + + serde = JsonPlusSerializer() + type_tag, blob = serde.dumps_typed(DELTA_SENTINEL) + assert type_tag == "msgpack" + assert blob # non-empty ext envelope + loaded = serde.loads_typed((type_tag, blob)) + assert loaded is DELTA_SENTINEL diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 4198f600c..617b23f68 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig from pydantic import BaseModel from langgraph.checkpoint.base import ( + DELTA_SENTINEL, Checkpoint, CheckpointMetadata, create_checkpoint, @@ -208,8 +209,6 @@ class TestMemorySaver: async def test_memory_saver() -> None: - from langgraph.checkpoint.memory import InMemorySaver - memory_saver = InMemorySaver() assert isinstance(memory_saver, InMemorySaver) @@ -320,3 +319,347 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None: assert direct is not None expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict() assert direct.checkpoint["channel_values"]["foo"] == expected + + +class TestInMemorySaverDeltaChannel: + def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None: + """_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred).""" + saver = InMemorySaver() + serde = JsonPlusSerializer() + + thread_id, ns, channel = "t1", "", "messages" + v1 = "00000000000000000000000000000001.0000000000000000" + + saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL) + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" + cp1["channel_versions"][channel] = v1 + saver.storage[thread_id][ns] = { + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + } + + result = saver._load_blobs(thread_id, ns, {channel: v1}) + assert channel in result + assert result[channel] is DELTA_SENTINEL + + def test_get_channel_writes_collects_ancestor_writes_only(self) -> None: + """_get_channel_writes_history collects ancestor writes oldest→newest, + and excludes writes stored at the target checkpoint itself (those are + pending writes for the next step, applied separately by pregel).""" + saver = InMemorySaver() + serde = JsonPlusSerializer() + + thread_id, ns, channel = "t1", "", "messages" + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" + cp2 = empty_checkpoint() + cp2["id"] = "cp2" + saver.storage[thread_id][ns] = { + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + "cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), + } + # Writes stored at cp1 produced the cp1 snapshot; part of history. + saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = ( + "task1", + channel, + serde.dumps_typed({"content": "hi"}), + "", + ) + # Writes stored at cp2 are pending — they will produce cp3 when the + # step that loaded cp2 completes. They MUST NOT appear in the + # reconstructed snapshot value at cp2. + saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = ( + "task2", + channel, + serde.dumps_typed({"content": "pending"}), + "", + ) + + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": "cp2", + } + } + result = saver._get_channel_writes_history(config, channel) + assert result.seed is DELTA_SENTINEL + values = [v for _, _, v in result.writes] + assert values == [{"content": "hi"}] + + def test_get_channel_writes_at_root_returns_empty(self) -> None: + """Reconstructing the root checkpoint's state: no ancestors → [].""" + saver = InMemorySaver() + serde = JsonPlusSerializer() + thread_id, ns, channel = "t1", "", "messages" + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" + saver.storage[thread_id][ns] = { + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + } + saver.writes[(thread_id, ns, "cp1")][("task1", 0)] = ( + "task1", + channel, + serde.dumps_typed({"content": "pending"}), + "", + ) + + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": "cp1", + } + } + result = saver._get_channel_writes_history(config, channel) + assert result.seed is DELTA_SENTINEL + assert result.writes == [] + + +class TestBaseFallbackGetChannelWrites: + """Exercises the `BaseCheckpointSaver._get_channel_writes_history` default + implementation — the path third-party savers inherit when they don't + override `_get_channel_writes_history` themselves. + + Regression guard for a bug where the fallback passed the caller's config + (with `checkpoint_id`) straight to `self.list()`, which most savers + collapse to a single row — causing the fallback to return `[]`. + """ + + def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]: + """Build an InMemorySaver with a 3-checkpoint chain and per-step writes + for a `messages` channel. + + Returns `(saver, thread_id, namespace)`. The saver subclass deletes the + InMemorySaver override so the base class fallback is exercised. + """ + + class _ThirdPartyStyleSaver(InMemorySaver): + _get_channel_writes_history = ( + InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] + ) + _aget_channel_writes_history = ( + InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + ) + + saver = _ThirdPartyStyleSaver() + serde = JsonPlusSerializer() + thread_id, ns, channel = "t1", "", "messages" + + cp0 = empty_checkpoint() + cp0["id"] = "00000000000000000000000000000001.0000000000000000" + cp1 = empty_checkpoint() + cp1["id"] = "00000000000000000000000000000002.0000000000000000" + cp2 = empty_checkpoint() + cp2["id"] = "00000000000000000000000000000003.0000000000000000" + saver.storage[thread_id][ns] = { + cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None), + cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]), + cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]), + } + # Writes under cp0 produced cp1's state; writes under cp1 produced cp2's. + saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = ( + "task1", + channel, + serde.dumps_typed({"content": "first"}), + "", + ) + saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = ( + "task2", + channel, + serde.dumps_typed({"content": "second"}), + "", + ) + return saver, thread_id, ns + + def test_fallback_returns_ancestor_writes_oldest_first(self) -> None: + saver, thread_id, ns = self._build_saver_with_chain() + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + result = saver._get_channel_writes_history(config, "messages") + + assert result.seed is DELTA_SENTINEL + values = [v for _, _, v in result.writes] + assert values == [{"content": "first"}, {"content": "second"}] + + async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None: + saver, thread_id, ns = self._build_saver_with_chain() + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + result = await saver._aget_channel_writes_history(config, "messages") + + assert result.seed is DELTA_SENTINEL + values = [v for _, _, v in result.writes] + assert values == [{"content": "first"}, {"content": "second"}] + + async def test_async_fallback_concurrent_tasks_do_not_interfere(self) -> None: + """Regression: the re-entrancy guard must be task-local, not thread-local. + + Two concurrent `_aget_channel_writes_history` calls on the same + event-loop thread must each see their full reconstructed writes. A + `threading.local()` guard would let whichever task set it first + short-circuit the other to `writes=[]`. + """ + import asyncio + + saver, thread_id, ns = self._build_saver_with_chain() + + # Force the two tasks to interleave across the `set(True)` boundary: + # each `aget_tuple` yields control, so if the guard were thread-local + # the second task would observe `active=True` set by the first. + orig_aget_tuple = saver.aget_tuple + + async def slow_aget_tuple(config: RunnableConfig) -> Any: + await asyncio.sleep(0) + return await orig_aget_tuple(config) + + saver.aget_tuple = slow_aget_tuple # type: ignore[method-assign] + + target_id = "00000000000000000000000000000003.0000000000000000" + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target_id, + } + } + + results = await asyncio.gather( + saver._aget_channel_writes_history(config, "messages"), + saver._aget_channel_writes_history(config, "messages"), + ) + + expected_values = [{"content": "first"}, {"content": "second"}] + for result in results: + assert result.seed is DELTA_SENTINEL + values = [v for _, _, v in result.writes] + assert values == expected_values + + +class TestPreDeltaBlobTerminator: + """Verify the pre-delta blob terminator: when the ancestor walk hits a + checkpoint whose blob for the channel is a real value (not + DELTA_SENTINEL), reconstruction seeds from it and stops. This guards + + * back-compat: a thread written by pre-delta code, then extended under + delta — reconstruction must return the correct value without walking + past the last pre-delta ancestor; + * perf: without the terminator, every reconstruct-after-migration would + walk all the way to the thread root. + """ + + def _build_mixed_thread(self) -> tuple[InMemorySaver, str, str, str, str]: + """Three-checkpoint chain: cp1 (pre-delta, blob=[A]), cp2 (delta, + write=B), cp3 (delta, write=C). Reconstructing at cp3 must yield + seed=[A] + writes=[B, C]. + + Returns `(saver, thread_id, ns, channel, cp3_id)`. + """ + saver = InMemorySaver() + serde = JsonPlusSerializer() + thread_id, ns, channel = "t1", "", "messages" + + v1 = "00000000000000000000000000000001.0" + v2 = "00000000000000000000000000000002.0" + v3 = "00000000000000000000000000000003.0" + + # Pre-delta: cp1 stored a real blob for the channel. + saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"]) + # Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes. + saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL) + saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL) + + cp1 = empty_checkpoint() + cp1["id"] = "cp1" + cp1["channel_versions"][channel] = v1 + cp2 = empty_checkpoint() + cp2["id"] = "cp2" + cp2["channel_versions"][channel] = v2 + cp3 = empty_checkpoint() + cp3["id"] = "cp3" + cp3["channel_versions"][channel] = v3 + + saver.storage[thread_id][ns] = { + "cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None), + "cp2": (serde.dumps_typed(cp2), serde.dumps_typed({}), "cp1"), + "cp3": (serde.dumps_typed(cp3), serde.dumps_typed({}), "cp2"), + } + # Write under cp1 would be from the pre-delta era and MUST be ignored + # (the blob already captures it). We add one and assert it is not + # folded into the reconstructed result. + saver.writes[(thread_id, ns, "cp1")][("task0", 0)] = ( + "task0", + channel, + serde.dumps_typed("PRE-DELTA-WRITE"), + "", + ) + saver.writes[(thread_id, ns, "cp2")][("task2", 0)] = ( + "task2", + channel, + serde.dumps_typed("B"), + "", + ) + saver.writes[(thread_id, ns, "cp3")][("task3", 0)] = ( + "task3", + channel, + serde.dumps_typed("PENDING-AT-TARGET"), + "", + ) + return saver, thread_id, ns, channel, "cp3" + + def test_seed_from_pre_delta_ancestor_blob(self) -> None: + saver, thread_id, ns, channel, target = self._build_mixed_thread() + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target, + } + } + + result = saver._get_channel_writes_history(config, channel) + + # Seed came from the pre-delta blob at cp1. + assert result.seed == ["A"] + # Delta-era writes from cp2 replay through the reducer on top of seed. + # cp3 is the target — its own write is pending for the NEXT step and + # must be excluded. + values = [v for _, _, v in result.writes] + assert values == ["B"] + + def test_pre_delta_blob_terminates_walk_before_older_writes(self) -> None: + """Writes stored at the pre-delta ancestor itself must not be replayed + (the blob subsumes them).""" + saver, thread_id, ns, channel, target = self._build_mixed_thread() + config: RunnableConfig = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": ns, + "checkpoint_id": target, + } + } + + result = saver._get_channel_writes_history(config, channel) + + values = [v for _, _, v in result.writes] + # The pre-delta write under cp1 must not appear (the blob subsumes it). + assert "PRE-DELTA-WRITE" not in values + # And the pending write at the target is never folded in. + assert "PENDING-AT-TARGET" not in values diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index a69c230b5..cc7fcc471 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,6 +1,7 @@ from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue, LastValueAfterFinish from langgraph.channels.named_barrier_value import ( @@ -20,6 +21,7 @@ __all__ = ( "UntrackedValue", "EphemeralValue", "BinaryOperatorAggregate", + "DeltaChannel", "NamedBarrierValue", "NamedBarrierValueAfterFinish", # topics diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index b33ea52fa..5735ac65d 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -22,10 +22,9 @@ __all__ = ("BinaryOperatorAggregate",) def _strip_extras(t): # type: ignore[no-untyped-def] """Strips Annotated, Required and NotRequired from a given type.""" if hasattr(t, "__origin__"): + if t.__origin__ in (Required, NotRequired): + return _strip_extras(t.__args__[0]) return _strip_extras(t.__origin__) - if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired): - return _strip_extras(t.__args__[0]) - return t @@ -33,11 +32,22 @@ def _get_overwrite(value: Any) -> tuple[bool, Any]: """Inspects the given value and returns (is_overwrite, overwrite_value).""" if isinstance(value, Overwrite): return True, value.value - if isinstance(value, dict) and set(value.keys()) == {OVERWRITE}: + if isinstance(value, dict) and len(value) == 1 and OVERWRITE in value: return True, value[OVERWRITE] return False, None +def _operators_equal(a: Callable, b: Callable) -> bool: + """Return True if two reducer operators should be considered equal. + + Lambdas all share the name '' so identity comparison is + unreliable; treat any pairing that includes a lambda as equal. + """ + if a.__name__ == "" or b.__name__ == "": + return True + return a is b + + class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): """Stores the result of applying a binary operator to the current value and each new value. @@ -68,11 +78,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): self.value = MISSING def __eq__(self, value: object) -> bool: - return isinstance(value, BinaryOperatorAggregate) and ( - value.operator is self.operator - if value.operator.__name__ != "" - and self.operator.__name__ != "" - else True + return isinstance(value, BinaryOperatorAggregate) and _operators_equal( + self.operator, value.operator ) @property diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py new file mode 100644 index 000000000..ed08f083f --- /dev/null +++ b/libs/langgraph/langgraph/channels/delta.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import collections.abc +import copy as _copy +from collections.abc import Callable, Sequence +from typing import Any, Generic + +from langgraph.checkpoint.base import DELTA_SENTINEL, PendingWrite +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from typing_extensions import Self + +from langgraph._internal._typing import MISSING +from langgraph.channels.base import BaseChannel, Value +from langgraph.channels.binop import _get_overwrite, _operators_equal, _strip_extras +from langgraph.errors import ( + EmptyChannelError, + ErrorCode, + InvalidUpdateError, + create_error_message, +) + +__all__ = ("DeltaChannel",) + + +class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): + """Reducer channel that stores only a sentinel in checkpoint blobs and + reconstructs state by replaying ancestor writes through the reducer. + + The reducer receives the current accumulated value and a batch of writes + in one call: `reducer(state, [write1, write2, ...]) -> new_state`. + + Reducers must be deterministic and batching-invariant (associative across + folds): applying two consecutive write batches separately must produce the + same state as applying their concatenation once: + + reducer(reducer(state, xs), ys) == reducer(state, xs + ys) + + This lets LangGraph replay checkpointed writes in larger batches than they + were originally produced without changing reconstructed state. + + `snapshot_frequency=None` (default): pure delta; stores only + `DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes. + + `snapshot_frequency=N`: `create_checkpoint` writes a full `_DeltaSnapshot` + blob every N steps, bounding replay depth to N. + + Parameters: + reducer: `(state, list[writes]) -> new_state`. Must be deterministic + and batching-invariant as described above. + typ: The value type (e.g. `list`, `dict`). Inferred automatically + from the outer type when used inside `Annotated[T, DeltaChannel(...)]`. + snapshot_frequency: Every Nth pregel step writes a snapshot blob. + `None` (default) = pure delta, never snapshot. + """ + + __slots__ = ("value", "reducer", "snapshot_frequency") + value: Value | Any + + def __init__( + self, + reducer: Callable[[Any, Sequence[Any]], Any], + typ: type[Value] | None = None, + *, + snapshot_frequency: int | None = None, + ) -> None: + if typ is None: + typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel + super().__init__(typ) + self.reducer = reducer + self.snapshot_frequency = snapshot_frequency + typ = _strip_extras(typ) + if typ in (collections.abc.Sequence, collections.abc.MutableSequence): + typ = list + if typ in (collections.abc.Set, collections.abc.MutableSet): + typ = set + if typ in (collections.abc.Mapping, collections.abc.MutableMapping): + typ = dict + self.typ = typ + self.value: Any = MISSING + + def __eq__(self, other: object) -> bool: + if not isinstance(other, DeltaChannel): + return False + if self.snapshot_frequency != other.snapshot_frequency: + return False + return _operators_equal(self.reducer, other.reducer) + + @property + def ValueType(self) -> Any: + return self.typ + + @property + def UpdateType(self) -> Any: + return self.typ + + def is_snapshot_step(self, step: int) -> bool: + """True if pregel should write a snapshot blob at this step.""" + return ( + self.snapshot_frequency is not None + and step > 0 + and step % self.snapshot_frequency == 0 + ) + + def copy(self) -> Self: + new = self.__class__( + self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency + ) + new.key = self.key + new.value = self.value if self.value is MISSING else _copy.copy(self.value) + return new + + def from_checkpoint(self, checkpoint: Any) -> Self: + """Initialize from a stored blob or sentinel. + + Blob types (dispatched via serde ext code, not dict key inspection): + * `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes. + * `_DeltaSnapshot(value)`: restore value directly from snapshot. + * plain value (migration from old BinOp blobs): use directly. + """ + new = self.__class__( + self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency + ) + new.key = self.key + if checkpoint is MISSING or checkpoint is DELTA_SENTINEL: + new.value = self.typ() + elif isinstance(checkpoint, _DeltaSnapshot): + new.value = checkpoint.value + else: + new.value = checkpoint + return new + + def replay_writes(self, writes: Sequence[PendingWrite]) -> None: + """Apply ancestor writes oldest-to-newest via a single reducer call. + + If any write is an Overwrite, the last one in the sequence acts as + the reset point: its value becomes the new base and only writes + after it are passed to the reducer. + """ + values = [v for _, _, v in writes] + if not values: + return + base = self.value + start = 0 + for i, v in enumerate(values): + is_ow, ow_value = _get_overwrite(v) + if is_ow: + base = _copy.copy(ow_value) if ow_value is not None else self.typ() + start = i + 1 + remaining = values[start:] + self.value = self.reducer(base, remaining) if remaining else base + + def update(self, values: Sequence[Any]) -> bool: + if not values: + return False + overwrite_idx: int | None = None + for i, v in enumerate(values): + is_ow, _ = _get_overwrite(v) + if is_ow: + if overwrite_idx is not None: + msg = create_error_message( + message="Can receive only one Overwrite value per super-step.", + error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE, + ) + raise InvalidUpdateError(msg) + overwrite_idx = i + if overwrite_idx is not None: + _, overwrite_value = _get_overwrite(values[overwrite_idx]) + base = ( + _copy.copy(overwrite_value) + if overwrite_value is not None + else self.typ() + ) + remaining = [v for i, v in enumerate(values) if i != overwrite_idx] + self.value = self.reducer(base, remaining) if remaining else base + return True + base = self.typ() if self.value is MISSING else self.value + self.value = self.reducer(base, list(values)) + return True + + def get(self) -> Any: + if self.value is MISSING: + raise EmptyChannelError() + return self.value + + def is_available(self) -> bool: + return self.value is not MISSING + + def checkpoint(self) -> Any: + """Return stored representation: always `DELTA_SENTINEL`. + + Snapshot decisions are made by `create_checkpoint` in pregel (which + has the step number) via `is_snapshot_step`. `checkpoint()` is only + called for non-snapshot steps or when no checkpointer is available. + """ + if self.value is MISSING: + return MISSING + return DELTA_SENTINEL diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 18c96d436..3018d21da 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -244,6 +244,52 @@ def add_messages( return merged +def _messages_delta_reducer( + state: list[AnyMessage], writes: list[list[AnyMessage]] +) -> list[AnyMessage]: + """**Experimental.** Batch reducer for use with `DeltaChannel`. + + Processes all writes in one pass — dedup by ID, `RemoveMessage` + tombstoning — without calling `add_messages`. Assumes writes contain + already-typed `BaseMessage` objects (no raw-dict coercion). + + This reducer is batching-invariant, as required by `DeltaChannel`: + `reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`. + + Use `add_messages` as the reducer for `BinaryOperatorAggregate` or + anywhere raw message dicts / strings need to be coerced first. + + Example:: + + from typing import Annotated + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import _messages_delta_reducer + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + """ + from itertools import chain + + index: dict[str, int] = {m.id: i for i, m in enumerate(state) if m.id is not None} + result: list[AnyMessage | None] = list(state) + for msg in chain.from_iterable( + [w] if isinstance(w, BaseMessage) else w for w in writes + ): + mid = msg.id + if mid is None: + result.append(msg) + elif isinstance(msg, RemoveMessage): + if mid in index: + result[index[mid]] = None + del index[mid] + elif mid in index: + result[index[mid]] = msg + else: + index[mid] = len(result) + result.append(msg) + return [m for m in result if m is not None] + + @deprecated( "MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.", category=None, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 59704c607..63f344c94 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -50,6 +50,7 @@ from langgraph._internal._timeout import coerce_timeout_policy from langgraph._internal._typing import EMPTY_SEQ, MISSING, DeprecatedKwargs from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue, LastValueAfterFinish from langgraph.channels.named_barrier_value import ( @@ -1110,6 +1111,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): CompiledStateGraph: The compiled `StateGraph`. """ checkpointer = ensure_valid_checkpointer(checkpointer) + serde_allowlist: set[tuple[str, ...]] | None = None if _serde.STRICT_MSGPACK_ENABLED: schema_types: list[type[Any]] = [ @@ -1697,6 +1699,20 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None: # Search through all annotated medata to find channel annotations for item in meta: if isinstance(item, BaseChannel): + if isinstance(item, DeltaChannel) and hasattr(typ, "__origin__"): + origin = typ.__origin__ + # Unwrap parameterized Required[X]/NotRequired[X] to X + # (e.g. Annotated[NotRequired[dict[...]], ...]). + if hasattr(origin, "__origin__") and origin.__origin__ in ( + Required, + NotRequired, + ): + origin = origin.__args__[0] + item = item.__class__( + item.reducer, + origin, + snapshot_frequency=item.snapshot_frequency, + ) return item elif isclass(item) and issubclass(item, BaseChannel): # ex, Annotated[int, EphemeralValue, SomeOtherAnnotation] diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 3d510ac7d..dc7675975 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -1,17 +1,23 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from datetime import datetime, timezone +from typing import Any, cast -from langgraph.checkpoint.base import Checkpoint +from langchain_core.runnables import RunnableConfig +from langgraph.checkpoint.base import DELTA_SENTINEL, BaseCheckpointSaver, Checkpoint from langgraph.checkpoint.base.id import uuid6 +from langgraph.checkpoint.serde.types import _DeltaSnapshot from langgraph._internal._typing import MISSING from langgraph.channels.base import BaseChannel +from langgraph.channels.delta import DeltaChannel from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec LATEST_VERSION = 4 +GetNextVersion = Callable[[Any, None], Any] + def empty_checkpoint() -> Checkpoint: return Checkpoint( @@ -31,35 +37,87 @@ def create_checkpoint( *, id: str | None = None, updated_channels: set[str] | None = None, + get_next_version: GetNextVersion | None = None, + force_delta_snapshot: bool = False, ) -> Checkpoint: - """Create a checkpoint for the given channels.""" + """Create a checkpoint for the given channels. + + For `DeltaChannel` with `snapshot_frequency=N`, snapshot steps write a + `_DeltaSnapshot` blob rather than `DELTA_SENTINEL`, bounding the ancestor + walk to at most N steps. Snapshots are eager: even if the channel had no + write this step, a version bump is forced (via `get_next_version`) so the + blob is stored by `put()`. Without `get_next_version` (e.g. static + contexts), snapshot steps gracefully fall back to sentinel. + + `force_delta_snapshot` writes available `DeltaChannel` values as snapshots + regardless of `snapshot_frequency`. This is used by `durability="exit"`, + where intermediate writes are not stored as ancestor `checkpoint_writes`. + """ ts = datetime.now(timezone.utc).isoformat() if channels is None: values = checkpoint["channel_values"] + channel_versions = checkpoint["channel_versions"] else: values = {} + channel_versions = dict(checkpoint["channel_versions"]) for k in channels: - if k not in checkpoint["channel_versions"]: + if k not in channel_versions: continue - v = channels[k].checkpoint() - if v is not MISSING: - values[k] = v + ch = channels[k] + if ( + isinstance(ch, DeltaChannel) + and (force_delta_snapshot or ch.is_snapshot_step(step)) + and ch.is_available() + ): + # Eager snapshot: bump version if not already written this step + # so put() includes this channel in new_versions and stores blob. + if get_next_version is not None and ( + updated_channels is None or k not in updated_channels + ): + channel_versions[k] = get_next_version(channel_versions[k], None) + values[k] = _DeltaSnapshot(ch.get()) + else: + v = ch.checkpoint() + if v is not MISSING: + values[k] = v return Checkpoint( v=LATEST_VERSION, ts=ts, id=id or str(uuid6(clock_seq=step)), channel_values=values, - channel_versions=checkpoint["channel_versions"], + channel_versions=channel_versions, versions_seen=checkpoint["versions_seen"], updated_channels=None if updated_channels is None else sorted(updated_channels), ) +def _needs_replay(spec: BaseChannel, stored: object) -> bool: + """True if `spec` is a `DeltaChannel` and the stored blob is a sentinel, + requiring an ancestor walk to reconstruct. + + `_DeltaSnapshot` blobs and plain values (migration) resolve directly via + `from_checkpoint` — only `DELTA_SENTINEL` / `MISSING` trigger replay. + """ + if not isinstance(spec, DeltaChannel): + return False + return stored is MISSING or stored is DELTA_SENTINEL + + def channels_from_checkpoint( specs: Mapping[str, BaseChannel | ManagedValueSpec], checkpoint: Checkpoint, + *, + saver: BaseCheckpointSaver | None = None, + config: RunnableConfig | None = None, ) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]: - """Get channels from a checkpoint.""" + """Hydrate channels from a checkpoint. + + For most channels, `spec.from_checkpoint(checkpoint["channel_values"][k])` + is sufficient. `DeltaChannel` is the exception: sentinel blobs require an + ancestor walk via `saver._get_channel_writes_history`. The walk terminates + at the nearest `_DeltaSnapshot` blob (step-based) or a pre-migration plain + value, so read depth is bounded by `snapshot_frequency`. + """ channel_specs: dict[str, BaseChannel] = {} managed_specs: dict[str, ManagedValueSpec] = {} for k, v in specs.items(): @@ -67,13 +125,53 @@ def channels_from_checkpoint( channel_specs[k] = v else: managed_specs[k] = v - return ( - { - k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - for k, v in channel_specs.items() - }, - managed_specs, - ) + + channels: dict[str, BaseChannel] = {} + for k, spec in channel_specs.items(): + ch: BaseChannel + stored = checkpoint["channel_values"].get(k, MISSING) + if _needs_replay(spec, stored) and saver is not None and config is not None: + delta_spec = cast(DeltaChannel, spec) + history = saver._get_channel_writes_history(config, k) + replay_ch = delta_spec.from_checkpoint(history.seed) + replay_ch.replay_writes(history.writes) + ch = replay_ch + else: + ch = spec.from_checkpoint(stored) + channels[k] = ch + return channels, managed_specs + + +async def achannels_from_checkpoint( + specs: Mapping[str, BaseChannel | ManagedValueSpec], + checkpoint: Checkpoint, + *, + saver: BaseCheckpointSaver | None = None, + config: RunnableConfig | None = None, +) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]: + """Async version of `channels_from_checkpoint`. See docstring there.""" + channel_specs: dict[str, BaseChannel] = {} + managed_specs: dict[str, ManagedValueSpec] = {} + for k, v in specs.items(): + if isinstance(v, BaseChannel): + channel_specs[k] = v + else: + managed_specs[k] = v + + channels: dict[str, BaseChannel] = {} + for k, spec in channel_specs.items(): + ch: BaseChannel + stored = checkpoint["channel_values"].get(k, MISSING) + if _needs_replay(spec, stored) and saver is not None and config is not None: + delta_spec = cast(DeltaChannel, spec) + history = await saver._aget_channel_writes_history(config, k) + replay_ch = delta_spec.from_checkpoint(history.seed) + replay_ch.replay_writes(history.writes) + ch = replay_ch + else: + ch = spec.from_checkpoint(stored) + channels[k] = ch + return channels, managed_specs def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 976c1382b..e6aea17d8 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -68,6 +68,7 @@ from langgraph.callbacks import ( GraphResumeEvent, ) from langgraph.channels.base import BaseChannel +from langgraph.channels.delta import DeltaChannel from langgraph.channels.untracked_value import UntrackedValue from langgraph.constants import TAG_HIDDEN from langgraph.errors import ( @@ -92,6 +93,7 @@ from langgraph.pregel._algo import ( task_path_str, ) from langgraph.pregel._checkpoint import ( + achannels_from_checkpoint, channels_from_checkpoint, copy_checkpoint, create_checkpoint, @@ -188,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 @@ -406,7 +410,7 @@ class PregelLoop: task = self.tasks.get(task_id) else: task = None - self.submit( + fut = self.submit( self.checkpointer_put_writes, config, writes_to_save, @@ -414,12 +418,16 @@ class PregelLoop: task_path_str(task.path) if task else "", ) else: - self.submit( + fut = self.submit( self.checkpointer_put_writes, config, writes_to_save, task_id, ) + 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) # output writes if hasattr(self, "tasks"): self.output_writes(task_id, writes) @@ -890,6 +898,10 @@ class PregelLoop: self.step, id=self.checkpoint["id"] if exiting else None, updated_channels=self.updated_channels, + get_next_version=self.checkpointer_get_next_version + if do_checkpoint + else None, + force_delta_snapshot=exiting and self.durability == "exit", ) # sanitize TASK channel in the checkpoint before saving (durability=="exit") if TASKS in self.checkpoint["channel_values"] and any( @@ -1273,7 +1285,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): ) self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) self.channels, self.managed = channels_from_checkpoint( - self.specs, self.checkpoint + self.specs, + self.checkpoint, + saver=self.checkpointer, + config=self.checkpoint_config, ) self.stack.push(self._suppress_interrupt) self.status = "input" @@ -1368,6 +1383,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): metadata: CheckpointMetadata, new_versions: ChannelVersions, ) -> RunnableConfig: + # Drain DeltaChannel write futures before committing the checkpoint so + # DELTA_SENTINEL blobs are never saved ahead of their backing writes. + if self._delta_write_futs: + futs, self._delta_write_futs = self._delta_write_futs, [] + await asyncio.gather(*futs) try: if prev is not None: await prev @@ -1473,11 +1493,15 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if saved.pending_writes is not None else [] ) + self._delta_write_futs = [] self.submit = await self.stack.enter_async_context( AsyncBackgroundExecutor(self.config) ) - self.channels, self.managed = channels_from_checkpoint( - self.specs, self.checkpoint + self.channels, self.managed = await achannels_from_checkpoint( + self.specs, + self.checkpoint, + saver=self.checkpointer, + config=self.checkpoint_config, ) self.stack.push(self._suppress_interrupt) self.status = "input" diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index feedff3f0..d4a22ae05 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -125,6 +125,7 @@ from langgraph.pregel._algo import ( ) from langgraph.pregel._call import identifier from langgraph.pregel._checkpoint import ( + achannels_from_checkpoint, channels_from_checkpoint, copy_checkpoint, create_checkpoint, @@ -1140,6 +1141,10 @@ class Pregel( channels, managed = channels_from_checkpoint( self.channels, saved.checkpoint, + saver=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + config=saved.config, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1256,9 +1261,13 @@ class Pregel( step = saved.metadata.get("step", -1) + 1 stop = step + 2 - channels, managed = channels_from_checkpoint( + channels, managed = await achannels_from_checkpoint( self.channels, saved.checkpoint, + saver=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + config=saved.config, ) # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1629,6 +1638,11 @@ class Pregel( channels, managed = channels_from_checkpoint( self.channels, checkpoint, + saver=self.checkpointer + if saved is not None + and isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + config=saved.config if saved is not None else None, ) values, as_node = updates[0][:2] @@ -2072,9 +2086,14 @@ class Pregel( ) if saved: checkpoint_config = patch_configurable(config, saved.config[CONF]) - channels, managed = channels_from_checkpoint( + channels, managed = await achannels_from_checkpoint( self.channels, checkpoint, + saver=self.checkpointer + if saved is not None + and isinstance(self.checkpointer, BaseCheckpointSaver) + else None, + config=saved.config if saved is not None else None, ) values, as_node = updates[0][:2] # no values, just clear all tasks diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index b6bfa54c0..d720c0b3b 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ ] dependencies = [ "langchain-core>=1.3.2,<2", - "langgraph-checkpoint>=2.1.0,<5.0.0", + "langgraph-checkpoint>=4.0.3,<5.0.0", "langgraph-sdk>=0.3.0,<0.4.0", "langgraph-prebuilt>=1.0.12,<1.1.0", "xxhash>=3.5.0", diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 37a5eb432..f967cbdbd 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -1,18 +1,34 @@ import operator from collections.abc import Sequence +from typing import Annotated import pytest +from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage +from langgraph.checkpoint.base import DELTA_SENTINEL +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from typing_extensions import NotRequired, TypedDict from langgraph._internal._typing import MISSING from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.errors import EmptyChannelError, InvalidUpdateError +from langgraph.graph import START, StateGraph +from langgraph.graph.message import _messages_delta_reducer +from langgraph.graph.state import _get_channel +from langgraph.types import Overwrite pytestmark = pytest.mark.anyio +# --------------------------------------------------------------------------- +# Core channel primitives +# --------------------------------------------------------------------------- + + def test_last_value() -> None: channel = LastValue(int).from_checkpoint(MISSING) assert channel.ValueType is int @@ -95,25 +111,543 @@ def test_untracked_value() -> None: assert channel.ValueType is dict assert channel.UpdateType is dict - # UntrackedValue should start empty with pytest.raises(EmptyChannelError): channel.get() - # Should be able to update with a value test_data = {"session": "test", "temp": "dir"} channel.update([test_data]) assert channel.get() == test_data - # Update with new value new_data = {"session": "updated", "temp": "newdir"} channel.update([new_data]) assert channel.get() == new_data - # On checkpoint, UntrackedValue should return MISSING checkpoint = channel.checkpoint() assert checkpoint is MISSING - # Creating from checkpoint with MISSING should start empty new_channel = UntrackedValue(dict).from_checkpoint(checkpoint) with pytest.raises(EmptyChannelError): new_channel.get() + + +# --------------------------------------------------------------------------- +# DeltaChannel — message reducer +# --------------------------------------------------------------------------- + + +def test_delta_channel_basic_two_steps() -> None: + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) + + ch.update([HumanMessage(content="hi", id="h1")]) + d1 = ch.checkpoint() + assert d1 is DELTA_SENTINEL + + ch.update([AIMessage(content="hello", id="a1")]) + d2 = ch.checkpoint() + assert d2 is DELTA_SENTINEL + + assert len(ch.get()) == 2 + assert ch.get()[0].content == "hi" + assert ch.get()[1].content == "hello" + + +def test_delta_channel_from_checkpoint_writes_list() -> None: + """replay_writes on a fresh channel replays through the operator.""" + spec = DeltaChannel(_messages_delta_reducer, list) + ch = spec.from_checkpoint(DELTA_SENTINEL) + ch.replay_writes( + [ + ("t0", "messages", HumanMessage(content="hi", id="h1")), + ("t1", "messages", AIMessage(content="hello", id="a1")), + ("t2", "messages", HumanMessage(content="bye", id="h2")), + ] + ) + msgs = ch.get() + assert len(msgs) == 3 + assert msgs[0].content == "hi" + assert msgs[1].content == "hello" + assert msgs[2].content == "bye" + + +def test_delta_channel_from_checkpoint_backwards_compat() -> None: + spec = DeltaChannel(_messages_delta_reducer, list) + old_value = [HumanMessage(content="old", id="h1")] + ch = spec.from_checkpoint(old_value) + assert ch.get() == old_value + + +def test_delta_channel_overwrite() -> None: + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) + ch.update([HumanMessage(content="old", id="h1")]) + + ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) + d = ch.checkpoint() + assert d is DELTA_SENTINEL + assert len(ch.get()) == 1 + assert ch.get()[0].content == "new" + + +def test_delta_channel_remove_message_and_replay() -> None: + """RemoveMessage must round-trip correctly when writes are replayed.""" + spec = DeltaChannel(_messages_delta_reducer, list) + ch = spec.from_checkpoint(MISSING) + + ch.update([HumanMessage(content="hi", id="h1")]) + ch.update([AIMessage(content="hello", id="a1")]) + assert ch.get() == [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + ] + + ch.update([RemoveMessage(id="a1")]) + assert ch.get() == [HumanMessage(content="hi", id="h1")] + + ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2.replay_writes( + [ + ("t0", "messages", HumanMessage(content="hi", id="h1")), + ("t1", "messages", AIMessage(content="hello", id="a1")), + ("t2", "messages", RemoveMessage(id="a1")), + ] + ) + assert ch2.get() == [HumanMessage(content="hi", id="h1")] + + +def test_delta_channel_update_by_id_and_replay() -> None: + """Updating a message by ID must round-trip correctly through writes replay.""" + spec = DeltaChannel(_messages_delta_reducer, list) + ch = spec.from_checkpoint(MISSING) + + ch.update([HumanMessage(content="original", id="h1")]) + ch.update([HumanMessage(content="updated", id="h1")]) + assert ch.get() == [HumanMessage(content="updated", id="h1")] + + ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2.replay_writes( + [ + ("t0", "messages", HumanMessage(content="original", id="h1")), + ("t1", "messages", HumanMessage(content="updated", id="h1")), + ] + ) + assert len(ch2.get()) == 1 + assert ch2.get()[0].content == "updated" + + +def test_delta_channel_checkpoint_returns_sentinel() -> None: + """checkpoint() always returns DELTA_SENTINEL regardless of state.""" + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) + assert ch.checkpoint() is DELTA_SENTINEL + + ch.update([HumanMessage(content="hi", id="h1")]) + assert ch.checkpoint() is DELTA_SENTINEL + + +# --------------------------------------------------------------------------- +# DeltaChannel — snapshot frequency +# --------------------------------------------------------------------------- + + +def test_delta_channel_snapshot_step_based() -> None: + """Snapshots fire on every Nth step regardless of whether the channel was written. + + With snapshot_frequency=N, every Nth pregel step produces a _DeltaSnapshot + blob — even if the channel had no write that step (eager snapshot). This + bounds the ancestor walk to at most N steps on any read. + """ + + class State(TypedDict): + messages: Annotated[ + list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=5) + ] + other: str + + def node_a(state: State) -> dict: + i = len(state["messages"]) // 2 + return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]} + + def node_b(state: State) -> dict: + return {"other": "y"} + + g = StateGraph(State) + g.add_node("a", node_a) + g.add_node("b", node_b) + g.add_edge(START, "a") + g.add_edge("a", "b") + saver = InMemorySaver() + graph = g.compile(checkpointer=saver) + + config = {"configurable": {"thread_id": "t1"}} + for i in range(6): + graph.invoke( + {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""}, + config, + ) + + msg_blob_values = [ + saver.serde.loads_typed((type_tag, blob)) + for k, (type_tag, blob) in saver.blobs.items() + if k[2] == "messages" and type_tag == "msgpack" and blob + ] + snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)] + assert snapshots, "expected at least one _DeltaSnapshot blob for messages" + + state = graph.get_state(config) + assert len(state.values["messages"]) == 12 # 6 human + 6 AI + + +def test_delta_channel_snapshot_fires_even_when_not_written() -> None: + """Eager snapshot: _DeltaSnapshot stored at snapshot step even when the + channel had no write that step (node_b doesn't touch messages). + """ + + class State(TypedDict): + messages: Annotated[ + list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3) + ] + tick: int + + def writer(state: State) -> dict: + i = len(state["messages"]) // 2 + return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]} + + def ticker(state: State) -> dict: + return {"tick": state["tick"] + 1} + + g = StateGraph(State) + g.add_node("writer", writer) + g.add_node("ticker", ticker) + g.add_edge(START, "writer") + g.add_edge("writer", "ticker") + saver = InMemorySaver() + graph = g.compile(checkpointer=saver) + + config = {"configurable": {"thread_id": "t1"}} + for i in range(5): + graph.invoke( + {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "tick": 0}, + config, + ) + + msg_blobs = { + k: saver.serde.loads_typed((t, b)) + for k, (t, b) in saver.blobs.items() + if k[2] == "messages" and t == "msgpack" and b + } + snapshots = {k: v for k, v in msg_blobs.items() if isinstance(v, _DeltaSnapshot)} + assert snapshots, ( + "eager snapshots must fire even on steps where messages wasn't written" + ) + + state = graph.get_state(config) + assert len(state.values["messages"]) == 10 # 5 human + 5 AI + + +# --------------------------------------------------------------------------- +# DeltaChannel — end-to-end (InMemorySaver) +# --------------------------------------------------------------------------- + + +def test_delta_channel_inmemory_saver_assembles_writes() -> None: + """InMemorySaver assembles writes from checkpoint_writes inside get_tuple.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer, list)] + + n = {"v": 0} + + def respond(state: State) -> dict: + n["v"] += 1 + return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + saver = InMemorySaver() + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "t1"}} + + graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config) + graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config) + + saved = saver.get_tuple(config) + assert saved is not None + assert "messages" in saved.checkpoint["channel_values"] + assert saved.checkpoint["channel_values"]["messages"] is DELTA_SENTINEL + + state = graph.get_state(config) + assert len(state.values["messages"]) == 4 # 2 human + 2 AI + + +# --------------------------------------------------------------------------- +# DeltaChannel — dict reducer +# --------------------------------------------------------------------------- + + +def _delta_channel_with_type(op, typ): + """Build a DeltaChannel with an explicit type via the Annotated injection path.""" + return _get_channel("_test", Annotated[typ, DeltaChannel(op)]) + + +def test_delta_channel_dict_reducer_fresh_channel() -> None: + """DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) + assert ch.is_available() + assert ch.get() == {} + + +def test_delta_channel_dict_reducer_basic_updates() -> None: + """DeltaChannel with a dict reducer accumulates key/value pairs across steps.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) + + ch.update([{"a": 1}]) + d1 = ch.checkpoint() + assert d1 is DELTA_SENTINEL + + ch.update([{"b": 2}]) + d2 = ch.checkpoint() + assert d2 is DELTA_SENTINEL + + assert ch.get() == {"a": 1, "b": 2} + + +def test_delta_channel_dict_reducer_writes_reconstruction() -> None: + """replay_writes on a fresh channel replays through a dict merge reducer.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + spec = _delta_channel_with_type(merge_dicts, dict) + ch = spec.from_checkpoint(DELTA_SENTINEL) + ch.replay_writes( + [ + ("t0", "files", {"a": 1}), + ("t1", "files", {"b": 2}), + ("t2", "files", {"c": 3}), + ] + ) + assert ch.get() == {"a": 1, "b": 2, "c": 3} + + +def test_delta_channel_dict_reducer_with_deletions() -> None: + """Dict reducer that treats None values as deletions works end-to-end.""" + + def merge_files(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + for k, v in w.items(): + if v is None: + result.pop(k, None) + else: + result[k] = v + return result + + ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING) + ch.update([{"file1.py": "content1", "file2.py": "content2"}]) + ch.update([{"file1.py": None, "file3.py": "content3"}]) + assert ch.get() == {"file2.py": "content2", "file3.py": "content3"} + + spec = _delta_channel_with_type(merge_files, dict) + ch2 = spec.from_checkpoint(DELTA_SENTINEL) + ch2.replay_writes( + [ + ("t0", "files", {"file1.py": "content1", "file2.py": "content2"}), + ("t1", "files", {"file1.py": None, "file3.py": "content3"}), + ] + ) + assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"} + + +def test_delta_channel_dict_reducer_overwrite_in_update() -> None: + """Overwrite(dict) in update() must preserve dict shape, not coerce to list.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) + ch.update([{"a": 1}]) + ch.update([Overwrite({"b": 2, "c": 3})]) + assert ch.get() == {"b": 2, "c": 3} + + +def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None: + """Overwrite(dict) embedded in replayed writes must reconstruct as dict.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + spec = _delta_channel_with_type(merge_dicts, dict) + ch = spec.from_checkpoint(DELTA_SENTINEL) + ch.replay_writes( + [ + ("t0", "files", {"a": 1}), + ("t1", "files", Overwrite({"x": 10, "y": 20})), + ("t2", "files", {"z": 30}), + ] + ) + assert ch.get() == {"x": 10, "y": 20, "z": 30} + + +def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None: + """DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + annotation = Annotated[NotRequired[dict[str, int]], DeltaChannel(merge_dicts)] + ch = _get_channel("files", annotation).from_checkpoint(MISSING) + assert ch.get() == {} + ch.update([{"a": 1}]) + ch.update([{"b": 2}]) + assert ch.get() == {"a": 1, "b": 2} + + +def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None: + """End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel.""" + + def merge_files(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + for k, v in w.items(): + if v is None: + result.pop(k, None) + else: + result[k] = v + return result + + class State(TypedDict): + files: Annotated[dict[str, str], DeltaChannel(merge_files)] + + turn = {"v": 0} + + def write_file(state: State) -> dict: + turn["v"] += 1 + n = turn["v"] + return {"files": {f"/doc_{n}.txt": f"content for turn {n}"}} + + builder = StateGraph(State) + builder.add_node("write_file", write_file) + builder.add_edge(START, "write_file") + saver = InMemorySaver() + graph = builder.compile(checkpointer=saver) + config = {"configurable": {"thread_id": "fs"}} + + for _ in range(3): + graph.invoke({"files": {}}, config) + + saved = saver.get_tuple(config) + assert saved is not None + assert saved.checkpoint["channel_values"]["files"] is DELTA_SENTINEL + state = graph.get_state(config) + assert state.values["files"] == { + "/doc_1.txt": "content for turn 1", + "/doc_2.txt": "content for turn 2", + "/doc_3.txt": "content for turn 3", + } + + def delete_file(state: State) -> dict: + return {"files": {"/doc_1.txt": None}} + + builder2 = StateGraph(State) + builder2.add_node("write_file", write_file) + builder2.add_node("delete_file", delete_file) + builder2.add_edge(START, "write_file") + builder2.add_edge("write_file", "delete_file") + turn["v"] = 0 + saver2 = InMemorySaver() + graph2 = builder2.compile(checkpointer=saver2) + config2 = {"configurable": {"thread_id": "fs2"}} + graph2.invoke({"files": {}}, config2) + state2 = graph2.get_state(config2) + assert state2.values["files"] == {} + + +def test_delta_channel_dict_reducer_backwards_compat() -> None: + """A pre-DeltaChannel dict checkpoint must load as a dict, not be listified.""" + + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result + + spec = _delta_channel_with_type(merge_dicts, dict) + old_value = {"a": 1, "b": 2} + ch = spec.from_checkpoint(old_value) + assert ch.get() == {"a": 1, "b": 2} + + +# --------------------------------------------------------------------------- +# DeltaChannel — seed / pre-delta migration +# --------------------------------------------------------------------------- + + +def test_delta_channel_from_checkpoint_honors_seed() -> None: + """A non-sentinel value to from_checkpoint is used as the pre-delta seed. + + Guards the pre-delta migration path: when the saver's ancestor walk hits + a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs + the post-migration state correctly rather than replaying from empty. + """ + spec = DeltaChannel(_messages_delta_reducer, list) + seed = [HumanMessage(content="pre-delta", id="p1")] + ch = spec.from_checkpoint(seed) + ch.replay_writes( + [ + ("t0", "messages", AIMessage(content="delta-1", id="d1")), + ("t1", "messages", HumanMessage(content="delta-2", id="d2")), + ] + ) + msgs = ch.get() + assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"] + + +def test_delta_channel_from_checkpoint_seed_without_writes() -> None: + """Reconstruction at a pre-delta ancestor with no newer deltas returns + just the seed — the saver's terminator fired immediately.""" + spec = DeltaChannel(_messages_delta_reducer, list) + seed = [HumanMessage(content="only-snap", id="s1")] + ch = spec.from_checkpoint(seed) + ch.replay_writes([]) + assert ch.get() == seed + + +def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> None: + """`seed=None` must start replay from None, not from an empty channel. + + The DELTA_SENTINEL / MISSING sentinels mean 'no seed'; passing `None` + explicitly should feed None to the reducer as the left operand. + """ + + def replace(state, writes): + return writes[-1] if writes else state + + spec = DeltaChannel(replace, list) + ch = spec.from_checkpoint(None) + ch.replay_writes([("t0", "x", "after")]) + assert ch.get() == "after" diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py new file mode 100644 index 000000000..cf17d0f12 --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -0,0 +1,478 @@ +"""Benchmark: DeltaChannel snapshot_frequency — storage vs. read-depth tradeoff. + +Run directly: python tests/test_delta_channel_benchmark.py +Run via pytest: pytest tests/test_delta_channel_benchmark.py -s + +Part 1 — baseline (original): DeltaChannel(inf) vs add_messages (BinOp). +Part 2 — snapshot_frequency sweep: shows the storage/read-latency tradeoff + across frequencies [1, 5, 10, 50, inf] at scale. + +Key insight: + snapshot_frequency=inf → O(N) storage, O(N) read depth (pure delta) + snapshot_frequency=N → O(N²/N) storage, O(N) read depth bounded by freq + snapshot_frequency=1 → O(N²) storage, O(1) read depth (full snapshot) +""" + +from __future__ import annotations + +import contextlib +import math +import os +import sys +import time +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import MemorySaver +from typing_extensions import TypedDict + +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import END, StateGraph +from langgraph.graph.message import _messages_delta_reducer, add_messages + +try: + from langgraph.checkpoint.postgres import PostgresSaver + + _POSTGRES_AVAILABLE = True + _POSTGRES_URI = os.environ.get( + "LANGGRAPH_BENCH_POSTGRES_URI", + "postgres://postgres@localhost:5432/postgres?sslmode=disable", + ) +except ImportError: + _POSTGRES_AVAILABLE = False + +# --------------------------------------------------------------------------- +# Realistic message payload (~100 tokens / ~400 chars each) +# --------------------------------------------------------------------------- + +_HUMAN_TEMPLATE = ( + "I need help understanding the implications of {topic} on our system architecture. " + "Specifically, I'm concerned about how this interacts with our existing {concern} " + "and whether we need to refactor the {component} layer before proceeding. " + "We've had prior incidents in this area and want to be deliberate. " + "What should we prioritize first, and are there known failure modes we should design around from the start?" +) + +_AI_TEMPLATE = ( + "Great question about {topic}. The key insight here is that {concern} introduces " + "a subtle ordering dependency that most teams overlook until they hit it in production. " + "For your {component} layer specifically, I'd recommend starting with a careful audit " + "of the interface boundaries before making any structural changes. This will give you " + "a clear picture of the blast radius and let you sequence the migration safely." +) + +_TOPICS = [ + "distributed tracing", + "eventual consistency", + "schema migration", + "backpressure handling", + "idempotency guarantees", + "cache invalidation", + "connection pooling", + "rate limiting", + "circuit breaking", + "observability pipelines", +] + +_CONCERNS = [ + "concurrency model", + "retry semantics", + "state management", + "error propagation", + "latency budget", +] + +_COMPONENTS = [ + "persistence", + "routing", + "ingestion", + "aggregation", + "serialization", +] + + +def _human_content(i: int) -> str: + return _HUMAN_TEMPLATE.format( + topic=_TOPICS[i % len(_TOPICS)], + concern=_CONCERNS[i % len(_CONCERNS)], + component=_COMPONENTS[i % len(_COMPONENTS)], + ) + + +def _ai_content(i: int) -> str: + return _AI_TEMPLATE.format( + topic=_TOPICS[i % len(_TOPICS)], + concern=_CONCERNS[i % len(_CONCERNS)], + component=_COMPONENTS[i % len(_COMPONENTS)], + ) + + +# --------------------------------------------------------------------------- +# State definitions +# --------------------------------------------------------------------------- + + +class BinaryState(TypedDict): + messages: Annotated[list, add_messages] + + +class DeltaState(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + +def _make_delta_state(snapshot_frequency: int | float) -> type: + """Create a TypedDict with DeltaChannel at the given snapshot_frequency.""" + channel = DeltaChannel( + _messages_delta_reducer, snapshot_frequency=snapshot_frequency + ) + # Use the functional TypedDict form so the Annotated type is stored as an + # already-evaluated object rather than a forward-reference string (which + # would fail when get_type_hints tries to resolve 'snapshot_frequency'). + return TypedDict( # type: ignore[return-value] + f"DeltaState_freq{snapshot_frequency}", + {"messages": Annotated[list, channel]}, + ) + + +# --------------------------------------------------------------------------- +# Graph factory +# --------------------------------------------------------------------------- + + +def _make_graph(state_cls: type, checkpointer: Any = None) -> Any: + def human_node(state: Any) -> dict: + return {} + + def ai_node(state: Any) -> dict: + i = len(state["messages"]) // 2 + return {"messages": [AIMessage(content=_ai_content(i), id=f"a{i}")]} + + g = StateGraph(state_cls) + g.add_node("human", human_node) + g.add_node("ai", ai_node) + g.add_edge("human", "ai") + g.add_edge("ai", END) + g.set_entry_point("human") + return g.compile(checkpointer=checkpointer or MemorySaver()) + + +# --------------------------------------------------------------------------- +# Measurement helpers +# --------------------------------------------------------------------------- + + +def _total_blob_bytes(saver: MemorySaver) -> int: + total = 0 + for (_, _, _, _), (type_tag, blob) in saver.blobs.items(): + if blob is not None: + total += len(blob) + return total + + +def _run_turns( + n_turns: int, + state_cls: type, + checkpointer: Any = None, +) -> tuple[float, float, int]: + """Run n_turns conversation turns. + + Returns (write_elapsed_s, read_elapsed_s, total_blob_bytes). + Read latency is the average of 5 get_state calls after the full history + is built — forces state rehydration including ancestor replay if needed. + """ + graph = _make_graph(state_cls, checkpointer) + config = {"configurable": {"thread_id": "bench"}} + + t0 = time.perf_counter() + for i in range(n_turns): + graph.invoke( + {"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]}, + config, + ) + write_elapsed = time.perf_counter() - t0 + + t1 = time.perf_counter() + for _ in range(5): + graph.get_state(config) + read_elapsed = (time.perf_counter() - t1) / 5 + + blob_bytes = ( + _total_blob_bytes(graph.checkpointer) + if isinstance(graph.checkpointer, MemorySaver) + else -1 + ) + return write_elapsed, read_elapsed, blob_bytes + + +def _fmt_bytes(n: int) -> str: + if n >= 1_000_000: + return f"{n / 1_000_000:.1f} MB" + if n >= 1_000: + return f"{n / 1_000:.1f} KB" + return f"{n} B" + + +def _approx_tokens(n_turns: int) -> str: + tokens = n_turns * 200 + if tokens >= 1_000_000: + return f"~{tokens / 1_000_000:.1f}M tok" + if tokens >= 1_000: + return f"~{tokens / 1_000:.0f}K tok" + return f"~{tokens} tok" + + +# --------------------------------------------------------------------------- +# Checkpointer factories +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _pg_saver(thread_id: str = "bench"): + """Context manager that yields a fresh PostgresSaver and cleans up after.""" + with PostgresSaver.from_conn_string(_POSTGRES_URI) as saver: + saver.setup() + with saver._cursor() as cur: + for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"): + cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,)) + yield saver + with saver._cursor() as cur: + for tbl in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"): + cur.execute(f"DELETE FROM {tbl} WHERE thread_id = %s", (thread_id,)) + + +def _checkpointers() -> list[tuple[str, Any]]: + """Return (label, saver_or_None) pairs for available checkpointers.""" + result: list[tuple[str, Any]] = [("InMemory", None)] + if _POSTGRES_AVAILABLE: + try: + import psycopg + + psycopg.connect(_POSTGRES_URI).close() + result.append(("Postgres", "postgres")) + except Exception: + pass + return result + + +# --------------------------------------------------------------------------- +# Part 1: baseline DeltaChannel(inf) vs add_messages +# --------------------------------------------------------------------------- + +BASELINE_TURN_COUNTS = [10, 25, 50, 100, 500] +DELTA_ONLY_TURN_COUNTS = [1000] + + +def _run_baseline_for_checkpointer(cp_label: str, cp_hint: Any) -> None: + W = 72 + + def _make_saver(): + if cp_hint is None: + return contextlib.nullcontext(None) + return _pg_saver() + + rows: list[tuple[int, Any, Any, Any, Any, Any, Any]] = [] + for turns in BASELINE_TURN_COUNTS: + with _make_saver() as saver: + b_wt, b_rt, b_bytes = _run_turns(turns, BinaryState, saver) + with _make_saver() as saver: + d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) + rows.append((turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt)) + for turns in DELTA_ONLY_TURN_COUNTS: + with _make_saver() as saver: + d_wt, d_rt, d_bytes = _run_turns(turns, DeltaState, saver) + rows.append((turns, None, d_bytes, None, d_rt, None, d_wt)) + + def _bytes_or_na(v: Any) -> str: + if v is None or v < 0: + return "n/a" + return _fmt_bytes(v) + + def _ms_or_na(v: Any) -> str: + return "n/a" if v is None else f"{v * 1000:.1f}ms" + + print(f"\n [{cp_label}] Storage (blob bytes)") + print( + f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}" + ) + print(" " + "-" * (W - 2)) + for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows: + if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0: + ratio_str = "n/a" + else: + ratio = b_bytes / d_bytes if d_bytes else float("inf") + ratio_str = f"{ratio:.0f}x" + print( + f" {turns:>6} {_approx_tokens(turns):>10} " + f"{_bytes_or_na(b_bytes):>12} {_bytes_or_na(d_bytes):>12} {ratio_str:>8}" + ) + + print(f"\n [{cp_label}] Read latency (avg of 5 get_state calls)") + print(f" {'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12}") + print(" " + "-" * (W - 2)) + for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows: + print( + f" {turns:>6} {_approx_tokens(turns):>10} " + f"{_ms_or_na(b_rt):>12} {_ms_or_na(d_rt):>12}" + ) + + +def run_baseline_benchmark() -> None: + print() + print("Part 1 — DeltaChannel(inf) vs add_messages: storage & latency") + print("=" * 72) + for cp_label, cp_hint in _checkpointers(): + _run_baseline_for_checkpointer(cp_label, cp_hint) + print() + + +# --------------------------------------------------------------------------- +# Part 2: snapshot_frequency sweep +# --------------------------------------------------------------------------- + +# Frequencies to test. 1 = always snapshot (like BinOp), inf = pure delta. +SNAPSHOT_FREQUENCIES: list[int | float] = [1, 5, 10, 50, math.inf] + +# Turn counts for the sweep — high enough to show storage divergence. +SWEEP_TURN_COUNTS = [50, 100, 500] + + +def _freq_label(freq: int | float) -> str: + if freq == math.inf: + return "inf" + return str(int(freq)) + + +def _run_sweep_for_checkpointer(cp_label: str, cp_hint: Any) -> None: + def _make_saver(): + if cp_hint is None: + return contextlib.nullcontext(None) + return _pg_saver() + + # Collect results: {turns: {freq_label: (write_s, read_s, bytes)}} + results: dict[int, dict[str, tuple[float, float, int]]] = {} + for turns in SWEEP_TURN_COUNTS: + results[turns] = {} + for freq in SNAPSHOT_FREQUENCIES: + state_cls = _make_delta_state(freq) + with _make_saver() as saver: + wt, rt, bb = _run_turns(turns, state_cls, saver) + results[turns][_freq_label(freq)] = (wt, rt, bb) + + freq_labels = [_freq_label(f) for f in SNAPSHOT_FREQUENCIES] + col_w = 12 + + header = f" {'turns':>6} {'ctx':>10}" + "".join( + f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels + ) + + print(f"\n [{cp_label}] Storage (blob bytes) — lower is better") + print(header) + print(" " + "-" * (len(header) - 2)) + for turns in SWEEP_TURN_COUNTS: + row = f" {turns:>6} {_approx_tokens(turns):>10}" + for label in freq_labels: + _, _, bb = results[turns][label] + row += f" {_fmt_bytes(bb) if bb >= 0 else 'n/a':>{col_w}}" + print(row) + + print(f"\n [{cp_label}] Read latency (avg of 5 get_state) — lower is better") + print(header) + print(" " + "-" * (len(header) - 2)) + for turns in SWEEP_TURN_COUNTS: + row = f" {turns:>6} {_approx_tokens(turns):>10}" + for label in freq_labels: + _, rt, _ = results[turns][label] + row += f" {f'{rt * 1000:.1f}ms':>{col_w}}" + print(row) + + print( + f"\n [{cp_label}] Per-invoke write latency (total / turns) — lower is better" + ) + print(header) + print(" " + "-" * (len(header) - 2)) + for turns in SWEEP_TURN_COUNTS: + row = f" {turns:>6} {_approx_tokens(turns):>10}" + for label in freq_labels: + wt, _, _ = results[turns][label] + row += f" {f'{(wt / turns) * 1000:.1f}ms':>{col_w}}" + print(row) + + +def run_snapshot_freq_benchmark() -> None: + print() + print("Part 2 — DeltaChannel snapshot_frequency sweep") + print("Lower freq → fewer snapshots → less storage but deeper read replay") + print("=" * 80) + for cp_label, cp_hint in _checkpointers(): + _run_sweep_for_checkpointer(cp_label, cp_hint) + print() + print("Legend:") + print( + " freq=1 snapshot every write (full blob always — same as add_messages / BinOp)" + ) + print(" freq=N snapshot every N writes; read walks at most N ancestor writes") + print(" freq=inf pure delta; read walks entire ancestor chain") + print() + + +# --------------------------------------------------------------------------- +# Pytest entry points +# --------------------------------------------------------------------------- + + +@pytest.mark.skip( + reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py" +) +def test_delta_channel_baseline_benchmark(capsys: Any) -> None: + """DeltaChannel(inf) uses less storage than add_messages at scale.""" + with capsys.disabled(): + run_baseline_benchmark() + + for turns in [25, 50]: + _, _, b_bytes = _run_turns(turns, BinaryState) + _, _, d_bytes = _run_turns(turns, DeltaState) + assert d_bytes < b_bytes, ( + f"DeltaChannel should use less storage at {turns} turns, " + f"got delta={d_bytes} binary={b_bytes}" + ) + + +@pytest.mark.skip( + reason="slow benchmark — run manually with: python tests/test_delta_channel_benchmark.py" +) +def test_snapshot_freq_benchmark(capsys: Any) -> None: + """snapshot_frequency trades storage for bounded read depth.""" + with capsys.disabled(): + run_snapshot_freq_benchmark() + + # Correctness: results at all frequencies should agree on final state. + n_turns = 20 + states: dict[str, list] = {} + for freq in SNAPSHOT_FREQUENCIES: + state_cls = _make_delta_state(freq) + graph = _make_graph(state_cls) + config = {"configurable": {"thread_id": "correctness"}} + for i in range(n_turns): + graph.invoke( + {"messages": [HumanMessage(content=_human_content(i), id=f"h{i}")]}, + config, + ) + state = graph.get_state(config) + states[_freq_label(freq)] = [m.id for m in state.values["messages"]] + + ref = states["inf"] + for label, msg_ids in states.items(): + assert msg_ids == ref, ( + f"freq={label} produced different message IDs than freq=inf" + ) + + +# --------------------------------------------------------------------------- +# Script entry point +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + run_baseline_benchmark() + run_snapshot_freq_benchmark() + sys.exit(0) diff --git a/libs/langgraph/tests/test_delta_channel_migration.py b/libs/langgraph/tests/test_delta_channel_migration.py new file mode 100644 index 000000000..245ba2561 --- /dev/null +++ b/libs/langgraph/tests/test_delta_channel_migration.py @@ -0,0 +1,613 @@ +"""Tests for the BinaryOperatorAggregate -> DeltaChannel migration path. + +A thread written under `BinaryOperatorAggregate(...)` must keep working +after its annotation is swapped to `DeltaChannel(...)` on the same +checkpointer — pre-migration state visible at each *settled* ancestor +checkpoint is preserved, and post-migration writes fold on top through +the reducer. + +Mechanism under test: the saver's `_get_channel_writes_history(config, +channel)` walks the parent chain; when it encounters an ancestor whose +`channel_values[channel]` is a real value (not `DELTA_SENTINEL`), it +returns that as the `seed`. `DeltaChannel.from_checkpoint(seed)` uses +it as the base value, and `replay_writes(writes)` folds on-path deltas. + +Scenarios covered: + +1. **Basic migration (sync + async)**: build pre-migration state with + `BinaryOperatorAggregate`, swap the annotation to `DeltaChannel` on + the same checkpointer, and verify that every settled pre-migration + super-step boundary (`next=('__start__',)`) round-trips exactly + under the delta-channel view. +2. **Time travel into a pre-migration checkpoint** after migration — + `graph.get_state(pre_migration_config)` at a settled ancestor + returns the same state as under the binop channel. +3. **Continuing a migrated thread**: driving one more super-step after + migration produces a state that includes the pre-migration settled + prefix plus the new delta write — proving `from_checkpoint(seed)` + + `replay_writes` correctly fold post-migration deltas onto the + pre-migration seed. +4. **Base-saver fallback path**: a third-party-style subclass that + removes the optimized `InMemorySaver` override and falls back to + `BaseCheckpointSaver._get_channel_writes_history` must produce the + same result as the optimized path. +5. **Channel-type isolation across threads**: two threads on the same + checkpointer under the delta-channel graph — one freshly-started, + one migrated from pre-migration state — don't cross-contaminate. + The parent-chain walk is scoped to the thread. + +TODO: add postgres variants in the existing `libs/checkpoint-postgres` +test files (different fixture setup; not this file). +""" + +from __future__ import annotations + +import operator +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict + +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import _messages_delta_reducer, add_messages + +pytestmark = pytest.mark.anyio + + +# --------------------------------------------------------------------------- +# Graph factories +# +# A minimal reducer (`operator.add` on lists of str) with a noop node keeps +# state change localized to the HumanMessage-like payload passed through +# `invoke`. That isolates the pre/post-migration parity assertions to +# channel-hydration semantics. +# --------------------------------------------------------------------------- + + +def _noop(_state: Any) -> dict: + return {} + + +def _list_concat(state: list, writes: list) -> list: + result = list(state) + for w in writes: + result.extend(w if isinstance(w, list) else [w]) + return result + + +def _binop_graph(checkpointer: Any) -> Any: + class BinopState(TypedDict): + items: Annotated[list, BinaryOperatorAggregate(list, operator.add)] + + return ( + StateGraph(BinopState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _delta_graph(checkpointer: Any) -> Any: + class DeltaState(TypedDict): + items: Annotated[list, DeltaChannel(_list_concat)] + + return ( + StateGraph(DeltaState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _drive(graph: Any, config: dict, tag: str, n: int) -> None: + for i in range(n): + graph.invoke({"items": [f"{tag}{i}"]}, config) + + +async def _adrive(graph: Any, config: dict, tag: str, n: int) -> None: + for i in range(n): + await graph.ainvoke({"items": [f"{tag}{i}"]}, config) + + +def _settled_boundaries(history: list) -> list[tuple[dict, list]]: + """Return `[(config, items), ...]` for every checkpoint in `history` + whose `next == ('__start__',)` — the stable boundaries between invokes. + """ + return [ + (s.config, list(s.values.get("items", []))) + for s in history + if s.next == ("__start__",) + ] + + +# --------------------------------------------------------------------------- +# 1. Basic migration (sync + async) +# --------------------------------------------------------------------------- + + +def test_basic_migration_preserves_pre_migration_state() -> None: + """Build state under `BinaryOperatorAggregate`, migrate to + `DeltaChannel` on the same checkpointer, and verify that every + settled pre-migration super-step boundary round-trips exactly. + + Settled boundaries (`next=('__start__',)`) are the stable hydration + targets for the migration path: writes that produced the NEXT + super-step are kept as `pending_writes` on the ancestor, so walking + from a descendant finds the ancestor's blob as the seed and + reconstructs the correct state. + """ + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "basic-sync"}} + + # Pre-migration: accumulate items across 3 invokes. + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 3) + + pre_boundaries = _settled_boundaries(list(binop.get_state_history(config))) + assert len(pre_boundaries) >= 2, "expected multiple settled boundaries" + + # Migrate: swap the annotation on the same checkpointer. + delta = _delta_graph(checkpointer) + + for cfg, items in pre_boundaries: + snap = delta.get_state(cfg) + assert list(snap.values.get("items", [])) == items, ( + f"snapshot mismatch at {cfg['configurable']['checkpoint_id']}: " + f"expected {items}, got {snap.values.get('items', [])}" + ) + + +async def test_basic_migration_preserves_pre_migration_state_async() -> None: + """Async variant of the basic migration scenario.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "basic-async"}} + + binop = _binop_graph(checkpointer) + await _adrive(binop, config, "u", 3) + + pre_history = [s async for s in binop.aget_state_history(config)] + pre_boundaries = _settled_boundaries(pre_history) + assert len(pre_boundaries) >= 2 + + delta = _delta_graph(checkpointer) + + for cfg, items in pre_boundaries: + snap = await delta.aget_state(cfg) + assert list(snap.values.get("items", [])) == items, ( + f"async snapshot mismatch at {cfg['configurable']['checkpoint_id']}" + ) + + +# --------------------------------------------------------------------------- +# 2. Time travel into a pre-migration checkpoint after migration +# --------------------------------------------------------------------------- + + +def test_time_travel_into_pre_migration_checkpoint() -> None: + """After migration, `graph.get_state(pre_migration_config)` at a + settled ancestor returns the state as stored at that point.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "time-travel"}} + + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 3) + + pre_boundaries = _settled_boundaries(list(binop.get_state_history(config))) + assert pre_boundaries, "no settled ancestors to time-travel to" + + delta = _delta_graph(checkpointer) + + # Pick the oldest non-empty boundary — a long distance to walk back. + non_empty = [(cfg, items) for cfg, items in pre_boundaries if items] + assert non_empty, "expected at least one non-empty boundary" + target_cfg, expected_items = non_empty[-1] + + snap = delta.get_state(target_cfg) + assert list(snap.values.get("items", [])) == expected_items + + +# --------------------------------------------------------------------------- +# 3. Continuing a migrated thread: deltas fold onto pre-migration seed +# --------------------------------------------------------------------------- + + +def test_continuing_migrated_thread_folds_deltas_on_seed() -> None: + """Resume a pre-migration settled ancestor via `invoke(None, cfg)` + under the delta-channel graph. Since the pre-migration checkpoint + has an existing `pending_writes` entry (the input for the NEXT + super-step), re-running from that ancestor reproduces the same + post-ancestor state as the original binop run. + + This proves the seed-terminator + write-replay pipeline works + end-to-end across the migration boundary. + """ + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "continue"}} + + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 2) + + # Pick the oldest settled boundary with non-empty state. + pre_boundaries = _settled_boundaries(list(binop.get_state_history(config))) + target_cfg, seed_items = next( + (cfg, items) for cfg, items in reversed(pre_boundaries) if items + ) + assert seed_items, "need a non-empty seed boundary" + + # Migrate and resume from the pre-migration ancestor. `invoke(None, + # cfg)` replays the pending writes staged at `cfg` under the new + # channel; the reducer folds those deltas onto the seed. + delta = _delta_graph(checkpointer) + result = delta.invoke(None, target_cfg) + + # The resumed state must include the pre-migration seed items in order. + result_items = list(result.get("items", [])) + for idx, prefix_item in enumerate(seed_items): + assert result_items[idx] == prefix_item, ( + f"pre-migration seed item at {idx} not preserved: " + f"got {result_items[: idx + 1]}, expected {seed_items}" + ) + + +# --------------------------------------------------------------------------- +# 4. Base-saver fallback path +# --------------------------------------------------------------------------- + + +class _ThirdPartyStyleSaver(InMemorySaver): + """Simulates a third-party saver that inherits the reference + `_get_channel_writes_history` implementation from + `BaseCheckpointSaver` rather than overriding it. + + We rebind the two methods to the base-class versions (via MRO) so + the fallback path is exercised even though the storage layer is + still the in-memory one. + """ + + # MRO: [_ThirdPartyStyleSaver, InMemorySaver, BaseCheckpointSaver, ...] + _get_channel_writes_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1]._get_channel_writes_history # type: ignore[attr-defined] + ) + _aget_channel_writes_history = ( # type: ignore[assignment] + InMemorySaver.__mro__[1]._aget_channel_writes_history # type: ignore[attr-defined] + ) + + +def test_base_saver_fallback_matches_optimized_override() -> None: + """The reference `BaseCheckpointSaver` implementation must produce + the same migration behavior as the optimized `InMemorySaver` + override. We drive the same migration scenario through both savers + and assert per-snapshot parity in the delta-channel view.""" + + # Fast path: optimized InMemorySaver override. + fast_saver = InMemorySaver() + fast_config = {"configurable": {"thread_id": "fast"}} + fast_binop = _binop_graph(fast_saver) + _drive(fast_binop, fast_config, "u", 3) + fast_delta = _delta_graph(fast_saver) + fast_history = [ + (s.next, list(s.values.get("items", []))) + for s in fast_delta.get_state_history(fast_config) + ] + + # Slow path: base-class fallback. + slow_saver = _ThirdPartyStyleSaver() + slow_config = {"configurable": {"thread_id": "slow"}} + slow_binop = _binop_graph(slow_saver) + _drive(slow_binop, slow_config, "u", 3) + slow_delta = _delta_graph(slow_saver) + slow_history = [ + (s.next, list(s.values.get("items", []))) + for s in slow_delta.get_state_history(slow_config) + ] + + assert slow_history == fast_history, ( + "base-saver fallback should match optimized-override behavior; " + f"fast={fast_history}, slow={slow_history}" + ) + + +# --------------------------------------------------------------------------- +# 5. Thread isolation under mixed-generation storage +# --------------------------------------------------------------------------- + + +def test_delta_and_migrated_threads_do_not_cross_contaminate() -> None: + """Two threads sharing a checkpointer — one migrated from + pre-migration state, one freshly-started under DeltaChannel — must + maintain independent state. The parent-chain walk in + `_get_channel_writes_history` must be scoped to the target thread. + """ + + checkpointer = InMemorySaver() + migrated_cfg = {"configurable": {"thread_id": "migrated"}} + fresh_cfg = {"configurable": {"thread_id": "fresh"}} + + # Thread A: pre-migration build-up. + binop = _binop_graph(checkpointer) + _drive(binop, migrated_cfg, "m", 2) + + # Thread B: fresh delta-channel run. + delta = _delta_graph(checkpointer) + _drive(delta, fresh_cfg, "f", 2) + + # Thread A: migrate and confirm its state is anchored in its own + # thread's pre-migration history (tag 'm'), never mixing in tag 'f'. + migrated_boundaries = _settled_boundaries( + list(delta.get_state_history(migrated_cfg)) + ) + assert migrated_boundaries, "migrated thread has no settled boundaries" + for _, items in migrated_boundaries: + for it in items: + assert it.startswith("m"), ( + f"migrated thread leaked item from other thread: {it}" + ) + + # Thread B: settled boundaries must only contain 'f' tags. + fresh_boundaries = _settled_boundaries(list(delta.get_state_history(fresh_cfg))) + assert fresh_boundaries, "fresh thread has no settled boundaries" + for _, items in fresh_boundaries: + for it in items: + assert it.startswith("f"), ( + f"fresh thread leaked item from migrated thread: {it}" + ) + + +# --------------------------------------------------------------------------- +# 6. Tip-of-pre-migration hydration: the latest checkpoint from a binop-run +# thread has a real accumulated value in its own `channel_values["items"]`. +# When hydrated under the delta-channel graph via `get_state(config)` with no +# `checkpoint_id`, the short-circuit must use that value directly instead of +# walking ancestors (which would skip the tip's own blob). +# --------------------------------------------------------------------------- + + +def test_tip_of_pre_migration_hydrates_directly() -> None: + """`graph.get_state(config)` at the latest (pre-migration) checkpoint + returns the full accumulated list stored in that checkpoint's own + `channel_values`. The hydration must not walk ancestors past it.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "tip-sync"}} + + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 3) + + binop_tip = binop.get_state(config) + expected_items = list(binop_tip.values.get("items", [])) + assert expected_items == ["u0", "u1", "u2"], ( + f"sanity: pre-migration tip should accumulate all 3 items, got {expected_items}" + ) + + delta = _delta_graph(checkpointer) + + snap = delta.get_state(config) + assert list(snap.values.get("items", [])) == expected_items, ( + f"tip hydration mismatch: expected {expected_items}, " + f"got {snap.values.get('items', [])}" + ) + + +async def test_tip_of_pre_migration_hydrates_directly_async() -> None: + """Async variant of the tip-of-pre-migration hydration scenario.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "tip-async"}} + + binop = _binop_graph(checkpointer) + await _adrive(binop, config, "u", 3) + + binop_tip = await binop.aget_state(config) + expected_items = list(binop_tip.values.get("items", [])) + assert expected_items == ["u0", "u1", "u2"] + + delta = _delta_graph(checkpointer) + + snap = await delta.aget_state(config) + assert list(snap.values.get("items", [])) == expected_items, ( + f"async tip hydration mismatch: expected {expected_items}, " + f"got {snap.values.get('items', [])}" + ) + + +# --------------------------------------------------------------------------- +# 7. `update_state` after migration writes a real value to the new +# checkpoint's `channel_values` (not a sentinel). Hydration must use it +# directly — the ancestor walk would skip this blob and return stale state. +# --------------------------------------------------------------------------- + + +def test_update_state_after_migration_uses_written_value() -> None: + """After migrating and running at least one post-migration super-step + (so the thread's tip has a `DELTA_SENTINEL`), `update_state` writes a + concrete value to a new checkpoint's `channel_values`. `get_state` + must reflect that concrete value.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "update-state"}} + + # Pre-migration: accumulate a little state. + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 2) + + # Migrate and run one more super-step so the tip is a post-migration + # checkpoint with `DELTA_SENTINEL` in its own `channel_values`. + delta = _delta_graph(checkpointer) + delta.invoke({"items": ["post"]}, config) + + # `update_state` writes a concrete value into a new checkpoint's blob + # via the reducer against the hydrated prior state. + delta.update_state(config, {"items": ["x", "y"]}) + + snap = delta.get_state(config) + updated_items = list(snap.values.get("items", [])) + # Must include the "x","y" update; without the hydration fix, the + # update_state-written blob would be skipped in favor of an ancestor + # walk, and the update values would disappear. + assert "x" in updated_items and "y" in updated_items, ( + f"update_state values missing from snapshot: {updated_items}" + ) + # The "x","y" items should be folded onto the prior accumulated state, + # not stand alone. This verifies the update-written blob is used + # directly by `get_state` (no ancestor walk past it). + assert len(updated_items) >= 4, ( + f"update_state snapshot should preserve pre-update state, got {updated_items}" + ) + assert updated_items[-2:] == ["x", "y"], ( + f"update_state deltas should be at the tail, got {updated_items}" + ) + + +# --------------------------------------------------------------------------- +# 8. Fork from an `update_state` checkpoint: a new run branched off the +# update_state-produced checkpoint must see that checkpoint's concrete +# `channel_values` as its base, with new deltas folded on top. +# --------------------------------------------------------------------------- + + +def test_fork_from_update_state_checkpoint() -> None: + """Branching a new run from the checkpoint produced by `update_state` + must use that checkpoint's concrete blob as the base. Additional + deltas from the forked run fold onto it through the reducer.""" + + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "fork"}} + + # Pre-migration build-up, then migrate and add one post-migration step. + binop = _binop_graph(checkpointer) + _drive(binop, config, "u", 2) + delta = _delta_graph(checkpointer) + delta.invoke({"items": ["post"]}, config) + + # Apply `update_state` and capture the returned config (references + # the new checkpoint produced by the update). + update_cfg = delta.update_state(config, {"items": ["x", "y"]}) + + update_snap = delta.get_state(update_cfg) + base_items = list(update_snap.values.get("items", [])) + assert "x" in base_items and "y" in base_items, ( + f"update_state values missing from snapshot: {base_items}" + ) + assert base_items[-2:] == ["x", "y"], ( + f"sanity: update_state deltas should be at the tail, got {base_items}" + ) + + # Fork: invoke from the update_state checkpoint with a new delta. + forked = delta.invoke({"items": ["fork0"]}, update_cfg) + forked_items = list(forked.get("items", [])) + # The fork must see the update_state-written blob as its base (not + # walk past it), and the new delta must fold on top of it. + assert forked_items[: len(base_items)] == base_items, ( + f"fork lost update_state base: base={base_items}, forked={forked_items}" + ) + assert forked_items[-1] == "fork0", f"fork delta not appended: {forked_items}" + + +# --------------------------------------------------------------------------- +# 9. Migration from `add_messages` → `DeltaChannel(_messages_delta_reducer)` +# +# `add_messages` is the primary real-world use case: it creates a +# BinaryOperatorAggregate with dedup-by-ID and RemoveMessage semantics. +# After swapping the annotation to DeltaChannel, pre-migration blobs +# (plain lists of Message objects) must be used directly as the seed. +# --------------------------------------------------------------------------- + + +def _add_messages_graph(checkpointer: Any) -> Any: + class MessagesState(TypedDict): + messages: Annotated[list, add_messages] + + return ( + StateGraph(MessagesState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def _delta_messages_graph(checkpointer: Any) -> Any: + class DeltaMessagesState(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + return ( + StateGraph(DeltaMessagesState) + .add_node("noop", _noop) + .add_edge(START, "noop") + .add_edge("noop", END) + .compile(checkpointer=checkpointer) + ) + + +def test_add_messages_to_delta_migration_preserves_message_history() -> None: + """Migration from `add_messages` to `DeltaChannel(_messages_delta_reducer)` + preserves message ordering and IDs at both the tip and settled ancestor + boundaries. + + The pre-migration blob is a plain list of Message objects; DeltaChannel + must use it directly as the seed without walking ancestors past it. + """ + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "add-messages-migration"}} + + pre_graph = _add_messages_graph(checkpointer) + pre_graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) + pre_graph.invoke({"messages": [AIMessage(content="hi", id="a1")]}, config) + pre_graph.invoke({"messages": [HumanMessage(content="thanks", id="h2")]}, config) + + pre_tip = pre_graph.get_state(config) + assert [m.id for m in pre_tip.values["messages"]] == ["h1", "a1", "h2"] + + delta_graph = _delta_messages_graph(checkpointer) + + # Tip: latest checkpoint has a full list blob — must use it directly. + snap = delta_graph.get_state(config) + assert [m.id for m in snap.values["messages"]] == ["h1", "a1", "h2"], ( + f"tip hydration mismatch: got {[m.id for m in snap.values['messages']]}" + ) + + # Settled ancestor boundaries must also match. + pre_settled = [ + [m.id for m in s.values.get("messages", [])] + for s in pre_graph.get_state_history(config) + if s.next == ("__start__",) + ] + delta_settled = [ + [m.id for m in s.values.get("messages", [])] + for s in delta_graph.get_state_history(config) + if s.next == ("__start__",) + ] + assert delta_settled == pre_settled, ( + f"settled boundary mismatch after migration: " + f"pre={pre_settled}, delta={delta_settled}" + ) + + +async def test_add_messages_to_delta_migration_preserves_message_history_async() -> ( + None +): + """Async variant of the add_messages migration test.""" + checkpointer = InMemorySaver() + config = {"configurable": {"thread_id": "add-messages-migration-async"}} + + pre_graph = _add_messages_graph(checkpointer) + await pre_graph.ainvoke( + {"messages": [HumanMessage(content="hello", id="h1")]}, config + ) + await pre_graph.ainvoke({"messages": [AIMessage(content="hi", id="a1")]}, config) + + delta_graph = _delta_messages_graph(checkpointer) + snap = await delta_graph.aget_state(config) + assert [m.id for m in snap.values["messages"]] == ["h1", "a1"], ( + f"async tip hydration mismatch: got {[m.id for m in snap.values['messages']]}" + ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e4d49c46d..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, @@ -41,6 +42,7 @@ from typing_extensions import NotRequired, TypedDict from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.delta import DeltaChannel from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -49,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, @@ -9400,3 +9402,254 @@ def test_fork_does_not_apply_pending_writes( # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 assert result == {"value": 121} + + +async def test_delta_channel_end_to_end_inmemory() -> None: + """Full graph run: DeltaChannel accumulates correctly across multiple turns.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def respond(state: State) -> dict: + n = len(state["messages"]) + return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-test-1"}} + + # Turn 1 + graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) + # Turn 2 + graph.invoke({"messages": [HumanMessage(content="world", id="h2")]}, config) + # Turn 3 + graph.invoke({"messages": [HumanMessage(content="bye", id="h3")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + # 3 human + 3 AI = 6 total + assert len(msgs) == 6, f"expected 6 messages, got {len(msgs)}: {msgs}" + assert msgs[0].content == "hello" + assert msgs[2].content == "world" + assert msgs[4].content == "bye" + assert msgs[1].content == "reply-1" + assert msgs[3].content == "reply-3" + assert msgs[5].content == "reply-5" + + +async def test_delta_channel_time_travel() -> None: + """Time-travel back to turn-1 checkpoint and resume; continuation must not include turn-2 deltas.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + counter = {"n": 0} + + def respond(state: State) -> dict: + counter["n"] += 1 + return { + "messages": [ + AIMessage(content=f"ai-{counter['n']}", id=f"ai-{counter['n']}") + ] + } + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + saver = InMemorySaver() + graph = builder.compile(checkpointer=saver) + + config = {"configurable": {"thread_id": "diff-time-travel"}} + + # Run 2 turns: h1→ai-1, h2→ai-2 + graph.invoke({"messages": [HumanMessage(content="h1", id="h1")]}, config) + graph.invoke({"messages": [HumanMessage(content="h2", id="h2")]}, config) + + # Find the checkpoint after turn 1 (2 messages: h1 + ai-1) + history = list(graph.get_state_history(config)) + after_turn1 = next(h for h in history if len(h.values.get("messages", [])) == 2) + + assert len(after_turn1.values["messages"]) == 2 + assert after_turn1.values["messages"][0].content == "h1" + assert after_turn1.values["messages"][1].content == "ai-1" + + # Resume from turn-1 checkpoint: inject h3, expect 3 messages total (h1, ai-1, ai-N) + # NOT 5 messages (turn-2 deltas must not bleed into the resumed run) + result = graph.invoke( + {"messages": [HumanMessage(content="h3", id="h3")]}, + after_turn1.config, + ) + msgs = result["messages"] + # Should be: h1, ai-1, h3, ai-N — 4 messages total + assert len(msgs) == 4, ( + f"expected 4 messages after time-travel resume, got {len(msgs)}: {msgs}" + ) + assert msgs[0].content == "h1" + assert msgs[1].content == "ai-1" + assert msgs[2].content == "h3" + + +async def test_delta_channel_remove_message_end_to_end() -> None: + """RemoveMessage inside a DeltaChannel graph must persist and reload correctly.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def respond(state: State) -> dict: + return {"messages": [AIMessage(content="reply", id="ai-1")]} + + def delete_first(state: State) -> dict: + # removes the first message + return {"messages": [RemoveMessage(id=state["messages"][0].id)]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_node("delete_first", delete_first) + builder.add_edge(START, "respond") + builder.add_edge("respond", "delete_first") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-remove-test"}} + graph.invoke({"messages": [HumanMessage(content="hello", id="h1")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + # h1 was removed, only ai-1 should remain + assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}" + assert msgs[0].id == "ai-1" + + # A subsequent turn must reconstruct from the checkpoint correctly + graph.invoke({"messages": [HumanMessage(content="again", id="h2")]}, config) + state = graph.get_state(config) + msgs = state.values["messages"] + # ai-1 + h2 + ai-1(second reply, same id overwrites) + h2 removed + # more simply: after second run we expect ai-1 updated + h2 remaining minus deleted h2 + # just assert h1 is still gone + assert all(m.id != "h1" for m in msgs), ( + "h1 should still be absent after second turn" + ) + + +async def test_delta_channel_update_by_id_end_to_end() -> None: + """Updating a message by ID via DeltaChannel must persist and reload correctly.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def update_msg(state: State) -> dict: + # re-send h1 with updated content + return {"messages": [HumanMessage(content="updated", id="h1")]} + + builder = StateGraph(State) + builder.add_node("update_msg", update_msg) + builder.add_edge(START, "update_msg") + graph = builder.compile(checkpointer=InMemorySaver()) + + config = {"configurable": {"thread_id": "diff-update-id-test"}} + graph.invoke({"messages": [HumanMessage(content="original", id="h1")]}, config) + + state = graph.get_state(config) + msgs = state.values["messages"] + assert len(msgs) == 1, f"expected 1 message, got {len(msgs)}: {msgs}" + assert msgs[0].content == "updated" + assert msgs[0].id == "h1" + + # Second turn: verify the updated state is the base for further accumulation + graph.invoke({"messages": [HumanMessage(content="new", id="h2")]}, config) + state = graph.get_state(config) + msgs = state.values["messages"] + ids = [m.id for m in msgs] + assert "h1" in ids # h1 persists (updated, not duplicated) + assert "h2" in ids + assert ids.count("h1") == 1, "h1 must not be duplicated" + + +async def test_delta_channel_durability_exit_stores_snapshot() -> None: + """DeltaChannel must reload from a durability='exit' checkpoint.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def respond(state: State) -> dict: + return {"messages": [AIMessage(content="reply", id="ai1")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "delta-exit-test"}} + + result = graph.invoke( + {"messages": [HumanMessage(content="hello", id="h1")]}, + config, + durability="exit", + ) + assert [m.content for m in result["messages"]] == ["hello", "reply"] + + state = graph.get_state(config) + assert [m.content for m in state.values["messages"]] == ["hello", "reply"] + + +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.""" + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + def respond(state: State) -> dict: + i = len(state["messages"]) + return {"messages": [AIMessage(content=f"r{i}", id=f"ai{i}")]} + + order: list[str] = [] + original_aput_writes = InMemorySaver.aput_writes + original_aput = InMemorySaver.aput + + async def tracked_aput_writes(self, config, writes, task_id, task_path=""): + result = await original_aput_writes(self, config, writes, task_id, task_path) + order.append("aput_writes") + return result + + async def tracked_aput(self, config, checkpoint, metadata, new_versions): + has_sentinel = any( + v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values() + ) + order.append("aput_sentinel" if has_sentinel else "aput_other") + return await original_aput(self, config, checkpoint, metadata, new_versions) + + InMemorySaver.aput_writes = tracked_aput_writes + InMemorySaver.aput = tracked_aput + try: + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "async-ordering-test"}} + + for i in range(3): + await graph.ainvoke( + {"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")]}, config + ) + + # Every aput_sentinel must be preceded by at least one aput_writes + for i, event in enumerate(order): + if event == "aput_sentinel": + preceding = order[:i] + assert "aput_writes" in preceding, ( + f"aput_sentinel at {i} had no preceding aput_writes: {order}" + ) + last_write_idx = max( + j for j, e in enumerate(order[:i]) if e == "aput_writes" + ) + assert last_write_idx < i, ( + f"aput_writes at {last_write_idx} should precede aput_sentinel at {i}: {order}" + ) + finally: + InMemorySaver.aput_writes = original_aput_writes + InMemorySaver.aput = original_aput + + state = await graph.aget_state(config) + assert len(state.values["messages"]) == 6 # 3 human + 3 AI diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b6ef6b088..19f2ce88c 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6101,6 +6101,36 @@ async def test_parent_command( ) +async def test_delta_channel_durability_exit_stores_snapshot_async() -> None: + """DeltaChannel must reload from an async durability='exit' checkpoint.""" + from langchain_core.messages import AIMessage + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import _messages_delta_reducer + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + async def respond(state: State) -> dict: + return {"messages": [AIMessage(content="reply", id="ai1")]} + + builder = StateGraph(State) + builder.add_node("respond", respond) + builder.add_edge(START, "respond") + graph = builder.compile(checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "delta-exit-async-test"}} + + result = await graph.ainvoke( + {"messages": [HumanMessage(content="hello", id="h1")]}, + config, + durability="exit", + ) + assert [m.content for m in result["messages"]] == ["hello", "reply"] + + state = await graph.aget_state(config) + assert [m.content for m in state.values["messages"]] == ["hello", "reply"] + + @NEEDS_CONTEXTVARS async def test_interrupt_subgraph(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict):