release(checkpoint-sqlite): 3.1.2 (#9288)

Version bump only. No library code changes in this PR.

Unreleased since `checkpointsqlite==3.1.1`:

- #8617 - widen `Store` `put` value type to `Mapping[str, Any]`
- #8557 - walk delta ancestors by parent pointer
- #8544 - order delta channel replay by task path
- #9217 - validate namespace labels before store dispatch
- #9262 - treat an empty `checkpoint_id` as the latest in delta history

Follows `checkpoint==4.3.0`, which is already live on PyPI. #8544 raised
the floor to `langgraph-checkpoint>=4.3.0`, so this release needs it.

libs/checkpoint-sqlite/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-checkpoint-sqlite"
-version = "3.1.1"
+version = "3.1.2"
 description = "Library with a SQLite implementation of LangGraph checkpoint saver."
 authors = []
 requires-python = ">=3.10"

libs/langgraph/langgraph/graph/message.py

@@ -24,6 +24,9 @@ from typing_extensions import TypedDict, deprecated
 
 from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
 from langgraph.graph.state import StateGraph
+from langgraph.pregel._messages import (
+    _messages_delta_reducer as _messages_delta_reducer,
+)
 from langgraph.warnings import LangGraphDeprecatedSinceV10
 
 __all__ = (
@@ -244,71 +247,6 @@ 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`.
-
-    This reducer is batching-invariant, as required by `DeltaChannel`:
-    `reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
-
-    Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
-    objects so that HTTP-driven graphs work without a separate coercion
-    step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
-    unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
-    `BaseMessageChunk` conversion are not handled here.
-
-    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)]
-    """
-
-    # Each write is either a list of message-likes or a single message-like
-    # (BaseMessage / dict / str / tuple). Only lists flatten; everything
-    # else is one message.
-    flat: list[Any] = []
-    for w in writes:
-        if isinstance(w, list):
-            flat.extend(w)
-        else:
-            flat.append(w)
-    # Steady state: the reducer's own output is already typed, so skip
-    # `convert_to_messages` on state when the first element is a BaseMessage.
-    # Only raw input (initial dicts, deserialized blobs) hits the slow path.
-    if state and isinstance(state[0], BaseMessage):
-        state_msgs = state
-    else:
-        state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
-    msgs = cast("list[AnyMessage]", convert_to_messages(flat))
-
-    index: dict[str, int] = {
-        m.id: i for i, m in enumerate(state_msgs) if m.id is not None
-    }
-    result: list[AnyMessage | None] = list(state_msgs)
-    for msg in msgs:
-        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,

libs/langgraph/langgraph/pregel/_loop.py

@@ -161,6 +161,12 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
     return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
 
 
+def _cacheable(writes: WritesT) -> bool:
+    # A cache hit skips the node, so a cached error would skip it without
+    # raising: only a node that finished goes to the cache.
+    return not any(c in (INTERRUPT, ERROR) for c, _ in writes)
+
+
 class PregelLoop:
     config: RunnableConfig
     store: BaseStore | None
@@ -443,13 +449,18 @@ class PregelLoop:
             return None
         return self._graph_lifecycle_events.popleft()
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
         if not writes:
             return
         # deduplicate writes to special channels, last write wins
         if all(w[0] in WRITES_IDX_MAP for w in writes):
             writes = list({w[0]: w for w in writes}.values())
+        # Before the writes are saved, kept in `checkpoint_pending_writes` or
+        # applied, so all of them hold the same message ids.
+        ensure_message_ids(self.specs, writes)
         if task_id == NULL_TASK_ID:
             # writes for the null task are accumulated
             self.checkpoint_pending_writes = [
@@ -486,14 +497,6 @@ class PregelLoop:
 
         # save writes
         self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
-        # Assign stable IDs to any id=None BaseMessages in DeltaChannel writes
-        # before the background thread serialises them. Without this, reducers
-        # that assign IDs inside apply_writes() race with serialisation and
-        # store id=None, causing get_state() replays to produce a different UUID
-        # on every call.
-        for c, v in writes_to_save:
-            if isinstance(self.specs.get(c), DeltaChannel):
-                ensure_message_ids(v)
         if self.durability != "exit" and self.checkpointer_put_writes is not None:
             config = patch_configurable(
                 self.checkpoint_config,
@@ -536,7 +539,7 @@ class PregelLoop:
                 self._error_handler_write_futs.append(fut)
         # output writes
         if hasattr(self, "tasks"):
-            self.output_writes(task_id, writes)
+            self.output_writes(task_id, writes, cached=cached)
 
     def _put_pending_writes(self) -> None:
         if self.checkpointer_put_writes is None:
@@ -685,10 +688,6 @@ class PregelLoop:
             self.status = "done"
             return False
 
-        if self.control is not None and self.control.drain_requested:
-            self.status = "draining"
-            return False
-
         # if there are pending writes from a previous loop, apply them
         if self._reapplies_pending_writes and self.checkpoint_pending_writes:
             self._reapply_writes_to_succeeded_nodes(self.tasks)
@@ -701,6 +700,13 @@ class PregelLoop:
             self.status = "interrupt_before"
             raise GraphInterrupt()
 
+        # stop before running the next tasks if a drain was requested. after
+        # the breakpoint check: a resume passes the next step's breakpoints,
+        # so draining in front of one would skip it
+        if self.control is not None and self.control.drain_requested:
+            self.status = "draining"
+            return False
+
         # produce debug output
         self._emit("tasks", map_debug_tasks, self.tasks.values())
 
@@ -1012,6 +1018,8 @@ class PregelLoop:
             # save writes
             for tid, ws in writes.items():
                 if forks:
+                    # These skip `put_writes`, which gives messages their ids.
+                    ensure_message_ids(self.specs, ws)
                     carried.extend((tid, c, v) for c, v in ws)
                 else:
                     self.put_writes(tid, ws)
@@ -1076,6 +1084,9 @@ class PregelLoop:
             )
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
+            # Before the input is applied, so this run reads the message ids
+            # that are saved with it.
+            ensure_message_ids(self.specs, input_writes)
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
@@ -1360,6 +1371,14 @@ class PregelLoop:
         ):
             return
 
+        # A checkpoint a `checkpoint_id` addressed may already have children,
+        # which read the writes stored on it, so then the channels this run
+        # wrote snapshot instead.
+        if self._has_persisted_parent and not self._loaded_latest:
+            self._delta_channels_forced_snapshot.update(
+                w[3] for w in self._exit_delta_writes
+            )
+
         counters = dict(
             self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
         )
@@ -1699,7 +1718,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in self.match_cached_writes():
-                self.output_writes(task.id, task.writes, cached=True)
+                self.put_writes(task.id, task.writes, cached=True)
         return pushed
 
     def schedule_error_handler(
@@ -1736,16 +1755,18 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         if self._reapplies_pending_writes:
             self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
         for task in self.match_cached_writes():
-            self.output_writes(task.id, task.writes, cached=True)
+            self.put_writes(task.id, task.writes, cached=True)
         return handler_task
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
-        super().put_writes(task_id, writes)
-        if not writes or self.cache is None or not hasattr(self, "tasks"):
+        super().put_writes(task_id, writes, cached=cached)
+        if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
             return
         task = self.tasks.get(task_id)
-        if task is None or task.cache_key is None:
+        if task is None or task.cache_key is None or not _cacheable(writes):
             return
         self.submit(
             self.cache.set,
@@ -1953,7 +1974,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     ) -> PregelExecutableTask | None:
         if pushed := super().accept_push(task, write_idx, call):
             for task in await self.amatch_cached_writes():
-                self.output_writes(task.id, task.writes, cached=True)
+                self.put_writes(task.id, task.writes, cached=True)
         return pushed
 
     async def aschedule_error_handler(
@@ -1990,19 +2011,18 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         if self._reapplies_pending_writes:
             self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
         for task in await self.amatch_cached_writes():
-            self.output_writes(task.id, task.writes, cached=True)
+            self.put_writes(task.id, task.writes, cached=True)
         return handler_task
 
-    def put_writes(self, task_id: str, writes: WritesT) -> None:
+    def put_writes(
+        self, task_id: str, writes: WritesT, *, cached: bool = False
+    ) -> None:
         """Put writes for a task, to be read by the next tick."""
-        super().put_writes(task_id, writes)
-        if not writes or self.cache is None or not hasattr(self, "tasks"):
+        super().put_writes(task_id, writes, cached=cached)
+        if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
             return
         task = self.tasks.get(task_id)
-        if task is None or task.cache_key is None:
-            return
-        if writes[0][0] in (INTERRUPT, ERROR):
-            # only cache successful tasks
+        if task is None or task.cache_key is None or not _cacheable(writes):
             return
         self.submit(
             self.cache.aset,

libs/langgraph/langgraph/pregel/_messages.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from collections.abc import (
+    AsyncIterator,
+    Callable,
+    Iterator,
+    Mapping,
+    MutableSequence,
+    Sequence,
+)
 from dataclasses import fields, is_dataclass
 from typing import (
     Any,
@@ -10,16 +17,18 @@ from typing import (
 from uuid import UUID, uuid4
 
 from langchain_core.callbacks import BaseCallbackHandler
-from langchain_core.messages import BaseMessage, ToolMessage
+from langchain_core.messages import AnyMessage, BaseMessage, RemoveMessage, ToolMessage
 from langchain_core.messages.utils import convert_to_messages
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
 from pydantic import BaseModel
 
 from langgraph._internal._config import filter_to_user_tags
 from langgraph._internal._constants import NS_SEP
+from langgraph.channels.binop import _get_overwrite
+from langgraph.channels.delta import DeltaChannel
 from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
 from langgraph.pregel.protocol import StreamChunk
-from langgraph.types import Command
+from langgraph.types import Command, Overwrite
 
 try:
     from langchain_core.tracers._streaming import _StreamingCallbackHandler
@@ -408,6 +417,74 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
             self.stream((meta[0], "messages", (event, v2_meta)))
 
 
+# Defined here rather than in `langgraph.graph.message`, which re-exports it:
+# `ensure_message_ids` has to recognize it, and importing that module from the
+# loop would be an import cycle.
+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`.
+
+    This reducer is batching-invariant, as required by `DeltaChannel`:
+    `reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
+
+    Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
+    objects so that HTTP-driven graphs work without a separate coercion
+    step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
+    unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
+    `BaseMessageChunk` conversion are not handled here.
+
+    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)]
+    """
+
+    # Each write is either a list of message-likes or a single message-like
+    # (BaseMessage / dict / str / tuple). Only lists flatten; everything
+    # else is one message.
+    flat: list[Any] = []
+    for w in writes:
+        if isinstance(w, list):
+            flat.extend(w)
+        else:
+            flat.append(w)
+    # Steady state: the reducer's own output is already typed, so skip
+    # `convert_to_messages` on state when the first element is a BaseMessage.
+    # Only raw input (initial dicts, deserialized blobs) hits the slow path.
+    if state and isinstance(state[0], BaseMessage):
+        state_msgs = state
+    else:
+        state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
+    msgs = cast("list[AnyMessage]", convert_to_messages(flat))
+
+    index: dict[str, int] = {
+        m.id: i for i, m in enumerate(state_msgs) if m.id is not None
+    }
+    result: list[AnyMessage | None] = list(state_msgs)
+    for msg in msgs:
+        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]
+
+
 # Known role values (OpenAI-style) and type values (LangChain serialisation)
 # that identify a dict as a message. Checked before coercing to BaseMessage so
 # we don't accidentally touch unrelated dicts that happen to have a "role" key.
