fix(langgraph): replay a resumed exit-mode run's delta writes in live order (#9114)

Depends on #8544

Resuming an exit-durability run after a parallel interrupt replays the
resumed checkpoint's `DeltaChannel` writes wrong: twice on `main`, out
of order if the duplicate is simply dropped. `p` finished before `q`
interrupted, and `r` runs after `q`:

```
live:            ['in1', 'p', 'r']
main, reload:    ['in1', 'p', 'r', 'p']
```

No fork is involved. Any exit-mode run that starts from a checkpoint
with pending delta writes hits it: resuming an interrupt, or retrying
after a task error.

## Cause

Exit durability stores a run's delta writes on the checkpoint it started
from, under step-prefixed task ids (`exit_delta_task_id`), so they
replay in superstep order. On a resume, that checkpoint already holds
the finished tasks' writes under their real task ids. The accumulator
stored those loaded writes again under a step-prefixed id, so they
replayed twice. Skipping them alone reorders the replay, because the
real ids sort after the step-prefixed ones.

## Fix

At exit:

- writes loaded with the checkpoint are skipped, since they are already
stored on it: `NULL_TASK_ID` writes by task id (this run's own `Command`
writes are recorded separately in `_first`), the rest by `(task_id,
channel)`;
- the checkpoint's own superstep keeps its real task paths, so it
interleaves with the loaded writes as live execution did, but gets a
step-prefixed task id. Under the real id, a resume whose final
checkpoint fails to save would leave the resumed task looking done, and
the retry would skip it and lose its other writes;
- a resume addressed by `checkpoint_id` reruns tasks whose writes are
already stored on the checkpoint; their writes to those channels are
skipped, and a write to a channel the task did not write before is kept;
- later supersteps get task id `ffffffff-<step>-...` and task path
`~~<step>...`, which sort after every real task id and every real task
path, in step order.

