diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index dc491e0806..91e9a23d7d 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/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`. diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 2f3023b5a5..25ea918b85 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/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)}) diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index b456a197a7..32ac0d2588 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/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()