@@ -423,39 +500,103 @@ def _is_message_dict(item: dict) -> bool:
     return item.get("role") in _MESSAGE_ROLES or item.get("type") in _MESSAGE_TYPES
 
 
-def ensure_message_ids(value: Any) -> None:
-    """Coerce message-like write values to typed BaseMessages with stable IDs.
+def ensure_message_ids(
+    channels: Mapping[str, Any], writes: Sequence[tuple[str, Any]]
+) -> None:
+    """Give the messages written to DeltaChannels stable ids.
 
-    Called in put_writes() before DeltaChannel writes are submitted to the
-    checkpointer. Without this the checkpoint may store raw dicts or id=None
-    BaseMessages; every get_state() replay then produces a different UUID and
-    the same message appears with a different ID in each LangSmith trace.
+    A `DeltaChannel` replays its saved writes on every read, and its reducer
+    may not assign missing ids (`_messages_delta_reducer` doesn't), so a
+    message saved without an id comes back without one, or with a new one on
+    every read, and nothing can target it by id, for example `RemoveMessage`.
 
-    Handles three input shapes:
-    - BaseMessage objects: assign a UUID if id is None.
-    - Dicts with a known "role" (OpenAI-style) or "type" (LangChain format) at
-      the root level: stamp "id" into the dict in-place. The reducer's
-      convert_to_messages call will forward the id to the resulting BaseMessage.
-    - Lists of the above: apply the same logic to each element, replacing dict
-      items with coerced BaseMessages so the shared list reference seen by
-      checkpoint_pending_writes and the background thread both get typed messages.
+    In a channel whose reducer is `_messages_delta_reducer`, every message-like
+    it accepts (a message, a dict, a string, a `(role, content)` tuple; alone,
+    in a list or in an `Overwrite`) becomes a message with an id. Other
+    DeltaChannels may hold other data, so there only message objects get ids,
+    plus message dicts outside an `Overwrite`: one in a list becomes a message.
+    Anything that can't become a message is left as it is, so it fails where
+    it always has: a message dict in a list raises here, anything else when
+    the reducer reads it. An empty id counts as missing.
 
-    Mutating synchronously here (before the background thread is submitted) is
-    safe: the serialised bytes always reflect the post-coercion state.
+    A write that changes is replaced in `writes`: lists and dicts the caller
+    passed in are never changed, only the ids of message objects.
     """
-    if isinstance(value, BaseMessage):
-        if value.id is None:
-            value.id = str(uuid4())
-    elif isinstance(value, dict) and _is_message_dict(value):
-        if not value.get("id"):
-            value["id"] = str(uuid4())
-    elif isinstance(value, list):
-        for i, item in enumerate(value):
-            if isinstance(item, BaseMessage):
-                if item.id is None:
-                    item.id = str(uuid4())
-            elif isinstance(item, dict) and _is_message_dict(item):
-                msg = convert_to_messages([item])[0]
-                if msg.id is None:
-                    msg.id = str(uuid4())
-                value[i] = msg
+    for i, (ch, value) in enumerate(writes):
+        spec = channels.get(ch)
+        if not isinstance(spec, DeltaChannel):
+            continue
+        if spec.reducer is _messages_delta_reducer:
+            new_value = _with_message_ids(value)
+        else:
+            new_value = _with_base_message_ids(value)
+        if new_value is not value:
+            cast(MutableSequence[tuple[str, Any]], writes)[i] = (ch, new_value)
+
+
+def _with_message_ids(value: Any) -> Any:
+    is_overwrite, overwrite_value = _get_overwrite(value)
+    if not is_overwrite:
+        if isinstance(value, list):
+            return _map_messages(value, _list_item_with_id)
+        return _message_with_id(value)
+    if isinstance(overwrite_value, (list, tuple)):
+        items = _map_messages(overwrite_value, _message_with_id)
+        if items is not overwrite_value:
+            return Overwrite(items)
+    return value
+
+
+def _with_base_message_ids(value: Any) -> Any:
+    is_overwrite, overwrite_value = _get_overwrite(value)
+    if is_overwrite:
+        if isinstance(overwrite_value, (list, tuple)):
+            for item in overwrite_value:
+                _stamp_message_id(item)
+        return value
+    if isinstance(value, list):
+        return _map_messages(value, _base_list_item_with_id)
+    if isinstance(value, dict) and _is_message_dict(value) and not value.get("id"):
+        return {**value, "id": str(uuid4())}
+    _stamp_message_id(value)
+    return value
+
+
+def _map_messages(items: list | tuple, fn: Callable[[Any], Any]) -> list | tuple:
+    mapped = [fn(item) for item in items]
+    if all(new is old for new, old in zip(mapped, items)):
+        return items
+    return mapped if isinstance(items, list) else tuple(mapped)
+
+
+def _list_item_with_id(item: Any) -> Any:
+    if isinstance(item, dict) and _is_message_dict(item):
+        return _message_with_id(item, strict=True)
+    return _message_with_id(item)
+
+
+def _base_list_item_with_id(item: Any) -> Any:
+    if isinstance(item, dict) and _is_message_dict(item):
+        return _message_with_id(item, strict=True)
+    _stamp_message_id(item)
+    return item
+
+
+def _message_with_id(item: Any, *, strict: bool = False) -> Any:
+    if isinstance(item, BaseMessage):
+        _stamp_message_id(item)
+        return item
+    try:
+        message = convert_to_messages([item])[0]
+    except Exception:
+        if strict:
+            raise
+        return item
+    if not message.id:
+        message.id = str(uuid4())
+    return message
+
+
+def _stamp_message_id(item: Any) -> None:
+    if isinstance(item, BaseMessage) and item.id is None:
+        item.id = str(uuid4())

libs/langgraph/langgraph/pregel/_retry.py

@@ -7,7 +7,7 @@ import sys
 import threading
 import time
 import weakref
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Callable, Sequence
 from contextlib import suppress
 from dataclasses import dataclass, replace
 from datetime import datetime, timedelta, timezone
@@ -686,8 +686,6 @@ async def arun_with_retry(
     task: PregelExecutableTask,
     retry_policy: Sequence[RetryPolicy] | None,
     stream: bool = False,
-    match_cached_writes: Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
-    | None = None,
     configurable: dict[str, Any] | None = None,
 ) -> None:
     """Run a task asynchronously with retries."""
@@ -711,11 +709,6 @@ async def arun_with_retry(
                 )
             },
         )
-    if match_cached_writes is not None and task.cache_key is not None:
-        for t in await match_cached_writes():
-            if t is task:
-                # if the task is already cached, return
-                return
     while True:
         runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
         if isinstance(runtime, Runtime):

libs/langgraph/langgraph/pregel/main.py

@@ -2059,7 +2059,7 @@ class Pregel(
                 checkpoint_config,
                 self.channels,
                 [
-                    (task_id, [w for w in task.writes if w[0] != PUSH], task)
+                    (task_id, task.writes, task)
                     for task_id, task in zip(run_task_ids, run_tasks)
                 ],
                 fork_pending,
@@ -2557,7 +2557,7 @@ class Pregel(
                 checkpoint_config,
                 self.channels,
                 [
-                    (task_id, [w for w in task.writes if w[0] != PUSH], task)
+                    (task_id, task.writes, task)
                     for task_id, task in zip(run_task_ids, run_tasks)
                 ],
                 fork_pending,
@@ -3054,7 +3054,7 @@ class Pregel(
                 # with channel updates applied only at the transition between steps.
                 while loop.tick():
                     for task in loop.match_cached_writes():
-                        loop.output_writes(task.id, task.writes, cached=True)
+                        loop.put_writes(task.id, task.writes, cached=True)
                     for _ in runner.tick(
                         [t for t in loop.tasks.values() if not t.writes],
                         timeout=self.step_timeout,
@@ -3525,7 +3525,7 @@ class Pregel(
                 try:
                     while loop.tick():
                         for task in await loop.amatch_cached_writes():
-                            loop.output_writes(task.id, task.writes, cached=True)
+                            loop.put_writes(task.id, task.writes, cached=True)
                         async for _ in runner.atick(
                             [t for t in loop.tasks.values() if not t.writes],
                             timeout=self.step_timeout,
@@ -4274,22 +4274,6 @@ def _task_path_kwarg(put_writes: Callable[..., Any], task: PregelTaskWrites) ->
 _UpdateWrites = Sequence[tuple[str, Sequence[tuple[str, Any]], PregelTaskWrites | None]]
 
 
-def _delta_writes(
-    channels: Mapping[str, BaseChannel | ManagedValueSpec], writes: _UpdateWrites
-) -> list[tuple[str, Any]]:
-    """Give the update's DeltaChannel writes message ids, as the loop's
-    `put_writes` does, so every read of them returns the same ids."""
-    delta = [
-        (ch, value)
-        for _, task_writes, _ in writes
-        for ch, value in task_writes
-        if isinstance(channels.get(ch), DeltaChannel)
-    ]
-    for _, value in delta:
-        ensure_message_ids(value)
-    return delta
-
-
 def _store_or_fork_delta_writes(
     checkpointer: BaseCheckpointSaver,
     config: RunnableConfig,
@@ -4307,19 +4291,28 @@ def _store_or_fork_delta_writes(
     ancestors, so the writes go on the checkpoint the update builds on. If the
     thread already moved past that checkpoint, its other children would read
     them too, so the new checkpoint snapshots those channels instead. `writes`
-    holds `(task_id, writes, task)` per task; `task` is `None` for input.
+    holds `(task_id, writes, task)` per task; `task` is `None` for input. A
+    task's `PUSH` writes are left out: the update saves them on the checkpoint
+    it creates. The messages in `writes` get their ids first, in place, as the
+    loop's `put_writes` does, so pass the writes the update then applies.
     """
-    delta = _delta_writes(channels, writes)
+    for _, task_writes, _ in writes:
+        ensure_message_ids(channels, task_writes)
     if saved is None:
         return
     if is_first and checkpoint_superseded(checkpointer, config, saved):
-        fork_pending.update(ch for ch, _ in delta)
+        fork_pending.update(
+            ch
+            for _, task_writes, _ in writes
+            for ch, _ in task_writes
+            if isinstance(channels.get(ch), DeltaChannel)
+        )
         return
     for task_id, task_writes, task in writes:
-        if task_writes:
+        if channel_writes := [w for w in task_writes if w[0] != PUSH]:
             checkpointer.put_writes(
                 checkpoint_config,
-                task_writes,
+                channel_writes,
                 task_id,
                 **(_task_path_kwarg(checkpointer.put_writes, task) if task else {}),
             )
@@ -4337,17 +4330,23 @@ async def _astore_or_fork_delta_writes(
     is_first: bool,
 ) -> None:
     """Async `_store_or_fork_delta_writes`."""
-    delta = _delta_writes(channels, writes)
+    for _, task_writes, _ in writes:
+        ensure_message_ids(channels, task_writes)
     if saved is None:
         return
     if is_first and await acheckpoint_superseded(checkpointer, config, saved):
-        fork_pending.update(ch for ch, _ in delta)
+        fork_pending.update(
+            ch
+            for _, task_writes, _ in writes
+            for ch, _ in task_writes
+            if isinstance(channels.get(ch), DeltaChannel)
+        )
         return
     for task_id, task_writes, task in writes:
-        if task_writes:
+        if channel_writes := [w for w in task_writes if w[0] != PUSH]:
             await checkpointer.aput_writes(
                 checkpoint_config,
-                task_writes,
+                channel_writes,
                 task_id,
                 **(_task_path_kwarg(checkpointer.aput_writes, task) if task else {}),
             )

libs/langgraph/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph"
-version = "1.2.14"
+version = "1.2.15"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 requires-python = ">=3.10"

libs/langgraph/tests/test_delta_channel_exit_mode.py

@@ -625,3 +625,33 @@ def test_exit_resume_retried_after_its_final_checkpoint_fails_reruns_the_resumed
     state = graph.get_state(config)
     assert state.values["log"] == state.values["plain"]
     assert sorted(state.values["plain"]) == ["ask", "done", "in"]
+
+
+class _TurnState(_ResumeState):
+    turn: int
+
+
+@pytest.mark.parametrize("exit_runs", [1, 2])
+def test_exit_runs_from_a_checkpoint_the_thread_moved_past_leave_its_newer_branch_alone(
+    sync_checkpointer: BaseCheckpointSaver, exit_runs: int
+) -> None:
+    builder = StateGraph(_TurnState)
+    builder.add_node("a", lambda state: _both(f"a{state['turn']}"))
+    builder.add_node("b", lambda state: _both(f"b{state['turn']}"))
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    graph = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"turn": 1, **_both("in1")}, config)
+    end_of_turn_1 = graph.get_state(config).config
+    graph.invoke({"turn": 2, **_both("in2")}, config)
+
+    for turn in range(3, 3 + exit_runs):
+        graph.invoke(
+            {"turn": turn, **_both(f"in{turn}")}, end_of_turn_1, durability="exit"
+        )
+
+    history = list(graph.get_state_history(config))
+    assert [s.values.get("log", []) for s in history] == [
+        s.values.get("plain", []) for s in history
+    ]