So superstep order holds whether a saver orders by task path (#8544) or
only by task id.

## Limits

- Within the resumed superstep, a saver that orders only by task id
replays this run's writes before the ones loaded from the checkpoint.
#8544 fixes that for the OSS savers.
- If the final checkpoint fails to save, the writes of supersteps after
the resumed one stay on the checkpoint and replay once more after the
retry (on `main`, every write does). That comes from exit durability
storing its writes before its checkpoint, not from this change.

## Tests

`test_delta_channel_exit_mode.py`, on the shared checkpointer fixture
(memory, sqlite, postgres) in all three durabilities:

- resume after a parallel interrupt with a later write, plain and with
the head's `checkpoint_id`
- the resumed superstep interleaved by task path: `a_asks` resumes after
`z_done` finished, and live order puts `a` first
- the same runs on a saver that orders writes by task id only
- a resume retried after its final checkpoint failed to save
- a resume with `Command(resume=..., update=...)` that writes the delta
channel replays that write once, in live order
- an addressed resume whose rerun task writes a delta channel it did not
write before keeps that write

The exit cases fail without the change, except the last one, which
guards against skipping by task id alone (an earlier version of this PR
did); sync and async pass either way.

#8548 marked exit durability in
`test_resume_on_an_interrupted_head_consumes_its_writes_without_a_snapshot`
as a strict expected failure for this bug; this PR drops the mark.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: ErenAta16 <149434812+ErenAta16@users.noreply.github.com>
Co-authored-by: ragnarok268 <58264829+ragnarok268@users.noreply.github.com>

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -168,8 +168,8 @@ class DeltaChannelHistory(TypedDict):
       a super-step's task writes in. `task_id` is a hash of the path, so
       ordering by it permutes parallel tasks writing one channel, and
       reducers need not be order-invariant. Writes stored without a
-      `task_path` (graph input, `update_state` updates, exit-durability runs,
-      rows predating the column) sort first, by `task_id`.
+      `task_path` (graph input, `update_state` and `Command` updates, rows
+      predating the column) sort first, by `task_id`.
     * `seed` — the stored value at the nearest ancestor whose
       `channel_values[ch]` is populated. Omitted if the walk reached the
       root without finding any stored value (consumer treats absence as

libs/langgraph/langgraph/pregel/_checkpoint.py

@@ -60,6 +60,16 @@ def exit_delta_task_id(step: int, task_id: str) -> str:
     return f"{step:08d}-{parts[1]}-{parts[2]}-{parts[3]}-{parts[4]}"
 
 
+def exit_delta_late_task_id(step: int, task_id: str) -> str:
+    """Synthetic task id for exit-mode writes of a superstep after the anchor's own.
+
+    Sorts after every real task id, in step order, so replay keeps them after
+    the anchor's own superstep whether a saver orders by task path or task id.
+    """
+    parts = str(uuid.UUID(task_id)).split("-")
+    return f"ffffffff-{step >> 16:04x}-{step & 0xFFFF:04x}-{parts[3]}-{parts[4]}"
+
+
 def delta_channels_to_snapshot(
     channels: Mapping[str, BaseChannel],
     counters_since_delta_snapshot: Mapping[str, tuple[int, int]],

libs/langgraph/langgraph/pregel/_loop.py

@@ -106,6 +106,7 @@ from langgraph.pregel._checkpoint import (
     delta_channels_to_snapshot,
     delta_channels_with_pending_writes,
     empty_checkpoint,
+    exit_delta_late_task_id,
     exit_delta_task_id,
 )
 from langgraph.pregel._executor import (
@@ -221,10 +222,18 @@ class PregelLoop:
     # `after_tick`). At exit, `_put_exit_delta_writes` filters out channels
     # that will snapshot, then persists the rest under an anchor parent.
     # `None` when not in exit mode (so the capture sites are no-ops).
-    # Each tuple is `(step, task_id, channel, value)` — `step` drives the
-    # synthetic step-prefixed task_id used to preserve chronological order
-    # under the saver's `ORDER BY task_id, idx` sorting.
-    _exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
+    # Each tuple is `(step, task_id, task_path, channel, value)`; see
+    # `_put_exit_delta_writes` for how they are ordered.
+    _exit_delta_writes: list[tuple[int, str, str, str, Any]] | None = None
+
+    # The (task_id, channel) pairs whose delta writes are stored on the loaded
+    # checkpoint, which a resume addressed by `checkpoint_id` can rerun; this
+    # run's `Command` delta writes, kept apart from the NULL_TASK_ID writes
+    # loaded with the checkpoint; and the checkpoint's own superstep, the
+    # first one this run ticks.
+    _stored_delta_writes: set[tuple[str, str]]
+    _exit_command_writes: list[tuple[str, Any]]
+    _exit_first_step: int | None = None
 
     # Delta channels that must snapshot at the next checkpoint, whatever their
     # cadence counters say:
@@ -737,9 +746,25 @@ class PregelLoop:
             )
         # capture delta-channel writes for exit-mode accumulator before clearing
         if self._exit_delta_writes is not None:
+            # On the first tick the pending writes still hold the ones loaded
+            # with the checkpoint, which are already stored on it.
+            first = self._exit_first_step is None
+            if first:
+                self._exit_first_step = self.step
+                self._exit_delta_writes.extend(
+                    (self.step, NULL_TASK_ID, "", ch, v)
+                    for ch, v in self._exit_command_writes
+                )
             for tid, ch, v in self.checkpoint_pending_writes:
-                if isinstance(self.specs.get(ch), DeltaChannel):
-                    self._exit_delta_writes.append((self.step, tid, ch, v))
+                if not isinstance(self.specs.get(ch), DeltaChannel):
+                    continue
+                if first and (
+                    tid == NULL_TASK_ID or (tid, ch) in self._stored_delta_writes
+                ):
+                    continue
+                task = self.tasks.get(tid)
+                path = task_path_str(task.path) if task else ""
+                self._exit_delta_writes.append((self.step, tid, path, ch, v))
         # clear pending writes
         self.checkpoint_pending_writes.clear()
         # only replay (re-execute) done tasks on the first tick
@@ -860,6 +885,12 @@ class PregelLoop:
     def _first(
         self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
     ) -> set[str] | None:
+        self._stored_delta_writes = {
+            (tid, ch)
+            for tid, ch, _ in self.checkpoint_pending_writes
+            if tid != NULL_TASK_ID and isinstance(self.specs.get(ch), DeltaChannel)
+        }
+        self._exit_command_writes = []
         # Resuming from a previous checkpoint requires two things:
         # 1. A prior checkpoint exists (channel_versions is non-empty)
         # 2. The input signals continuation (not a fresh run with new input)
@@ -977,6 +1008,12 @@ class PregelLoop:
                     carried.extend((tid, c, v) for c, v in ws)
                 else:
                     self.put_writes(tid, ws)
+                    if self._exit_delta_writes is not None and tid == NULL_TASK_ID:
+                        self._exit_command_writes.extend(
+                            (c, v)
+                            for c, v in ws
+                            if isinstance(self.specs.get(c), DeltaChannel)
+                        )
             self._delta_channels_forced_snapshot.update(
                 delta_channels_with_pending_writes(self.specs, carried)
             )
@@ -1074,7 +1111,9 @@ class PregelLoop:
             if self._exit_delta_writes is not None:
                 for c, v in input_writes:
                     if isinstance(self.specs.get(c), DeltaChannel):
-                        self._exit_delta_writes.append((self.step, NULL_TASK_ID, c, v))
+                        self._exit_delta_writes.append(
+                            (self.step, NULL_TASK_ID, "", c, v)
+                        )
             # Persist delta-channel input writes so sub-freq inputs are
             # recoverable via ancestor walk (mirrors the Command input path).
             if self.durability != "exit":
@@ -1309,9 +1348,7 @@ class PregelLoop:
         )
 
         pending = [
-            (step, tid, ch, v)
-            for (step, tid, ch, v) in self._exit_delta_writes
-            if ch not in channels_to_snapshot
+            w for w in self._exit_delta_writes if w[3] not in channels_to_snapshot
         ]
         if not pending:
             return
@@ -1346,11 +1383,21 @@ class PregelLoop:
             # sees the stub as its parent.
             self.checkpoint_config = anchor_config
 
-        # Step-prefixed synthetic task_id preserves chronological superstep
-        # order under the saver's ORDER BY task_id, idx sorting.
-        grouped: dict[tuple[int, str], list[tuple[str, Any]]] = {}
-        for step, tid, ch, v in pending:
-            grouped.setdefault((step, tid), []).append((ch, v))
+        # The checkpoint's own superstep keeps its real task paths, so it
+        # interleaves with the writes a resume loaded from it. Its task ids stay
+        # synthetic: under the real id, a run whose final checkpoint fails to
+        # save would leave the resumed task looking done to the next resume.
+        # Later supersteps sort after every real task path and task id, in step
+        # order, so this holds whether a saver orders by path or by id.
+        grouped: dict[tuple[str, str], list[tuple[str, Any]]] = {}
+        for step, tid, path, ch, v in pending:
+            if tid == NULL_TASK_ID:
+                key = (exit_delta_task_id(step, tid), "")
+            elif step == self._exit_first_step:
+                key = (exit_delta_task_id(step, tid), path)
+            else:
+                key = (exit_delta_late_task_id(step, tid), f"~~{step:010d}{path}")
+            grouped.setdefault(key, []).append((ch, v))
         anchor_write_config = patch_configurable(
             anchor_config,
             {
@@ -1360,22 +1407,21 @@ class PregelLoop:
                 CONFIG_KEY_CHECKPOINT_ID: anchor_config[CONF][CONFIG_KEY_CHECKPOINT_ID],
             },
         )
-        for (step, tid), entries in grouped.items():
-            synth_tid = exit_delta_task_id(step, tid)
+        for (tid, path), entries in grouped.items():
             if self.checkpointer_put_writes_accepts_task_path:
                 fut = self.submit(
                     self.checkpointer_put_writes,
                     anchor_write_config,
                     entries,
-                    synth_tid,
-                    "",
+                    tid,
+                    path,
                 )
             else:
                 fut = self.submit(
                     self.checkpointer_put_writes,
                     anchor_write_config,
                     entries,
-                    synth_tid,
+                    tid,
                 )
             if self._delta_write_futs is not None:
                 self._delta_write_futs.append(fut)

libs/langgraph/tests/test_delta_channel_exit_mode.py

@@ -6,11 +6,13 @@ channel), lazy stub creation when no parent exists, and proper read-path
 reconstruction via ancestor walks.
 """
 
+import operator
 import uuid
 from typing import Annotated, Any
 
 import pytest
 from langchain_core.messages import AIMessage, HumanMessage
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.checkpoint.serde.types import _DeltaSnapshot
 from typing_extensions import TypedDict
@@ -19,6 +21,7 @@ from langgraph.channels.delta import DeltaChannel
 from langgraph.graph import START, StateGraph
 from langgraph.graph.message import _messages_delta_reducer
 from langgraph.pregel._checkpoint import exit_delta_task_id
+from langgraph.types import Command, Durability, interrupt
 
 pytestmark = pytest.mark.anyio
 
@@ -389,3 +392,197 @@ async def test_exit_snapshot_then_tail_deltas() -> None:
     assert "seed-msg" in contents
     assert "tail-msg" in contents
     assert contents.index("seed-msg") < contents.index("tail-msg")
+
+
+def _append(current: list, writes: list) -> list:
+    out = list(current)
+    for write in writes:
+        out.extend(write)
+    return out
+
+
+class _ResumeState(TypedDict):
+    log: Annotated[list, DeltaChannel(_append)]
+    plain: Annotated[list, operator.add]
+
+
+def _both(marker: str) -> dict:
+    return {"log": [marker], "plain": [marker]}
+
+
+def _ask(marker: str) -> Any:
+    def ask(state: _ResumeState) -> dict:
+        interrupt("approve?")
+        return _both(marker)
+
+    return ask
+
+
+@pytest.mark.parametrize("addressed", [False, True])
+def test_resume_after_a_parallel_interrupt_replays_in_live_order(
+    sync_checkpointer: BaseCheckpointSaver, durability: Durability, addressed: bool
+) -> None:
+    builder = StateGraph(_ResumeState)
+    builder.add_node("done", lambda state: _both("done"))
+    builder.add_node("ask", _ask("ask"))
+    builder.add_node("after", lambda state: _both("after"))
+    builder.add_edge(START, "done")
+    builder.add_edge(START, "ask")
+    builder.add_edge("ask", "after")
+    graph = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability=durability)
+    head = graph.get_state(config).config
+
+    graph.invoke(
+        Command(resume="yes"), head if addressed else config, durability=durability
+    )
+
+    state = graph.get_state(config)
+    assert state.values["log"] == state.values["plain"]
+    assert sorted(state.values["log"]) == ["after", "ask", "done", "in"]
+
+
+def test_resume_with_a_command_update_replays_its_write_once(
+    sync_checkpointer: BaseCheckpointSaver, durability: Durability
+) -> None:
+    builder = StateGraph(_ResumeState)
+    builder.add_node("done", lambda state: _both("done"))
+    builder.add_node("ask", _ask("ask"))
+    builder.add_edge(START, "done")
+    builder.add_edge(START, "ask")
+    graph = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability=durability)
+
+    graph.invoke(
+        Command(resume="yes", update=_both("cmd")), config, durability=durability
+    )
+
+    state = graph.get_state(config)
+    assert state.values["log"] == state.values["plain"] == ["in", "cmd", "ask", "done"]
+
+
+class _FlagState(_ResumeState, total=False):
+    extra: Annotated[list, DeltaChannel(_append)]
+    flag: bool
+
+
+def test_addressed_resume_keeps_a_rerun_tasks_write_to_a_new_channel(
+    sync_checkpointer: BaseCheckpointSaver, durability: Durability
+) -> None:
+    def done(state: _FlagState) -> dict:
+        return {**_both("done"), **({"extra": ["new"]} if state.get("flag") else {})}
+
+    builder = StateGraph(_FlagState)
+    builder.add_node("done", done)
+    builder.add_node("ask", _ask("ask"))
+    builder.add_edge(START, "done")
+    builder.add_edge(START, "ask")
+    graph = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability=durability)
+    head = graph.get_state(config).config
+
+    live = graph.invoke(
+        Command(resume="yes", update={"flag": True}), head, durability=durability
+    )
+
+    state = graph.get_state(config)
+    assert live["extra"] == state.values["extra"] == ["new"]
+    assert state.values["log"] == state.values["plain"]
+
+
+def test_resume_interleaves_the_resumed_superstep_by_task_path(
+    sync_checkpointer: BaseCheckpointSaver, durability: Durability
+) -> None:
+    builder = StateGraph(_ResumeState)
+    builder.add_node("z_done", lambda state: _both("z"))
+    builder.add_node("a_asks", _ask("a"))
+    builder.add_edge(START, "z_done")
+    builder.add_edge(START, "a_asks")
+    graph = builder.compile(checkpointer=sync_checkpointer)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability=durability)
+
+    graph.invoke(Command(resume="yes"), config, durability=durability)
+
+    state = graph.get_state(config)
+    assert state.values["log"] == state.values["plain"] == ["in", "a", "z"]
+
+
+class _TaskIdOrderSaver(InMemorySaver):
+    """Replays each checkpoint's writes by task id, as savers without task path
+    ordering do."""
+
+    def get_tuple(self, config: Any) -> Any:
+        tup = super().get_tuple(config)
+        if tup and tup.pending_writes:
+            tup = tup._replace(pending_writes=sorted(tup.pending_writes))
+        return tup
+
+    get_delta_channel_history = BaseCheckpointSaver.get_delta_channel_history
+
+
+def test_exit_run_replays_supersteps_in_order_on_a_task_id_ordered_saver() -> None:
+    builder = StateGraph(_ResumeState)
+    builder.add_node("a", lambda state: _both("a"))
+    builder.add_node("b", lambda state: _both("b"))
+    builder.add_edge(START, "a")
+    builder.add_edge("a", "b")
+    graph = builder.compile(checkpointer=_TaskIdOrderSaver())
+    config = {"configurable": {"thread_id": "t"}}
+
+    graph.invoke(_both("in"), config, durability="exit")
+
+    assert graph.get_state(config).values["log"] == ["in", "a", "b"]
+
+
+def test_exit_resume_replays_supersteps_in_order_on_a_task_id_ordered_saver() -> None:
+    builder = StateGraph(_ResumeState)
+    builder.add_node("ask", _ask("ask"))
+    builder.add_node("after", lambda state: _both("after"))
+    builder.add_edge(START, "ask")
+    builder.add_edge("ask", "after")
+    graph = builder.compile(checkpointer=_TaskIdOrderSaver())
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability="exit")
+
+    graph.invoke(Command(resume="yes"), config, durability="exit")
+
+    assert graph.get_state(config).values["log"] == ["in", "ask", "after"]
+
+
+class _FailingPutSaver(InMemorySaver):
+    fail = False
+
+    def put(
+        self, config: Any, checkpoint: Any, metadata: Any, new_versions: Any
+    ) -> Any:
+        if self.fail:
+            raise RuntimeError("final checkpoint lost")
+        return super().put(config, checkpoint, metadata, new_versions)
+
+
+def test_exit_resume_retried_after_its_final_checkpoint_fails_reruns_the_resumed_task() -> (
+    None
+):
+    saver = _FailingPutSaver()
+    builder = StateGraph(_ResumeState)
+    builder.add_node("done", lambda state: _both("done"))
+    builder.add_node("ask", _ask("ask"))
+    builder.add_edge(START, "done")
+    builder.add_edge(START, "ask")
+    graph = builder.compile(checkpointer=saver)
+    config = {"configurable": {"thread_id": "t"}}
+    graph.invoke(_both("in"), config, durability="exit")
+    saver.fail = True
+    with pytest.raises(RuntimeError, match="final checkpoint lost"):
+        graph.invoke(Command(resume="yes"), config, durability="exit")
+    saver.fail = False
+
+    graph.invoke(Command(resume="yes"), config, durability="exit")
+
+    state = graph.get_state(config)
+    assert state.values["log"] == state.values["plain"]
+    assert sorted(state.values["plain"]) == ["ask", "done", "in"]

libs/langgraph/tests/test_delta_channel_fork.py

@@ -740,20 +740,6 @@ def _build_deferred_after_interrupt(checkpointer: BaseCheckpointSaver) -> Any:
     return builder.compile(checkpointer=checkpointer, interrupt_after=["a"])
 
 
-@pytest.mark.parametrize(
-    "durability",
-    [
-        "sync",
-        "async",
-        pytest.param(
-            "exit",
-            marks=pytest.mark.xfail(
-                reason="exit durability stores a resumed run's loaded writes twice",
-                strict=True,
-            ),
-        ),
-    ],
-)
 def test_resume_on_an_interrupted_head_consumes_its_writes_without_a_snapshot(
     sync_checkpointer: BaseCheckpointSaver, durability: Durability
 ) -> None: