fix(runtime): keep thread checkpoints when a rollback cancel precedes snapshot capture (#6513)

* fix(runtime): keep thread checkpoints when a rollback cancel precedes snapshot capture

A rollback cancel that lands after try_start() flips the run to running but
before _capture_rollback_point() completes left rollback_point=None with
snapshot_capture_failed=False. The rollback finalizer treated that state as
"the thread never had a pre-run checkpoint" and called adelete_thread(thread_id),
deleting the whole conversation's checkpoints.

Initialize snapshot_capture_failed fail-closed and clear it only after the
capture completes. A completed capture that returns None still resets an empty
thread as before. Covered by two regressions in tests/test_run_worker_rollback.py
(cancellation before and during the capture await).

* fix(runtime): drop the redundant snapshot_capture_failed assignment

The fail-closed initializer owns the capture state: `snapshot_capture_failed`
is True from the top of run_agent, and the capture block runs once, so the
except branch can only observe it still True. Re-assigning it there was dead
code (review nit from @willem-bd).

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

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

@@ -22,7 +22,7 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
 
 **Wholesale state replacement uses a state-only mutation graph + `Overwrite`.** `update_state` values pass through channel reducers (`add_messages` merge in full, append in delta), so replacing reducer values requires `Overwrite` rather than an ordinary update. Full-mode rollback and context compaction replace `messages`; delta resume and delta rollback replace every materialized channel and reset current-head-only channels to their schema default (or `None`). These writes go through `build_state_mutation_graph(as_node, mode, state_schema)`, and `state_schema` MUST be the thread's effective schema (`graph_state_schema(assistant_graph)`), because the base-ThreadState fallback silently discards written channels contributed by custom `AgentMiddleware.state_schema`. Channels absent from a full-mode fork write inherit the parent's channel blobs, so middleware channels survive rollback/compaction (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`). The compiled mutation graph has one no-op node (entry = finish) whose checkpoint machinery (channels/versions/metadata) is identical to the agent graph's but schedules no pending tasks, so the restored/compacted head stays idle instead of re-triggering the agent. Never hand-write checkpoints via `checkpointer.aput` for this; raw writers elsewhere must preserve checkpoint parentage — severed ancestry breaks delta replay (see `runtime/runs/worker.py` writer parenting and `checkpoint_patches.py`).
 
-**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Resume and rollback state rewrites copy only the selected/captured checkpoint's server-authored agent binding; missing or malformed bindings remain unbound. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. Edit replay runs (`metadata.replay_kind="edit"`) also restore the pre-run checkpoint on failed, timed-out, or interrupted completion and publish the restored `values` snapshot to the stream before `end`, so clients do not remain on a transient edited branch when the replay did not produce a successful replacement.
+**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure, including cancellation delivered before or during the capture await, disables rollback (fail-closed), never restores partial state, and never deletes the thread; only a capture that completes and finds no checkpoint may reset the thread to empty. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Resume and rollback state rewrites copy only the selected/captured checkpoint's server-authored agent binding; missing or malformed bindings remain unbound. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. Edit replay runs (`metadata.replay_kind="edit"`) also restore the pre-run checkpoint on failed, timed-out, or interrupted completion and publish the restored `values` snapshot to the stream before `end`, so clients do not remain on a transient edited branch when the replay did not produce a successful replacement.
 
 **Message sequence placement:** Keep backend and frontend message identity rules aligned. Details: `backend/docs/runtime-guidance-details.md`.
 

backend/packages/harness/deerflow/runtime/runs/worker.py

@@ -938,7 +938,11 @@ async def run_agent(
     pre_run_workspace_snapshot: WorkspaceSnapshot | None = None
     workspace_changes_user_id: str | None = None
     workspace_excluded_dir_names: frozenset[str] | None = None
-    snapshot_capture_failed = False
+    # Fail closed until a pre-run capture completes: a cancellation delivered
+    # before or during capture leaves no rollback point, and treating that as
+    # "the thread had no pre-run checkpoint" would delete a live thread's
+    # checkpoints on a rollback cancel.
+    snapshot_capture_failed = True
     llm_error_fallback_message: str | None = None
     checkpoint_rollback_completed = False
     # Message ids checkpointed *before* this run started. The stream loop uses
@@ -1281,8 +1285,12 @@ async def run_agent(
                 try:
                     rollback_point = await _capture_rollback_point(accessor, checkpointer, checkpoint_config)
                 except Exception:
-                    snapshot_capture_failed = True
                     logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True)
+                else:
+                    # Only a completed capture proves what the pre-run state was.
+                    # ``None`` now unambiguously means the thread had no
+                    # checkpoint, so rollback may safely reset it to empty.
+                    snapshot_capture_failed = False
                 if rollback_point is not None:
                     pre_run_checkpoint_id = rollback_point.config.get("configurable", {}).get("checkpoint_id")
                     pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(rollback_point.messages)})

backend/tests/test_run_worker_rollback.py

@@ -1625,6 +1625,77 @@ async def test_run_agent_marks_rollback_unusable_when_capture_fails():
     assert rollback_kwargs["rollback_point"] is None
 
 
+@pytest.mark.anyio
+async def test_run_agent_cancel_before_snapshot_capture_does_not_delete_thread(monkeypatch):
+    """A rollback cancel before the pre-run capture must not wipe the thread.
+
+    The worker has already transitioned to ``running`` but has not captured a
+    rollback point yet, so ``rollback_point`` is None. That must never be
+    mistaken for "the thread had no pre-run checkpoint" — the only case where
+    rollback resets the thread to empty state.
+    """
+    run_manager = RunManager()
+    record = await run_manager.create("thread-1")
+    checkpointer = FakeCheckpointer()
+
+    async def _cancel_before_capture(*_args, **_kwargs):
+        record.abort_action = "rollback"
+        record.abort_event.set()
+        raise asyncio.CancelledError()
+
+    monkeypatch.setattr(
+        "deerflow.runtime.runs.worker.aensure_checkpoint_mode_compatible",
+        _cancel_before_capture,
+    )
+
+    await run_agent(
+        _lease_test_bridge(),
+        run_manager,
+        record,
+        ctx=RunContext(checkpointer=checkpointer),
+        agent_factory=lambda **_kwargs: SimpleNamespace(),
+        graph_input={},
+        config={},
+    )
+
+    checkpointer.adelete_thread.assert_not_awaited()
+    assert record.status == RunStatus.error
+    assert record.error == "Rolled back by user"
+
+
+@pytest.mark.anyio
+async def test_run_agent_cancel_during_snapshot_capture_does_not_delete_thread():
+    """Cancellation raised inside the capture await still disables rollback.
+
+    ``_capture_rollback_point`` is an await chain; a ``CancelledError`` raised
+    there is a ``BaseException`` that the capture's ``except Exception`` does
+    not catch, so no rollback point exists afterwards.
+    """
+    run_manager = RunManager()
+    record = await run_manager.create("thread-1")
+    checkpointer = FakeCheckpointer()
+
+    class CancelledCaptureAgent:
+        async def aget_state(self, _config):
+            record.abort_action = "rollback"
+            record.abort_event.set()
+            raise asyncio.CancelledError()
+
+    await run_agent(
+        _lease_test_bridge(),
+        run_manager,
+        record,
+        ctx=RunContext(checkpointer=checkpointer),
+        agent_factory=lambda **_kwargs: CancelledCaptureAgent(),
+        graph_input={},
+        config={},
+    )
+
+    checkpointer.adelete_thread.assert_not_awaited()
+    assert record.status == RunStatus.error
+    assert record.error == "Rolled back by user"
+
+
 @pytest.mark.anyio
 async def test_run_agent_overrides_spoofed_pre_existing_message_ids_without_snapshot():
     run_manager = RunManager()