libs/langgraph/tests/test_delta_channel_id_stability.py

@@ -16,7 +16,11 @@ from langgraph.checkpoint.memory import InMemorySaver
 from typing_extensions import TypedDict
 
 from langgraph.channels.delta import DeltaChannel
+from langgraph.channels.last_value import LastValue
 from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import _messages_delta_reducer
+from langgraph.pregel import NodeBuilder, Pregel
+from langgraph.types import Command, Durability, Overwrite
 
 pytestmark = pytest.mark.anyio
 
@@ -34,18 +38,19 @@ def _append_reducer(
     return result
 
 
-def _build_graph(checkpointer: Any) -> Any:
+def _build_graph(
+    checkpointer: Any,
+    *,
+    reducer: Any = _append_reducer,
+    write: Any = lambda: [AIMessage(content="reply", id="ai-1")],
+) -> Any:
     State = TypedDict(  # noqa: UP013
         "State",
-        {
-            "messages": Annotated[
-                list, DeltaChannel(_append_reducer, snapshot_frequency=50)
-            ]
-        },
+        {"messages": Annotated[list, DeltaChannel(reducer, snapshot_frequency=50)]},
     )  # type: ignore[call-overload]
 
     def agent(state: dict) -> dict:
-        return {"messages": [AIMessage(content="reply", id="ai-1")]}
+        return {"messages": write()}
 
     return (
         StateGraph(State)
@@ -138,3 +143,193 @@ def test_delta_channel_dict_style_message_gets_stable_id() -> None:
     assert len(set(ids)) == 1, (
         f"dict-style HumanMessage id must be stable across get_state() calls; got {ids}"
     )
+
+
+@pytest.mark.parametrize("container", [list, tuple])
+def test_delta_channel_message_in_an_overwrite_gets_stable_id(container: type) -> None:
+    graph = _build_graph(InMemorySaver())
+    config = {"configurable": {"thread_id": "overwrite-id-stability"}}
+
+    graph.invoke(
+        {"messages": Overwrite(container([HumanMessage(content="hello")]))}, config
+    )
+
+    ids = [
+        next(
+            m.id
+            for m in graph.get_state(config).values["messages"]
+            if isinstance(m, HumanMessage)
+        )
+        for _ in range(3)
+    ]
+
+    assert ids[0] is not None, "a message inside an Overwrite should get a UUID"
+    assert len(set(ids)) == 1, (
+        f"its id must be stable across get_state() calls; got {ids}"
+    )
+
+
+# Writes of "hello" in every shape `_messages_delta_reducer` accepts.
+MESSAGE_SHAPES = [
+    pytest.param(lambda: [HumanMessage("hello")], id="message"),
+    pytest.param(lambda: HumanMessage("hello"), id="bare message"),
+    pytest.param(lambda: [{"role": "user", "content": "hello"}], id="dict"),
+    pytest.param(lambda: {"type": "human", "content": "hello"}, id="bare dict"),
+    pytest.param(
+        lambda: [{"role": "user", "content": "hello", "id": ""}], id="dict, empty id"
+    ),
+    pytest.param(lambda: [("user", "hello")], id="tuple"),
+    pytest.param(lambda: ("user", "hello"), id="bare tuple"),
+    pytest.param(lambda: ["hello"], id="string"),
+    pytest.param(lambda: "hello", id="bare string"),
+    pytest.param(lambda: Overwrite(["hello"]), id="Overwrite"),
+    pytest.param(lambda: Overwrite(("hello",)), id="Overwrite of a tuple"),
+    pytest.param(lambda: {"__overwrite__": [("user", "hello")]}, id="Overwrite dict"),
+]
+
+
+def _hello_ids(graph: Any, config: dict) -> list[str | None]:
+    return [
+        next(
+            m.id
+            for m in graph.get_state(config).values["messages"]
+            if m.content == "hello"
+        )
+        for _ in range(2)
+    ]
+
+
+@pytest.mark.parametrize("write", MESSAGE_SHAPES)
+def test_input_message_of_any_shape_gets_stable_id(write: Any) -> None:
+    graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke({"messages": write()}, config)
+
+    first, second = _hello_ids(graph, config)
+    assert first, "the message should get a UUID"
+    assert first == second, "its id must be the same on every read"
+
+
+@pytest.mark.parametrize("write", MESSAGE_SHAPES)
+def test_node_message_of_any_shape_gets_stable_id(write: Any) -> None:
+    graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer, write=write)
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke({"messages": []}, config)
+
+    first, second = _hello_ids(graph, config)
+    assert first, "the message should get a UUID"
+    assert first == second, "its id must be the same on every read"
+
+
+def test_a_list_written_twice_adds_its_messages_twice() -> None:
+    turn = [("user", "hello")]
+    graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke({"messages": turn}, config)
+    graph.invoke({"messages": turn}, config)
+
+    hellos = [
+        m.id for m in graph.get_state(config).values["messages"] if m.content == "hello"
+    ]
+    assert len(set(hellos)) == 2, f"each write should add a message; got {hellos}"
+    assert turn == [("user", "hello")], "the caller's list should be left as it was"
+
+
+def test_a_forked_command_update_gets_a_stable_id() -> None:
+    graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
+    before_agent = next(
+        s.config for s in graph.get_state_history(config) if s.next == ("agent",)
+    )
+
+    graph.invoke(Command(update={"messages": "hello"}), before_agent)
+
+    first, second = _hello_ids(graph, config)
+    assert first, "the message should get a UUID"
+    assert first == second, "its id must be the same on every read"
+
+
+def test_a_message_dict_the_reducer_rejects_fails_the_run_before_it_is_saved() -> None:
+    graph = _build_graph(
+        InMemorySaver(),
+        reducer=_messages_delta_reducer,
+        write=lambda: [{"role": "tool", "content": "x"}],
+    )
+    config = {"configurable": {"thread_id": "t"}}
+
+    with pytest.raises(KeyError):
+        graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
+
+    assert [m.id for m in graph.get_state(config).values["messages"]] == ["a"]
+
+
+def _pregel_graph(seen: list[list[str | None]]) -> Any:
+    def look(state: dict) -> list:
+        seen.append([m.id for m in state["messages"]])
+        return []
+
+    node = (
+        NodeBuilder()
+        .subscribe_to("go", read=False)
+        .read_from("messages")
+        .do(look)
+        .write_to("messages")
+    )
+    return Pregel(
+        nodes={"look": node},
+        channels={
+            "messages": DeltaChannel(_messages_delta_reducer),
+            "go": LastValue(int),
+        },
+        input_channels=["messages", "go"],
+        output_channels=["messages"],
+        checkpointer=InMemorySaver(),
+    )
+
+
+def test_pregel_input_message_gets_stable_id_on_a_new_thread(
+    durability: Durability,
+) -> None:
+    graph = _pregel_graph([])
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke(
+        {"messages": [HumanMessage("hello")], "go": 1}, config, durability=durability
+    )
+
+    first, second = _hello_ids(graph, config)
+    assert first, "the message should get a UUID"
+    assert first == second, "its id must be the same on every read"
+
+
+def test_a_node_reads_the_pregel_input_message_ids_that_are_saved() -> None:
+    seen: list[list[str | None]] = []
+    graph = _pregel_graph(seen)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"messages": [HumanMessage("first", id="first")], "go": 1}, config)
+
+    graph.invoke({"messages": [{"role": "user", "content": "hello"}], "go": 1}, config)
+
+    saved = [m.id for m in graph.get_state(config).values["messages"]]
+    assert None not in saved
+    assert seen[-1] == saved, "the node must see the ids that are saved"
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        pytest.param({"type": "tool", "name": "search"}, id="not a message"),
+        pytest.param({"role": "user", "content": "x"}, id="message-shaped"),
+    ],
+)
+def test_other_delta_channels_keep_the_dicts_in_an_overwrite(value: dict) -> None:
+    graph = _build_graph(InMemorySaver(), write=lambda: Overwrite([dict(value)]))
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke({"messages": []}, config)
+
+    assert graph.get_state(config).values["messages"] == [value]

