fix(runtime): drain sync checkpoint mutations across cancellation (#5664)

* fix(runtime): drain sync goal checkpoint commits across cancellation

* fix(runtime): drain sync rollback checkpointer mutations

* test(runtime): cover sync rollback mutation cancellation

* docs(runtime): generalize sync checkpoint mutation contract

backend/packages/harness/deerflow/runtime/AGENTS.md

@@ -14,7 +14,7 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
 
 **Compatibility is asymmetric and fail-closed.** Every checkpoint written in delta mode carries metadata marker `deerflow_checkpoint_channel_mode: "delta"` (injected via `inject_checkpoint_mode`; absence of marker = full, so pre-feature checkpoints need no migration). Before any state read/write, `ensure_checkpoint_mode_compatible` rejects a full-mode process opening a delta thread with `CheckpointModeMismatchError` (surfaced as HTTP 409 with the cause and thread id by the threads router; `CheckpointModeReconfigurationError` maps to 503) — a full-mode raw read of a delta blob would silently return empty/partial `messages`. The reverse direction is allowed: delta-mode processes read full checkpoints transparently (old full checkpoints seed the delta channel), so full → delta is the smooth migration path; delta → full requires materializing/converting the data first. Detection also honors upstream's `counters_since_delta_snapshot.messages` metadata, and an explicit config marker takes precedence over any ambient context value.
 
-**Never bypass `CheckpointStateAccessor` (`runtime/checkpoint_state.py`) for thread-state access.** It is the single choke point binding graph + checkpointer + mode: it injects the mode marker into configs, runs the compatibility check before every `get`/`update`/`history`, and returns materialized state (delta checkpoints lack `channel_values.messages` — raw `get_tuple` reads see a sentinel). Use `get_metadata` / `aget_metadata` when only persisted metadata is needed, avoiding delta-history materialization while retaining the mode gate. Gateway `services.py` builds and passes the accessor; thread-owned reads (state/history/regeneration) must use `build_thread_checkpoint_state_accessor` so the recorded assistant's middleware schema materializes every channel. `history(limit)` semantics: `0` means zero items (explicit empty), `None` means unlimited — do not pass `limit=0` through to `graph.get_state_history`. Assistant metadata lookup is fail-closed for mutation accessors so a store outage cannot silently select the default schema and discard extension channels. In `full` mode the read path degrades to a raw checkpointer read (`_RawCheckpointReadAccessor`) when the agent factory cannot build the graph (bad model config, MCP outage) — full checkpoints carry complete `channel_values`, so reads don't need the graph; degraded snapshots take `created_at` from the standard checkpoint `ts` field, falling back to metadata only for compatibility. The delta gate still applies on the degraded path; `next`/`tasks` degrade to empty and thread status falls back to the stored status because task presence is not derivable, while delta mode has no fallback (materialization needs the channel table).
+**Never bypass `CheckpointStateAccessor` (`runtime/checkpoint_state.py`) for thread-state access.** It is the single choke point binding graph + checkpointer + mode: it injects the mode marker into configs, runs the compatibility check before every `get`/`update`/`history`, and returns materialized state (delta checkpoints lack `channel_values.messages` — raw `get_tuple` reads see a sentinel). For metadata-only reads, use `get_metadata` / `aget_metadata` to retain the mode gate without materialization. Thread-owned reads must use `build_thread_checkpoint_state_accessor` so the recorded assistant schema materializes every channel. Sync checkpoint mutations must drain off-thread commits before cancellation propagates; sync `get_tuple` reads remain cancellable. `history(limit)` semantics: `0` means zero items (explicit empty), `None` means unlimited — do not pass `limit=0` through to `graph.get_state_history`. Assistant metadata lookup is fail-closed for mutation accessors so a store outage cannot silently select the default schema and discard extension channels. In `full` mode the read path degrades to a raw checkpointer read (`_RawCheckpointReadAccessor`) when the agent factory cannot build the graph (bad model config, MCP outage) — full checkpoints carry complete `channel_values`, so reads don't need the graph; degraded snapshots take `created_at` from the standard checkpoint `ts` field, falling back to metadata only for compatibility. The delta gate still applies on the degraded path; `next`/`tasks` degrade to empty and thread status falls back to the stored status because task presence is not derivable, while delta mode has no fallback (materialization needs the channel table).
 
 **Replay checkpoint lookup prefers lineage and degrades only for an explicitly missing legacy parent link.** Branch and regenerate paths first walk `parent_config`, which prevents a global chronological scan from selecting a sibling created by regeneration. `CheckpointParentMissingError` alone enables the bounded newest-first history fallback in `app/gateway/checkpoint_lineage.py`; cycles, dangling/non-addressable parents, target mismatches, and depth exhaustion raise `CheckpointLineageIntegrityError` and fail closed instead of selecting a sibling. The compatibility scans request 400 raw checkpoints so up to 200 duration-only entries do not consume the effective branch-history budget; the fallback scans oldest-to-newest internally, skips duration-only checkpoints, and accepts only checkpoints with an addressable id as the replay base. A source history with no discoverable pre-user checkpoint preserves the historical single-checkpoint branch behavior instead of rejecting the branch; regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are not mutated by regenerate preparation, and no raw checkpoint tuple is copied across threads because delta state depends on ancestry and pending writes. Regenerate source-run lookup uses the current thread's exact event, then the server-stamped `run_id` on the copied human message, then verified RunManager content matching; it does not read parent-thread events. When an interrupted response was streamed but never checkpointed, regeneration accepts only the latest visible human message's server-stamped `run_id` after verifying that it belongs to the same thread and still has `interrupted` status. Storage or checkpoint-mode failures are not treated as a missing base and still fail closed.
 
@@ -361,4 +361,4 @@ Drain tasks are named `jsonl-mutation:{thread_id}` for asyncio task dumps. Multi
 the admitted group keeps its records on success or completes rollback on failure.
 This is a store-local guarantee, not a change to RunJournal cancellation policy or
 JSONL's single-process deployment constraint. Regression coverage is in
-`tests/test_jsonl_event_store_cancellation.py`.
+`tests/test_jsonl_event_store_cancellation.py`.

backend/packages/harness/deerflow/runtime/goal.py

@@ -25,6 +25,7 @@ from deerflow.agents.goal_state import GoalBlocker, GoalEvaluation, GoalState
 from deerflow.models import create_chat_model
 from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
 from deerflow.tracing import inject_langfuse_metadata
+from deerflow.utils.file_io import await_drained
 from deerflow.utils.messages import message_to_text
 from deerflow.utils.time import now_iso
 
@@ -419,8 +420,11 @@ async def _call_checkpointer_method(checkpointer: Any, async_name: str, sync_nam
     if sync_method is None:
         raise AttributeError(f"Missing checkpointer method: {async_name}/{sync_name}")
     # Offload the synchronous checkpointer call so its blocking IO never runs on
-    # the event loop (backend/AGENTS.md blocking-IO gate).
-    result = await asyncio.to_thread(sync_method, *args, **kwargs)
+    # the event loop (backend/AGENTS.md blocking-IO gate). A sync checkpoint
+    # mutation must finish before cancellation propagates; otherwise the caller
+    # can observe cancellation while the worker commits state afterwards.
+    worker = asyncio.to_thread(sync_method, *args, **kwargs)
+    result = await await_drained(worker) if sync_name in {"put", "put_writes", "delete_thread"} else await worker
     return await result if inspect.isawaitable(result) else result
 
 

backend/tests/test_goal_sync_checkpointer_cancellation.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from deerflow.runtime.goal import build_goal_state, write_thread_goal
+
+
+class _BlockingSyncCheckpointer:
+    def __init__(self) -> None:
+        self.put_started = threading.Event()
+        self.allow_put = threading.Event()
+        self.put_finished = threading.Event()
+        self.saved_checkpoint = None
+
+    def get_tuple(self, _config):
+        return SimpleNamespace(
+            config={"configurable": {"checkpoint_id": "checkpoint-1"}},
+            checkpoint={
+                "id": "checkpoint-1",
+                "channel_values": {},
+                "channel_versions": {},
+            },
+            metadata={"step": 0},
+        )
+
+    def put(self, _config, checkpoint, _metadata, _new_versions):
+        self.put_started.set()
+        try:
+            assert self.allow_put.wait(5.0)
+            self.saved_checkpoint = checkpoint
+        finally:
+            self.put_finished.set()
+
+
+@pytest.mark.asyncio
+async def test_sync_goal_checkpoint_write_drains_across_repeated_cancellation() -> None:
+    checkpointer = _BlockingSyncCheckpointer()
+    task = asyncio.create_task(
+        write_thread_goal(
+            checkpointer,
+            "thread-1",
+            build_goal_state("Finish the migration"),
+        )
+    )
+
+    try:
+        assert await asyncio.to_thread(checkpointer.put_started.wait, 1.0)
+
+        task.cancel()
+        await asyncio.sleep(0)
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+
+        assert not task.done(), "goal write returned before the synchronous checkpoint commit finished"
+
+        checkpointer.allow_put.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert checkpointer.put_finished.is_set()
+        assert checkpointer.saved_checkpoint is not None
+    finally:
+        checkpointer.allow_put.set()
+        await asyncio.gather(task, return_exceptions=True)
+        await asyncio.to_thread(checkpointer.put_finished.wait, 1.0)

backend/tests/test_run_worker_sync_checkpointer_cancellation.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+
+from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
+from deerflow.runtime.runs.worker import RollbackPoint, _rollback_to_pre_run_checkpoint
+
+
+class _BlockingSyncRollbackCheckpointer:
+    def __init__(self) -> None:
+        self.mutation_started = threading.Event()
+        self.allow_mutation = threading.Event()
+        self.mutation_finished = threading.Event()
+
+    async def aget_tuple(self, _config: dict[str, Any]) -> None:
+        return None
+
+    def _block_mutation(self) -> None:
+        self.mutation_started.set()
+        try:
+            assert self.allow_mutation.wait(5.0)
+        finally:
+            self.mutation_finished.set()
+
+    def delete_thread(self, _thread_id: str) -> None:
+        self._block_mutation()
+
+    def put_writes(self, _config: dict[str, Any], _writes: list[tuple[str, Any]], *, task_id: str) -> None:
+        del task_id
+        self._block_mutation()
+
+
+def _rollback_point() -> RollbackPoint:
+    return RollbackPoint(
+        config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "ckpt-1"}},
+        state_values={},
+        messages=("before",),
+        metadata={"source": "input"},
+        pending_writes=(("task-a", "messages", "value"),),
+    )
+
+
+@pytest.mark.parametrize("mutation", ["delete_thread", "put_writes"])
+@pytest.mark.asyncio
+async def test_sync_rollback_mutations_drain_across_repeated_cancellation(monkeypatch, mutation: str) -> None:
+    checkpointer = _BlockingSyncRollbackCheckpointer()
+    rollback_point = None
+    if mutation == "put_writes":
+        graph = SimpleNamespace(aupdate_state=AsyncMock(return_value={"configurable": {"thread_id": "thread-1", "checkpoint_ns": "", "checkpoint_id": "restored-1"}}))
+        monkeypatch.setattr("deerflow.runtime.runs.worker.build_state_mutation_graph", lambda *_args, **_kwargs: graph)
+        rollback_point = _rollback_point()
+
+    accessor = CheckpointStateAccessor(graph=SimpleNamespace(), checkpointer=checkpointer, mode="full")
+    task = asyncio.create_task(
+        _rollback_to_pre_run_checkpoint(
+            accessor=accessor,
+            checkpointer=checkpointer,
+            thread_id="thread-1",
+            run_id="run-1",
+            rollback_point=rollback_point,
+            snapshot_capture_failed=False,
+        )
+    )
+
+    try:
+        assert await asyncio.to_thread(checkpointer.mutation_started.wait, 1.0)
+        task.cancel()
+        await asyncio.sleep(0)
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+
+        assert not task.done(), "rollback returned before the synchronous checkpoint mutation finished"
+
+        checkpointer.allow_mutation.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert checkpointer.mutation_finished.is_set()
+    finally:
+        checkpointer.allow_mutation.set()
+        await asyncio.gather(task, return_exceptions=True)
+        await asyncio.to_thread(checkpointer.mutation_finished.wait, 1.0)