fix(runs): serialize run rollback to prevent state overwrites

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

@@ -12,6 +12,16 @@
 
 logger = logging.getLogger(__name__)
 
+# Statuses that mean "a worker is still doing something for this run".
+# Used by has_inflight, create_or_reject, and any caller that needs to know
+# whether new mutable work for the same thread should wait for cleanup.
+_INFLIGHT_STATUSES = (
+    RunStatus.pending,
+    RunStatus.running,
+    RunStatus.cancelling,
+    RunStatus.rolling_back,
+)
+
 
 def _now_iso() -> str:
     return datetime.now(UTC).isoformat()
@@ -106,7 +116,10 @@ async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool:
             action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state.
 
         Sets the abort event with the action reason and cancels the asyncio task.
-        Returns ``True`` if the run was in-flight and cancellation was initiated.
+        Transitions the run to ``cancelling`` immediately; the worker is
+        responsible for the subsequent ``cancelling -> rolling_back ->
+        interrupted`` (or ``-> error``) progression. Returns ``True`` if the
+        run was in-flight and cancellation was initiated.
         """
         async with self._lock:
             record = self._runs.get(run_id)
@@ -118,9 +131,9 @@ async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool:
             record.abort_event.set()
             if record.task is not None and not record.task.done():
                 record.task.cancel()
-            record.status = RunStatus.interrupted
+            record.status = RunStatus.cancelling
             record.updated_at = _now_iso()
-        logger.info("Run %s cancelled (action=%s)", run_id, action)
+        logger.info("Run %s cancelling (action=%s)", run_id, action)
         return True
 
     async def create_or_reject(
@@ -147,30 +160,50 @@ async def create_or_reject(
 
         _supported_strategies = ("reject", "interrupt", "rollback")
 
+        # Hold the lock only for the inflight check + cancel-signal phase. We
+        # then release it to await the worker tasks (so other operations on
+        # other threads aren't blocked), and reacquire it to insert the new
+        # record. The new record is created with a fresh run_id and is keyed
+        # by run_id, so the brief unlocked window cannot collide.
+        tasks_to_await: list[asyncio.Task] = []
         async with self._lock:
             if multitask_strategy not in _supported_strategies:
                 raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}")
 
-            inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running)]
+            inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in _INFLIGHT_STATUSES]
 
             if multitask_strategy == "reject" and inflight:
                 raise ConflictError(f"Thread {thread_id} already has an active run")
 
             if multitask_strategy in ("interrupt", "rollback") and inflight:
                 for r in inflight:
-                    r.abort_action = multitask_strategy
-                    r.abort_event.set()
+                    # Skip runs that are already cancelling/rolling_back -- the
+                    # earlier signaller already armed them. We still wait on
+                    # their tasks below so the new run is serialized after them.
+                    if r.status in (RunStatus.pending, RunStatus.running):
+                        r.abort_action = multitask_strategy
+                        r.abort_event.set()
+                        if r.task is not None and not r.task.done():
+                            r.task.cancel()
+                        r.status = RunStatus.cancelling
+                        r.updated_at = now
                     if r.task is not None and not r.task.done():
-                        r.task.cancel()
-                    r.status = RunStatus.interrupted
-                    r.updated_at = now
+                        tasks_to_await.append(r.task)
                 logger.info(
-                    "Cancelled %d inflight run(s) on thread %s (strategy=%s)",
+                    "Cancelling %d inflight run(s) on thread %s (strategy=%s)",
                     len(inflight),
                     thread_id,
                     multitask_strategy,
                 )
 
+        # Wait for cancelled workers to finish their cleanup (rollback or
+        # interrupt) BEFORE creating the new run. This is the serialization
+        # point that closes the race in #2505: an old worker can no longer
+        # restore an older snapshot over the new run's state.
+        if tasks_to_await:
+            await asyncio.gather(*tasks_to_await, return_exceptions=True)
+
+        async with self._lock:
             record = RunRecord(
                 run_id=run_id,
                 thread_id=thread_id,
@@ -189,9 +222,13 @@ async def create_or_reject(
         return record
 
     async def has_inflight(self, thread_id: str) -> bool:
-        """Return ``True`` if *thread_id* has a pending or running run."""
+        """Return ``True`` if *thread_id* has a run that hasn't reached a terminal state.
+
+        ``cancelling`` and ``rolling_back`` count as inflight too -- a new
+        mutable run for the same thread must wait until cleanup completes.
+        """
         async with self._lock:
-            return any(r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running) for r in self._runs.values())
+            return any(r.thread_id == thread_id and r.status in _INFLIGHT_STATUSES for r in self._runs.values())
 
     async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
         """Remove a run record after an optional delay."""

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

@@ -8,6 +8,8 @@ class RunStatus(StrEnum):
 
     pending = "pending"
     running = "running"
+    cancelling = "cancelling"
+    rolling_back = "rolling_back"
     success = "success"
     error = "error"
     timeout = "timeout"

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

@@ -184,7 +184,10 @@ async def run_agent(
         if record.abort_event.is_set():
             action = record.abort_action
             if action == "rollback":
-                await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user")
+                # Mark rolling_back BEFORE the rollback work so create_or_reject
+                # for the same thread sees this run as still inflight and
+                # blocks until rollback completes (#2505).
+                await run_manager.set_status(run_id, RunStatus.rolling_back)
                 try:
                     await _rollback_to_pre_run_checkpoint(
                         checkpointer=checkpointer,
@@ -195,8 +198,10 @@ async def run_agent(
                         snapshot_capture_failed=snapshot_capture_failed,
                     )
                     logger.info("Run %s rolled back to pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
+                    await run_manager.set_status(run_id, RunStatus.interrupted, error="Rolled back by user")
                 except Exception:
                     logger.warning("Failed to rollback checkpoint for run %s", run_id, exc_info=True)
+                    await run_manager.set_status(run_id, RunStatus.error, error="Rollback failed")
             else:
                 await run_manager.set_status(run_id, RunStatus.interrupted)
         else:
@@ -205,7 +210,7 @@ async def run_agent(
     except asyncio.CancelledError:
         action = record.abort_action
         if action == "rollback":
-            await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user")
+            await run_manager.set_status(run_id, RunStatus.rolling_back)
             try:
                 await _rollback_to_pre_run_checkpoint(
                     checkpointer=checkpointer,
@@ -216,8 +221,10 @@ async def run_agent(
                     snapshot_capture_failed=snapshot_capture_failed,
                 )
                 logger.info("Run %s was cancelled and rolled back", run_id)
+                await run_manager.set_status(run_id, RunStatus.interrupted, error="Rolled back by user")
             except Exception:
                 logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True)
+                await run_manager.set_status(run_id, RunStatus.error, error="Rollback failed")
         else:
             await run_manager.set_status(run_id, RunStatus.interrupted)
             logger.info("Run %s was cancelled", run_id)

backend/tests/test_run_manager.py

@@ -1,5 +1,6 @@
 """Tests for RunManager."""
 
+import asyncio
 import re
 
 import pytest
@@ -53,14 +54,18 @@ async def test_status_transitions(manager: RunManager):
 
 @pytest.mark.anyio
 async def test_cancel(manager: RunManager):
-    """Cancel should set abort_event and transition to interrupted."""
+    """Cancel should set abort_event and transition to cancelling.
+
+    Final state (interrupted/error) is set by the worker after cleanup;
+    see #2505 for the rollback-serialization rationale.
+    """
     record = await manager.create("thread-1")
     await manager.set_status(record.run_id, RunStatus.running)
 
     cancelled = await manager.cancel(record.run_id)
     assert cancelled is True
     assert record.abort_event.is_set()
-    assert record.status == RunStatus.interrupted
+    assert record.status == RunStatus.cancelling
 
 
 @pytest.mark.anyio
@@ -141,3 +146,93 @@ async def test_create_defaults(manager: RunManager):
     assert record.kwargs == {}
     assert record.multitask_strategy == "reject"
     assert record.assistant_id is None
+
+
+# --- #2505: rollback-serialization regression tests ---
+
+
+@pytest.mark.anyio
+async def test_has_inflight_includes_cancelling_and_rolling_back(manager: RunManager):
+    """Cancelling/rolling_back must count as inflight so a new run waits."""
+    a = await manager.create("thread-1")
+    await manager.set_status(a.run_id, RunStatus.cancelling)
+    assert await manager.has_inflight("thread-1") is True
+
+    await manager.set_status(a.run_id, RunStatus.rolling_back)
+    assert await manager.has_inflight("thread-1") is True
+
+    await manager.set_status(a.run_id, RunStatus.interrupted)
+    assert await manager.has_inflight("thread-1") is False
+
+
+@pytest.mark.anyio
+async def test_create_or_reject_awaits_cancelled_workers_before_creating(manager: RunManager):
+    """Interrupt/rollback strategies must wait for the old worker to finish.
+
+    Before #2505 was fixed, create_or_reject would mark the inflight run
+    as ``interrupted`` and immediately insert the new run, leaving the old
+    worker free to write rollback state on top of the new run. The fix
+    awaits the worker's task before insertion so the new run is serialized
+    after cleanup.
+    """
+    # Old run with a still-running worker task.
+    old = await manager.create("thread-1", multitask_strategy="rollback")
+    await manager.set_status(old.run_id, RunStatus.running)
+
+    finished_in_correct_order: list[str] = []
+
+    async def fake_worker() -> None:
+        try:
+            # The cancel signal hits before this completes naturally.
+            await asyncio.sleep(0.5)
+        except asyncio.CancelledError:
+            # Simulate rollback cleanup: takes time, sets status during.
+            await manager.set_status(old.run_id, RunStatus.rolling_back)
+            await asyncio.sleep(0.05)
+            await manager.set_status(old.run_id, RunStatus.interrupted)
+            finished_in_correct_order.append("worker_done")
+            raise
+
+    old.task = asyncio.create_task(fake_worker())
+    # Yield once so the task starts.
+    await asyncio.sleep(0)
+
+    # Kick off create_or_reject. It must wait for the worker.
+    new_run = await manager.create_or_reject(
+        "thread-1",
+        multitask_strategy="rollback",
+    )
+    finished_in_correct_order.append("new_created")
+
+    # Worker finished BEFORE the new run was created.
+    assert finished_in_correct_order == ["worker_done", "new_created"]
+    assert old.status == RunStatus.interrupted
+    assert new_run.status == RunStatus.pending
+    assert new_run.thread_id == "thread-1"
+
+
+@pytest.mark.anyio
+async def test_create_or_reject_skips_already_cancelling_runs(manager: RunManager):
+    """Re-cancelling a cancelling run is a no-op for state, but still awaited.
+
+    If two cancellations race in for the same run, the second one must
+    not stomp on the first one's abort_action. We still wait on the task
+    so the new run is serialized after cleanup either way.
+    """
+    old = await manager.create("thread-1", multitask_strategy="rollback")
+    await manager.set_status(old.run_id, RunStatus.running)
+
+    # First signaller: rollback action.
+    cancelled = await manager.cancel(old.run_id, action="rollback")
+    assert cancelled is True
+    assert old.status == RunStatus.cancelling
+    assert old.abort_action == "rollback"
+
+    # Old has no real task; create_or_reject must still proceed and not
+    # overwrite abort_action.
+    new_run = await manager.create_or_reject(
+        "thread-1",
+        multitask_strategy="interrupt",
+    )
+    assert old.abort_action == "rollback"  # not stomped
+    assert new_run.status == RunStatus.pending