libs/langgraph/tests/test_delta_channel_node_cache.py

@@ -0,0 +1,122 @@
+"""A node served from the node cache saves its writes like a node that ran,
+without writing them back to the cache, and a node that fails isn't cached."""
+
+import operator
+from typing import Annotated, Any
+
+import pytest
+from langgraph.cache.memory import InMemoryCache
+from langgraph.checkpoint.memory import InMemorySaver
+from typing_extensions import TypedDict
+
+from langgraph.channels.delta import DeltaChannel
+from langgraph.graph import START, StateGraph
+from langgraph.types import CachePolicy, Durability
+
+pytestmark = pytest.mark.anyio
+
+INPUT = {"log": [], "plain": []}
+
+
+def _append(current: list, writes: list) -> list:
+    return [*current, *(item for write in writes for item in write)]
+
+
+class _State(TypedDict):
+    log: Annotated[list, DeltaChannel(_append)]
+    plain: Annotated[list, operator.add]
+
+
+class _CountsSets(InMemoryCache):
+    sets = 0
+
+    def set(self, keys: Any) -> None:
+        self.sets += 1
+        super().set(keys)
+
+
+def _a_then_cached_b_then_c(runs: list[str], cache: InMemoryCache) -> Any:
+    def node(name: str) -> Any:
+        def run(state: _State) -> dict:
+            runs.append(name)
+            return {"log": [name], "plain": [name]}
+
+        return run
+
+    builder = StateGraph(_State)
+    builder.add_node("a", node("a"))
+    builder.add_node("b", node("b"), cache_policy=CachePolicy())
+    builder.add_node("c", node("c"))
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    builder.add_edge("b", "c")
+    return builder.compile(checkpointer=InMemorySaver(), cache=cache)
+
+
+def test_a_cache_hit_saves_its_writes_without_caching_them_again(
+    durability: Durability,
+) -> None:
+    runs: list[str] = []
+    cache = _CountsSets()
+    graph = _a_then_cached_b_then_c(runs, cache)
+    graph.invoke(INPUT, {"configurable": {"thread_id": "1"}}, durability=durability)
+    config = {"configurable": {"thread_id": "2"}}
+
+    graph.invoke(INPUT, config, durability=durability)
+
+    assert runs == ["a", "b", "c", "a", "c"]
+    assert cache.sets == 1, "the cache hit was written back to the cache"
+    for state in graph.get_state_history(config):
+        assert state.values.get("log", []) == state.values.get("plain", [])
+
+
+async def test_a_cache_hit_saves_its_writes_without_caching_them_again_async(
+    durability: Durability,
+) -> None:
+    runs: list[str] = []
+    cache = _CountsSets()
+    graph = _a_then_cached_b_then_c(runs, cache)
+    await graph.ainvoke(
+        INPUT, {"configurable": {"thread_id": "1"}}, durability=durability
+    )
+    config = {"configurable": {"thread_id": "2"}}
+
+    await graph.ainvoke(INPUT, config, durability=durability)
+
+    assert runs == ["a", "b", "c", "a", "c"]
+    assert cache.sets == 1, "the cache hit was written back to the cache"
+    async for state in graph.aget_state_history(config):
+        assert state.values.get("log", []) == state.values.get("plain", [])
+
+
+def _cached_node_that_fails(runs: list[str]) -> Any:
+    def fail(state: _State) -> dict:
+        runs.append("b")
+        raise ValueError("b failed")
+
+    builder = StateGraph(_State)
+    builder.add_node("b", fail, cache_policy=CachePolicy())
+    builder.add_edge(START, "b")
+    return builder.compile(checkpointer=InMemorySaver(), cache=InMemoryCache())
+
+
+def test_a_node_that_fails_is_not_cached() -> None:
+    runs: list[str] = []
+    graph = _cached_node_that_fails(runs)
+
+    for thread in ("1", "2"):
+        with pytest.raises(ValueError, match="b failed"):
+            graph.invoke(INPUT, {"configurable": {"thread_id": thread}})
+
+    assert runs == ["b", "b"]
+
+
+async def test_a_node_that_fails_is_not_cached_async() -> None:
+    runs: list[str] = []
+    graph = _cached_node_that_fails(runs)
+
+    for thread in ("1", "2"):
+        with pytest.raises(ValueError, match="b failed"):
+            await graph.ainvoke(INPUT, {"configurable": {"thread_id": thread}})
+
+    assert runs == ["b", "b"]

