From 5548dbc3088eb15099e1222476b659aa392329b3 Mon Sep 17 00:00:00 2001 From: Will Fu-Hinthorn Date: Wed, 29 Apr 2026 13:30:59 -0700 Subject: [PATCH] fixup exit mode --- libs/checkpoint-postgres/pyproject.toml | 2 +- libs/langgraph/langgraph/channels/delta.py | 44 +++++++++++-------- libs/langgraph/langgraph/graph/message.py | 13 +++--- .../langgraph/langgraph/pregel/_checkpoint.py | 7 ++- libs/langgraph/langgraph/pregel/_loop.py | 1 + libs/langgraph/pyproject.toml | 4 +- libs/langgraph/tests/test_pregel.py | 31 +++++++++++++ libs/langgraph/tests/test_pregel_async.py | 30 +++++++++++++ 8 files changed, 105 insertions(+), 27 deletions(-) diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index e6f53189d..558292e72 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/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 2ca944e55..ed08f083f 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -26,23 +26,31 @@ 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 the full list of - new writes for that step in one call: - ``reducer(state, [write1, write2, ...]) -> new_state``. + The reducer receives the current accumulated value and a batch of writes + in one call: `reducer(state, [write1, write2, ...]) -> new_state`. - ``snapshot_frequency=None`` (default): pure delta — stores only - ``DELTA_SENTINEL`` in checkpoint blobs; reads replay all ancestor writes. + 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: - ``snapshot_frequency=N``: ``create_checkpoint`` writes a full - ``_DeltaSnapshot`` blob every N steps, bounding replay depth to N. + 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``. Receives the current - accumulated value and the list of all writes for this step. - typ: The value type (e.g. ``list``, ``dict``). Inferred automatically - from the outer type when used inside ``Annotated[T, DeltaChannel(...)]``. + 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. + `None` (default) = pure delta, never snapshot. """ __slots__ = ("value", "reducer", "snapshot_frequency") @@ -105,8 +113,8 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): """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. + * `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__( @@ -122,7 +130,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return new def replay_writes(self, writes: Sequence[PendingWrite]) -> None: - """Apply ancestor writes oldest→newest via a single reducer call. + """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 @@ -178,10 +186,10 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return self.value is not MISSING def checkpoint(self) -> Any: - """Return stored representation: always ``DELTA_SENTINEL``. + """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 + 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: diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 53388b812..3018d21da 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -247,13 +247,16 @@ def add_messages( def _messages_delta_reducer( state: list[AnyMessage], writes: list[list[AnyMessage]] ) -> list[AnyMessage]: - """**Experimental.** Batch reducer for use with ``DeltaChannel``. + """**Experimental.** Batch reducer for use with `DeltaChannel`. - Processes all writes for a step in one pass — dedup by ID, ``RemoveMessage`` - tombstoning — without calling ``add_messages``. Assumes writes contain - already-typed ``BaseMessage`` objects (no raw-dict coercion). + 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). - Use ``add_messages`` as the reducer for ``BinaryOperatorAggregate`` or + 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:: diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 41493f999..dc7675975 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -38,6 +38,7 @@ 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. @@ -47,6 +48,10 @@ def create_checkpoint( 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: @@ -61,7 +66,7 @@ def create_checkpoint( ch = channels[k] if ( isinstance(ch, DeltaChannel) - and ch.is_snapshot_step(step) + and (force_delta_snapshot or ch.is_snapshot_step(step)) and ch.is_available() ): # Eager snapshot: bump version if not already written this step diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 2a5aa5446..a3a3a4765 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -899,6 +899,7 @@ class PregelLoop: 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( diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 8114550f5..d35193762 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -24,8 +24,8 @@ classifiers = [ 'Programming Language :: Python :: 3.13', ] dependencies = [ - "langchain-core>=1.3.0,<2", - "langgraph-checkpoint>=2.1.0,<5.0.0", + "langchain-core>=1.3.2,<2", + "langgraph-checkpoint>=4.0.3,<5.0.0", "langgraph-sdk>=0.3.0,<0.4.0", "langgraph-prebuilt>=1.0.9,<1.1.0", "xxhash>=3.5.0", diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ee95fdfd7..677fe7cf3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9586,6 +9586,37 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: 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.""" + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + + from langgraph.graph import START, StateGraph + from langgraph.graph.message import _messages_delta_reducer + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + + 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.""" 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):