diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index cd5a14a19..537acd0fa 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -32,7 +32,7 @@ Conn = _internal.Conn # For backward compatibility class PostgresSaver(BasePostgresSaver): """Checkpointer that stores checkpoints in a Postgres database.""" - lock: threading.RLock + lock: threading.Lock def __init__( self, @@ -48,7 +48,7 @@ class PostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe - self.lock = threading.RLock() + self.lock = threading.Lock() self.supports_pipeline = Capabilities().has_pipeline() @classmethod @@ -179,7 +179,7 @@ class PostgresSaver(BasePostgresSaver): value["channel_values"], ) for value in values: - yield self._load_checkpoint_tuple(value) + yield self._load_checkpoint_tuple(value, cur) def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the database. @@ -250,7 +250,7 @@ class PostgresSaver(BasePostgresSaver): value["channel_values"], ) - return self._load_checkpoint_tuple(value) + return self._load_checkpoint_tuple(value, cur) def put( self, @@ -430,22 +430,26 @@ class PostgresSaver(BasePostgresSaver): with conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur - def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: + def _load_checkpoint_tuple( + self, value: DictRow, cur: Cursor[DictRow] + ) -> CheckpointTuple: """ Convert a database row into a CheckpointTuple object. Args: value: A row from the database containing checkpoint data. + cur: The cursor used by the caller; reused for DeltaChannel + reconstruction to avoid acquiring `self.lock` a second time. Returns: CheckpointTuple: A structured representation of the checkpoint, including its configuration, metadata, parent checkpoint (if any), and pending writes. """ - from langgraph.checkpoint.base import DeltaChannelSentinel + from langgraph.checkpoint.base import DELTA_SENTINEL channel_values = self._load_blobs(value["channel_values"]) - if any(isinstance(v, DeltaChannelSentinel) for v in channel_values.values()): + if any(v is DELTA_SENTINEL for v in channel_values.values()): cp_config = cast( RunnableConfig, { @@ -456,8 +460,7 @@ class PostgresSaver(BasePostgresSaver): } }, ) - with self._cursor() as cur: - self._resolve_delta_channels(cp_config, channel_values, cur) + self._resolve_delta_channels(cp_config, channel_values, cur) return CheckpointTuple( { "configurable": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 595e240f3..86070a909 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -8,12 +8,13 @@ from typing import Any from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, - DeltaChannelSentinel, + DeltaChannelWrites, get_checkpoint_id, get_serializable_checkpoint_metadata, ) @@ -169,7 +170,7 @@ class AsyncPostgresSaver(BasePostgresSaver): value["channel_values"], ) for value in values: - yield await self._load_checkpoint_tuple(value) + yield await self._load_checkpoint_tuple(value, cur) async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: """Get a checkpoint tuple from the database asynchronously. @@ -220,7 +221,7 @@ class AsyncPostgresSaver(BasePostgresSaver): value["channel_values"], ) - return await self._load_checkpoint_tuple(value) + return await self._load_checkpoint_tuple(value, cur) async def aput( self, @@ -444,12 +445,17 @@ class AsyncPostgresSaver(BasePostgresSaver): thread_id, checkpoint_ns, checkpoint_id, channel, cur ) - async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple: + async def _load_checkpoint_tuple( + self, value: DictRow, cur: AsyncCursor[DictRow] + ) -> CheckpointTuple: """ Convert a database row into a CheckpointTuple object. Args: value: A row from the database containing checkpoint data. + cur: The cursor used by the caller; reused for DeltaChannel + reconstruction to avoid re-entering `self.lock`, which is + `asyncio.Lock` and would deadlock. Returns: CheckpointTuple: A structured representation of the checkpoint, @@ -464,17 +470,13 @@ class AsyncPostgresSaver(BasePostgresSaver): channel_values: dict[str, Any] = {} if blob_values: channel_values = self._load_blobs(blob_values) - delta_channels = [ - ch - for ch, v in channel_values.items() - if isinstance(v, DeltaChannelSentinel) - ] - if delta_channels: - async with self._cursor() as cur: - for channel in delta_channels: - channel_values[channel] = await self._aget_channel_writes_cur( + for channel, v in channel_values.items(): + if v is DELTA_SENTINEL: + channel_values[channel] = DeltaChannelWrites( + await self._aget_channel_writes_cur( thread_id, checkpoint_ns, checkpoint_id, channel, cur ) + ) return CheckpointTuple( { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 0a17f052a..f03c5d63b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -9,10 +9,11 @@ from typing import Any, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + DELTA_SENTINEL, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, - DeltaChannelSentinel, + DeltaChannelWrites, get_checkpoint_id, ) from langgraph.checkpoint.serde.types import TASKS @@ -205,14 +206,16 @@ class BasePostgresSaver(BaseCheckpointSaver[str]): channel_values: dict[str, Any], cur: Any, ) -> None: - for channel, value in list(channel_values.items()): - if isinstance(value, DeltaChannelSentinel): - channel_values[channel] = self._get_channel_writes_cur( - config["configurable"]["thread_id"], - config["configurable"].get("checkpoint_ns", ""), - config["configurable"]["checkpoint_id"], - channel, - cur, + for channel, value in channel_values.items(): + if value is DELTA_SENTINEL: + channel_values[channel] = DeltaChannelWrites( + self._get_channel_writes_cur( + config["configurable"]["thread_id"], + config["configurable"].get("checkpoint_ns", ""), + config["configurable"]["checkpoint_id"], + channel, + cur, + ) ) def _get_channel_writes_cur( diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index f52b9b117..c32024445 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,14 +1,12 @@ from __future__ import annotations import copy -import dataclasses import logging import threading from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence -from typing import ( # noqa: UP035 +from typing import ( Any, Generic, - List, Literal, NamedTuple, TypedDict, @@ -22,28 +20,18 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_ from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( + DELTA_SENTINEL, ERROR, INTERRUPT, RESUME, SCHEDULED, ChannelProtocol, + DeltaChannelWrites, ) V = TypeVar("V", int, float, str) PendingWrite = tuple[str, str, Any] - -@dataclasses.dataclass -class DeltaChannelSentinel: - """Marker stored in checkpoint_blobs for a DeltaChannel field. - - No data is stored here — the actual per-step writes live in checkpoint_writes - and are replayed through the reducer at load time. - """ - - pass - - _DELTA_RECONSTRUCTION: threading.local = threading.local() logger = logging.getLogger(__name__) @@ -475,14 +463,14 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_channel_writes(self, config: RunnableConfig, channel: str) -> List[Any]: # noqa: UP006 + def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]: """Collect all writes for `channel` across this checkpoint's ancestry, oldest→newest. Default implementation walks the full thread history via `list()`. Savers can override with a more efficient query (InMemorySaver and PostgresSaver do this). """ # Guard against re-entrant calls: when list() triggers reconstruction which - # calls list() again, the inner call returns tuples with DeltaChannelSentinel + # calls list() again, the inner call returns tuples with DELTA_SENTINEL # in channel_values (which get_channel_writes ignores — it only reads # pending_writes). This breaks the recursion safely. if getattr(_DELTA_RECONSTRUCTION, "active", False): @@ -505,7 +493,7 @@ class BaseCheckpointSaver(Generic[V]): async def aget_channel_writes( self, config: RunnableConfig, channel: str - ) -> List[Any]: # noqa: UP006 + ) -> list[Any]: """Async version of get_channel_writes.""" if getattr(_DELTA_RECONSTRUCTION, "active", False): return [] diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index f265fe860..bc85e1020 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -20,7 +20,8 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, CheckpointTuple, - DeltaChannelSentinel, + DELTA_SENTINEL, + DeltaChannelWrites, SerializerProtocol, get_checkpoint_id, get_checkpoint_metadata, @@ -143,10 +144,14 @@ class InMemorySaver( config: RunnableConfig, channel_values: dict[str, Any], ) -> None: - """Replace DeltaChannelSentinel entries with reconstructed write lists.""" - for channel, value in list(channel_values.items()): - if isinstance(value, DeltaChannelSentinel): - channel_values[channel] = self.get_channel_writes(config, channel) + """Replace DELTA_SENTINEL entries with DeltaChannelWrites so + DeltaChannel.from_checkpoint can distinguish reconstructed writes from + a pre-DeltaChannel accumulated list.""" + for channel, value in channel_values.items(): + if value is DELTA_SENTINEL: + channel_values[channel] = DeltaChannelWrites( + self.get_channel_writes(config, channel) + ) def get_channel_writes(self, config: RunnableConfig, channel: str) -> list[Any]: thread_id = config["configurable"]["thread_id"] diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index a3bb23b59..48ef785ca 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -33,14 +33,13 @@ 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 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"" @@ -64,14 +63,6 @@ def _warn_once( logger.warning(msg, *args) -def _get_delta_sentinel_cls() -> type: - from langgraph.checkpoint.base import ( - DeltaChannelSentinel, - ) # lazy import avoids circular dep - - return DeltaChannelSentinel - - class JsonPlusSerializer(SerializerProtocol): """Serializer that uses ormsgpack, with optional fallbacks. @@ -260,12 +251,12 @@ class JsonPlusSerializer(SerializerProtocol): def dumps_typed(self, obj: Any) -> tuple[str, bytes]: if obj is None: return "null", EMPTY_BYTES + elif obj is DELTA_SENTINEL: + return "delta", EMPTY_BYTES elif isinstance(obj, bytes): return "bytes", obj elif isinstance(obj, bytearray): return "bytearray", obj - elif isinstance(obj, _get_delta_sentinel_cls()): - return "delta", b"" else: try: return "msgpack", _msgpack_enc(obj) @@ -289,9 +280,7 @@ class JsonPlusSerializer(SerializerProtocol): data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) elif type_ == "delta": - from langgraph.checkpoint.base import DeltaChannelSentinel - - return DeltaChannelSentinel() + return DELTA_SENTINEL elif self.pickle_fallback and type_ == "pickle": return pickle.loads(data_) else: diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 65a2b0c8e..7a8ba7484 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -1,3 +1,4 @@ +import dataclasses from collections.abc import Sequence from typing import ( Any, @@ -14,6 +15,34 @@ 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() + + +@dataclasses.dataclass +class DeltaChannelWrites: + """In-memory wrapper around per-step writes reconstructed by a saver. + Consumed by `DeltaChannel.from_checkpoint`. Never serialized — if this + reaches the wire, something upstream forgot to unwrap it. + """ + + writes: list[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 17a9a98b1..3384e5598 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -999,13 +999,14 @@ def test_msgpack_nested_pydantic_serializes_as_dict( assert result == obj -def test_delta_channel_sentinel_serde_round_trip() -> None: - from langgraph.checkpoint.base import DeltaChannelSentinel +def test_delta_sentinel_serde_round_trip() -> None: + from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer serde = JsonPlusSerializer() - original = DeltaChannelSentinel() - type_tag, blob = serde.dumps_typed(original) + type_tag, blob = serde.dumps_typed(DELTA_SENTINEL) + # Zero-byte "delta" tag — no allowlist change needed. assert type_tag == "delta" + assert blob == b"" loaded = serde.loads_typed((type_tag, blob)) - assert isinstance(loaded, DeltaChannelSentinel) + assert loaded is DELTA_SENTINEL diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 8c636bcb5..6890eeff8 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -324,9 +324,9 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None: class TestInMemorySaverDeltaChannel: def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None: - """_load_blobs returns DeltaChannelSentinel for delta channels (reconstruction deferred).""" + """_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred).""" from langgraph.checkpoint.base import ( - DeltaChannelSentinel, + DELTA_SENTINEL, empty_checkpoint, ) @@ -336,8 +336,7 @@ class TestInMemorySaverDeltaChannel: thread_id, ns, channel = "t1", "", "messages" v1 = "00000000000000000000000000000001.0000000000000000" - sentinel = DeltaChannelSentinel() - saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(sentinel) + saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL) cp1 = empty_checkpoint() cp1["id"] = "cp1" @@ -348,7 +347,7 @@ class TestInMemorySaverDeltaChannel: result = saver._load_blobs(thread_id, ns, {channel: v1}) assert channel in result - assert isinstance(result[channel], DeltaChannelSentinel) + assert result[channel] is DELTA_SENTINEL def test_get_channel_writes_collects_writes(self) -> None: """get_channel_writes collects per-step writes oldest→newest.""" diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 7a442bae0..9207aa2bb 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -119,12 +119,3 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC): Returns `True` if the channel was updated, `False` otherwise. """ return False - - def after_checkpoint(self, version: Any, checkpoint_id: str | None = None) -> None: - """Called after checkpoint() with the assigned version, and after - from_checkpoint() with the current channel version. - - No-op by default. Override in channels that track their own version - for incremental checkpointing (e.g. DeltaChannel). - """ - pass diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 3d5e02693..dbe1dddd5 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence from typing import Any, Generic -from langgraph.checkpoint.base import DeltaChannelSentinel +from langgraph.checkpoint.base import DELTA_SENTINEL, DeltaChannelWrites from typing_extensions import Self from langgraph._internal._typing import MISSING @@ -14,9 +14,14 @@ from langgraph.errors import EmptyChannelError __all__ = ("DeltaChannel",) -class DeltaChannel( - Generic[Value], BaseChannel[list[Value], Value, DeltaChannelSentinel] -): +def _empty(typ: Any) -> Any: + try: + return typ() + except Exception: + return [] + + +class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, Any]): """A channel that stores only a sentinel in checkpoints; per-step writes are stored in checkpoint_writes and replayed through the operator at load time. @@ -76,25 +81,20 @@ class DeltaChannel( new.typ = self.typ new.key = self.key if checkpoint is MISSING: - try: - new.value = new.typ() - except Exception: - new.value = [] - elif isinstance(checkpoint, list): - # Flat list of write values (oldest→newest) from get_channel_writes. - try: - value: Any = new.typ() - except Exception: - value = [] - for write in checkpoint: + new.value = _empty(new.typ) + elif isinstance(checkpoint, DeltaChannelWrites): + # Saver reconstructed per-step writes; replay through the operator. + value: Any = _empty(new.typ) + for write in checkpoint.writes: value = new.operator(value, write) new.value = value else: - # Backward compat: plain accumulated value (e.g. from a migrated thread). + # Backward compat: a pre-DeltaChannel thread stored the accumulated + # value directly. Trust it as-is. try: new.value = list(checkpoint) except Exception: - new.value = [] + new.value = _empty(new.typ) return new def update(self, values: Sequence[Any]) -> bool: @@ -133,5 +133,5 @@ class DeltaChannel( def is_available(self) -> bool: return self.value is not MISSING - def checkpoint(self) -> DeltaChannelSentinel: - return DeltaChannelSentinel() + def checkpoint(self) -> Any: + return DELTA_SENTINEL diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 8ebaba89f..3d510ac7d 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -1,6 +1,5 @@ from __future__ import annotations -import logging from collections.abc import Mapping from datetime import datetime, timezone @@ -13,8 +12,6 @@ from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec LATEST_VERSION = 4 -logger = logging.getLogger(__name__) - def empty_checkpoint() -> Checkpoint: return Checkpoint( @@ -70,12 +67,13 @@ def channels_from_checkpoint( channel_specs[k] = v else: managed_specs[k] = v - channels: dict[str, BaseChannel] = {} - for k, v in channel_specs.items(): - ch = v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) - ch.after_checkpoint(checkpoint["channel_versions"].get(k), checkpoint.get("id")) - channels[k] = ch - return channels, managed_specs + return ( + { + k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING)) + for k, v in channel_specs.items() + }, + 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 5e80f5d2b..976c1382b 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -891,12 +891,6 @@ class PregelLoop: id=self.checkpoint["id"] if exiting else None, updated_channels=self.updated_channels, ) - if do_checkpoint and self.channels: - for k, ch in self.channels.items(): - ch.after_checkpoint( - self.checkpoint["channel_versions"].get(k), - self.checkpoint.get("id"), - ) # sanitize TASK channel in the checkpoint before saving (durability=="exit") if TASKS in self.checkpoint["channel_values"] and any( isinstance(channel, UntrackedValue) for channel in self.channels.values() diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 3b4ebb3c2..24ff96571 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -121,7 +121,7 @@ def test_untracked_value() -> None: def test_delta_channel_basic_two_steps() -> None: from langchain_core.messages import AIMessage, HumanMessage - from langgraph.checkpoint.base import DeltaChannelSentinel + from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages @@ -131,12 +131,12 @@ def test_delta_channel_basic_two_steps() -> None: # Step 1: one message added ch.update([HumanMessage(content="hi", id="h1")]) d1 = ch.checkpoint() - assert isinstance(d1, DeltaChannelSentinel) + assert d1 is DELTA_SENTINEL # Step 2: another message ch.update([AIMessage(content="hello", id="a1")]) d2 = ch.checkpoint() - assert isinstance(d2, DeltaChannelSentinel) + assert d2 is DELTA_SENTINEL # Full accumulated value is preserved in memory assert len(ch.get()) == 2 @@ -145,19 +145,21 @@ def test_delta_channel_basic_two_steps() -> None: def test_delta_channel_from_checkpoint_writes_list() -> None: - """from_checkpoint with a flat list of individual writes replays them through the operator.""" + """from_checkpoint given DeltaChannelWrites replays through the operator.""" from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.base import DeltaChannelWrites from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages spec = DeltaChannel(add_messages) - # Each element is one write value (as stored in checkpoint_writes) - writes = [ - HumanMessage(content="hi", id="h1"), - AIMessage(content="hello", id="a1"), - HumanMessage(content="bye", id="h2"), - ] + writes = DeltaChannelWrites( + [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + HumanMessage(content="bye", id="h2"), + ] + ) ch = spec.from_checkpoint(writes) msgs = ch.get() assert len(msgs) == 3 @@ -181,7 +183,7 @@ def test_delta_channel_from_checkpoint_backwards_compat() -> None: def test_delta_channel_overwrite() -> None: from langchain_core.messages import HumanMessage - from langgraph.checkpoint.base import DeltaChannelSentinel + from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages @@ -192,7 +194,7 @@ def test_delta_channel_overwrite() -> None: ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) d = ch.checkpoint() - assert isinstance(d, DeltaChannelSentinel) + assert d is DELTA_SENTINEL # After overwrite, value is reset to only the new message assert len(ch.get()) == 1 assert ch.get()[0].content == "new" @@ -221,11 +223,15 @@ def test_delta_channel_remove_message_and_replay() -> None: assert ch.get() == [HumanMessage(content="hi", id="h1")] # Replay the writes list from scratch — must reproduce the post-remove state - writes = [ - HumanMessage(content="hi", id="h1"), - AIMessage(content="hello", id="a1"), - RemoveMessage(id="a1"), - ] + from langgraph.checkpoint.base import DeltaChannelWrites + + writes = DeltaChannelWrites( + [ + HumanMessage(content="hi", id="h1"), + AIMessage(content="hello", id="a1"), + RemoveMessage(id="a1"), + ] + ) ch2 = spec.from_checkpoint(writes) assert ch2.get() == [HumanMessage(content="hi", id="h1")] @@ -248,29 +254,33 @@ def test_delta_channel_update_by_id_and_replay() -> None: assert ch.get() == [HumanMessage(content="updated", id="h1")] # Replay writes — must produce the updated message, not the original - writes = [ - HumanMessage(content="original", id="h1"), - HumanMessage(content="updated", id="h1"), - ] + from langgraph.checkpoint.base import DeltaChannelWrites + + writes = DeltaChannelWrites( + [ + HumanMessage(content="original", id="h1"), + HumanMessage(content="updated", id="h1"), + ] + ) ch2 = spec.from_checkpoint(writes) assert len(ch2.get()) == 1 assert ch2.get()[0].content == "updated" def test_delta_channel_checkpoint_returns_sentinel() -> None: - """checkpoint() always returns DeltaChannelSentinel regardless of state.""" - from langgraph.checkpoint.base import DeltaChannelSentinel + """checkpoint() always returns DELTA_SENTINEL regardless of state.""" + from langgraph.checkpoint.base import DELTA_SENTINEL from langgraph.channels.delta import DeltaChannel from langgraph.graph.message import add_messages ch = DeltaChannel(add_messages).from_checkpoint(MISSING) - assert isinstance(ch.checkpoint(), DeltaChannelSentinel) + assert ch.checkpoint() is DELTA_SENTINEL from langchain_core.messages import HumanMessage ch.update([HumanMessage(content="hi", id="h1")]) - assert isinstance(ch.checkpoint(), DeltaChannelSentinel) + assert ch.checkpoint() is DELTA_SENTINEL def test_delta_channel_inmemory_saver_assembles_writes() -> None: @@ -304,16 +314,16 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None: graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config) graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config) - # get_tuple must return a resolved list (not DeltaChannelSentinel) - from langgraph.checkpoint.base import DeltaChannelSentinel + # get_tuple must return a DeltaChannelWrites wrapper (not the raw sentinel) + from langgraph.checkpoint.base import DELTA_SENTINEL, DeltaChannelWrites saved = saver.get_tuple(config) assert saved is not None assert "messages" in saved.checkpoint["channel_values"] - assert not isinstance( - saved.checkpoint["channel_values"]["messages"], DeltaChannelSentinel + assert saved.checkpoint["channel_values"]["messages"] is not DELTA_SENTINEL + assert isinstance( + saved.checkpoint["channel_values"]["messages"], DeltaChannelWrites ) - assert isinstance(saved.checkpoint["channel_values"]["messages"], list) state = graph.get_state(config) assert len(state.values["messages"]) == 4 # 2 human + 2 AI @@ -343,7 +353,7 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None: def test_delta_channel_dict_reducer_basic_updates() -> None: """DeltaChannel with a dict reducer accumulates key/value pairs across steps.""" - from langgraph.checkpoint.base import DeltaChannelSentinel + from langgraph.checkpoint.base import DELTA_SENTINEL def merge_dicts(left: dict, right: dict) -> dict: return {**left, **right} @@ -352,24 +362,24 @@ def test_delta_channel_dict_reducer_basic_updates() -> None: ch.update([{"a": 1}]) d1 = ch.checkpoint() - assert isinstance(d1, DeltaChannelSentinel) + assert d1 is DELTA_SENTINEL ch.update([{"b": 2}]) d2 = ch.checkpoint() - assert isinstance(d2, DeltaChannelSentinel) + assert d2 is DELTA_SENTINEL assert ch.get() == {"a": 1, "b": 2} def test_delta_channel_dict_reducer_writes_reconstruction() -> None: - """from_checkpoint with a writes list replays correctly through a dict merge reducer.""" + """from_checkpoint given DeltaChannelWrites replays through a dict merge reducer.""" + from langgraph.checkpoint.base import DeltaChannelWrites def merge_dicts(left: dict, right: dict) -> dict: return {**left, **right} spec = _delta_channel_with_type(merge_dicts, dict) - # Each element is one write value (oldest→newest) - writes = [{"a": 1}, {"b": 2}, {"c": 3}] + writes = DeltaChannelWrites([{"a": 1}, {"b": 2}, {"c": 3}]) ch = spec.from_checkpoint(writes) assert ch.get() == {"a": 1, "b": 2, "c": 3} @@ -398,10 +408,14 @@ def test_delta_channel_dict_reducer_with_deletions() -> None: assert ch.get() == {"file2.py": "content2", "file3.py": "content3"} # Confirm writes reconstruction produces the same result - writes = [ - {"file1.py": "content1", "file2.py": "content2"}, - {"file1.py": None, "file3.py": "content3"}, - ] + from langgraph.checkpoint.base import DeltaChannelWrites + + writes = DeltaChannelWrites( + [ + {"file1.py": "content1", "file2.py": "content2"}, + {"file1.py": None, "file3.py": "content3"}, + ] + ) spec = _delta_channel_with_type(merge_files, dict) ch2 = spec.from_checkpoint(writes) assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"} diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 3a641dddd..25dbf1d29 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -54,7 +54,9 @@ except ImportError: _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." + "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 = ( @@ -285,14 +287,13 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: rows.append((turns, b_bytes, d_bytes, b_rt, d_rt)) # ── Table 1: Storage ───────────────────────────────────────────────────── - W = 70 + W = 60 print("Storage (checkpoint blob bytes)") print("=" * W) print( f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12} {'savings':>8}" ) print("-" * W) - storage_results = [] for turns, b_bytes, d_bytes, b_rt, d_rt in rows: if b_bytes < 0: print( @@ -300,17 +301,15 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: ) else: ratio = b_bytes / d_bytes if d_bytes else float("inf") - storage_results.append((turns, b_bytes, d_bytes, ratio)) print( f"{turns:>6} {_approx_tokens(turns):>10} " - f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} " - f"{ratio:>7.0f}x" + f"{_fmt_bytes(b_bytes):>12} {_fmt_bytes(d_bytes):>12} {ratio:>7.0f}x" ) print("=" * W) print() # ── Table 2: Read latency ───────────────────────────────────────────────── - print("Read latency (avg of 5 get_state calls)") + print("Read latency (avg of 5 get_state calls = cost per invoke)") print("=" * W) print(f"{'turns':>6} {'ctx size':>10} {'add_msgs':>12} {'delta':>12}") print("-" * W) @@ -322,21 +321,9 @@ def _run_benchmark_for_checkpointer(cp_hint: Any) -> None: print("=" * W) print() - if storage_results: - turns, b_bytes, d_bytes, ratio = storage_results[-1] - b_rt = rows[-1][-2] - d_rt = rows[-1][-1] - print( - f"At {turns} turns: {_fmt_bytes(b_bytes)} → {_fmt_bytes(d_bytes)} ({ratio:.0f}x less storage); " - f"read {b_rt * 1000:.1f}ms → {d_rt * 1000:.1f}ms" - ) - print() - print("Legend:") print(" add_msgs = Annotated[list, add_messages] — O(N²) storage") - print( - " delta = DeltaChannel(add_messages) — O(N) storage, full chain replay" - ) + print(" delta = DeltaChannel(add_messages) — O(N) storage") print() diff --git a/libs/langgraph/tests/test_time_travel.py b/libs/langgraph/tests/test_time_travel.py index caf51f5a5..fe544274f 100644 --- a/libs/langgraph/tests/test_time_travel.py +++ b/libs/langgraph/tests/test_time_travel.py @@ -1161,9 +1161,7 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id( assert called == ["step_a", "ask_human"] # Resume with explicit head checkpoint_id in config - head_checkpoint_id = graph.get_state(config).config["configurable"][ - "checkpoint_id" - ] + head_checkpoint_id = graph.get_state(config).config["configurable"]["checkpoint_id"] called.clear() resume_config = { "configurable": {