libs/langgraph/tests/test_delta_channel_update_state.py

@@ -31,7 +31,7 @@ from langgraph.channels.last_value import LastValue
 from langgraph.graph import START, StateGraph
 from langgraph.graph.message import _messages_delta_reducer
 from langgraph.pregel import NodeBuilder, Pregel
-from langgraph.types import StateSnapshot, StateUpdate
+from langgraph.types import Overwrite, StateSnapshot, StateUpdate
 
 pytestmark = pytest.mark.anyio
 
@@ -661,24 +661,41 @@ def _messages_input_graph() -> Any:
     )
 
 
-def test_update_state_gives_a_message_an_id_that_every_read_keeps() -> None:
+# Writes of "b" in shapes `_messages_delta_reducer` accepts.
+MESSAGE_SHAPES = [
+    pytest.param(lambda: [HumanMessage("b")], id="message"),
+    pytest.param(lambda: ["b"], id="string"),
+    pytest.param(lambda: "b", id="bare string"),
+    pytest.param(lambda: [("user", "b")], id="tuple"),
+    pytest.param(lambda: {"role": "user", "content": "b"}, id="bare dict"),
+    pytest.param(lambda: Overwrite([HumanMessage("b")]), id="Overwrite"),
+]
+
+
+@pytest.mark.parametrize("write", MESSAGE_SHAPES)
+def test_update_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
+    write: Any,
+) -> None:
     graph = _build_graph(InMemorySaver())
     config = {"configurable": {"thread_id": "t"}}
     graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
 
-    graph.update_state(config, {"messages": [HumanMessage("b")]})
+    graph.update_state(config, {"messages": write()})
 
     first, second = _message_ids(graph, config), _message_ids(graph, config)
     assert first[-1] is not None
     assert first == second
 
 
-async def test_aupdate_state_gives_a_message_an_id_that_every_read_keeps() -> None:
+@pytest.mark.parametrize("write", MESSAGE_SHAPES)
+async def test_aupdate_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
+    write: Any,
+) -> None:
     graph = _build_graph(InMemorySaver())
     config = {"configurable": {"thread_id": "t"}}
     await graph.ainvoke({"messages": [HumanMessage("a", id="a")]}, config)
 
-    await graph.aupdate_state(config, {"messages": [HumanMessage("b")]})
+    await graph.aupdate_state(config, {"messages": write()})
 
     first, second = (
         await _amessage_ids(graph, config),
@@ -688,6 +705,19 @@ async def test_aupdate_state_gives_a_message_an_id_that_every_read_keeps() -> No
     assert first == second
 
 
+@pytest.mark.parametrize("write", MESSAGE_SHAPES)
+def test_update_as_input_gives_a_message_of_any_shape_an_id(write: Any) -> None:
+    graph = _messages_input_graph()
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
+
+    graph.update_state(config, {"messages": write()}, as_node="__input__")
+
+    first, second = _message_ids(graph, config), _message_ids(graph, config)
+    assert first[-1] is not None
+    assert first == second
+
+
 def test_update_state_on_an_older_checkpoint_gives_a_message_an_id() -> None:
     graph = _build_graph(InMemorySaver())
     config = {"configurable": {"thread_id": "t"}}
@@ -700,18 +730,6 @@ def test_update_state_on_an_older_checkpoint_gives_a_message_an_id() -> None:
     assert _message_ids(graph, branch)[-1] is not None
 
 
-def test_update_as_input_gives_a_message_an_id_that_every_read_keeps() -> None:
-    graph = _messages_input_graph()
-    config = {"configurable": {"thread_id": "t"}}
-    graph.invoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
-
-    graph.update_state(config, {"messages": [HumanMessage("b")]}, as_node="__input__")
-
-    first, second = _message_ids(graph, config), _message_ids(graph, config)
-    assert first[-1] is not None
-    assert first == second
-
-
 async def test_aupdate_as_input_gives_a_message_an_id_that_every_read_keeps() -> None:
     graph = _messages_input_graph()
     config = {"configurable": {"thread_id": "t"}}

libs/langgraph/tests/test_pregel.py

@@ -5932,9 +5932,9 @@ def test_no_redundant_put_writes_for_cached_task(
     put_writes_task_ids: list[str] = []
     orig = PregelLoop.put_writes
 
-    def spy(self, task_id, writes):
+    def spy(self, task_id, writes, **kwargs):
         put_writes_task_ids.append(task_id)
-        return orig(self, task_id, writes)
+        return orig(self, task_id, writes, **kwargs)
 
     with patch.object(PregelLoop, "put_writes", spy):
         result = workflow.invoke(Command(resume="ans"), config=config)

libs/langgraph/tests/test_pregel_async.py

@@ -8157,9 +8157,9 @@ async def test_no_redundant_put_writes_for_cached_task(
     put_writes_task_ids: list[str] = []
     orig = PregelLoop.put_writes
 
-    def spy(self, task_id, writes):
+    def spy(self, task_id, writes, **kwargs):
         put_writes_task_ids.append(task_id)
-        return orig(self, task_id, writes)
+        return orig(self, task_id, writes, **kwargs)
 
     with patch.object(PregelLoop, "put_writes", spy):
         result = await workflow.ainvoke(Command(resume="ans"), config=config)

libs/langgraph/tests/test_runtime.py

@@ -248,6 +248,95 @@ def test_drain_from_subgraph_can_resume_parent() -> None:
     }
 
 
+class _BreakpointState(TypedDict, total=False):
+    first: str
+    second: str
+
+
+def _drain_before_breakpoint_builder(
+    second_runs: list[str],
+) -> StateGraph[_BreakpointState]:
+    def first(state: _BreakpointState, runtime: Runtime) -> _BreakpointState:
+        runtime.control.request_drain("rollout")
+        return {"first": "done"}
+
+    def second(state: _BreakpointState) -> _BreakpointState:
+        second_runs.append("second")
+        return {"second": "done"}
+
+    builder = StateGraph(_BreakpointState)
+    builder.add_node("first", first)
+    builder.add_node("second", second)
+    builder.add_edge(START, "first")
+    builder.add_edge("first", "second")
+    return builder
+
+
+@pytest.mark.parametrize("at_compile", [True, False])
+def test_drain_before_interrupt_before_stops_at_breakpoint(at_compile: bool) -> None:
+    # A resume passes the next step's breakpoints, so a drain right before a
+    # breakpoint has to stop there instead, or the resumed run skips it.
+    second_runs: list[str] = []
+    breakpoints = {"interrupt_before": ["second"]}
+    compiled = _drain_before_breakpoint_builder(second_runs).compile(
+        checkpointer=MemorySaver(), **(breakpoints if at_compile else {})
+    )
+    run_kwargs = {} if at_compile else breakpoints
+    config = {"configurable": {"thread_id": "drain-breakpoint"}}
+
+    assert compiled.invoke({}, config, control=RunControl(), **run_kwargs) == {
+        "first": "done"
+    }
+    assert compiled.get_state(config).next == ("second",)
+    assert second_runs == []
+
+    assert compiled.invoke(None, config, **run_kwargs) == {
+        "first": "done",
+        "second": "done",
+    }
+    assert second_runs == ["second"]
+
+
+def test_drain_before_subgraph_interrupt_before_stops_at_breakpoint() -> None:
+    second_runs: list[str] = []
+    child = _drain_before_breakpoint_builder(second_runs).compile(
+        interrupt_before=["second"]
+    )
+    parent = StateGraph(_BreakpointState)
+    parent.add_node("child", child)
+    parent.add_edge(START, "child")
+    compiled = parent.compile(checkpointer=MemorySaver())
+    config = {"configurable": {"thread_id": "drain-subgraph-breakpoint"}}
+
+    compiled.invoke({}, config, control=RunControl())
+    state = compiled.get_state(config, subgraphs=True)
+    assert state.next == ("child",)
+    assert state.tasks[0].state.next == ("second",)
+    assert second_runs == []
+
+    assert compiled.invoke(None, config) == {"first": "done", "second": "done"}
+    assert second_runs == ["second"]
+
+
+@pytest.mark.anyio
+async def test_drain_before_interrupt_before_stops_at_breakpoint_async() -> None:
+    second_runs: list[str] = []
+    compiled = _drain_before_breakpoint_builder(second_runs).compile(
+        checkpointer=MemorySaver(), interrupt_before=["second"]
+    )
+    config = {"configurable": {"thread_id": "drain-breakpoint-async"}}
+
+    assert await compiled.ainvoke({}, config, control=RunControl()) == {"first": "done"}
+    assert (await compiled.aget_state(config)).next == ("second",)
+    assert second_runs == []
+
+    assert await compiled.ainvoke(None, config) == {
+        "first": "done",
+        "second": "done",
+    }
+    assert second_runs == ["second"]
+
+
 @pytest.mark.anyio
 async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None:
     class State(TypedDict, total=False):

libs/sdk-py/langgraph_sdk/__init__.py

@@ -3,7 +3,7 @@ from langgraph_sdk.client import get_client, get_sync_client
 from langgraph_sdk.encryption import Encryption
 from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
 
-__version__ = "0.4.6"
+__version__ = "0.4.7"
 
 __all__ = [
     "Auth",