feat: Blueapi plan pause

src/blueapi/worker/task_worker.py

@@ -8,6 +8,7 @@
 from threading import Event, RLock
 from typing import Any, TypeVar
 
+from bluesky import RunEngineInterrupted
 from bluesky._vendor.super_state_machine.errors import TransitionError
 from bluesky.protocols import Status
 from observability_utils.tracing import (
@@ -109,6 +110,7 @@ class TaskWorker:
 
     _task_channel: Queue  # type: ignore
     _current: TrackableTask | None
+    _pending_cancel: "CancelSignal | None"
     _status_lock: RLock
     _status_snapshot: dict[str, StatusView]
     _completed_statuses: set[str]
@@ -139,6 +141,7 @@ def __init__(
         self._warnings = []
         self._task_channel = Queue(maxsize=1)
         self._current = None
+        self._pending_cancel = None
         self._worker_events = EventPublisher()
         self._progress_events = EventPublisher()
         self._data_events = EventPublisher()
@@ -180,19 +183,44 @@ def cancel_active_task(
         Returns:
             The task_id of the active task
         """
-        if self._current is None:
+        current = self._current
+        if current is None:
             # Persuades type checker that self._current is not None
             # We only allow this method to be called if a Plan is active
             raise TransitionError("Attempted to cancel while no active Task")
+
+        if self._ctx.run_engine.state == "paused":
+            # abort()/stop() block until cleanup finishes, so defer to the worker
+            # thread. Also recorded in _pending_cancel so a queued ResumeSignal
+            # can't race ahead and finish the task before this is looked at.
+            signal = CancelSignal(failure=failure, reason=reason)
+            self._pending_cancel = signal
+            try:
+                self._task_channel.put_nowait(signal)
+            except Full:
+                pass  # a signal already queued will check _pending_cancel
+            return current.task_id
+
+        # RE.abort()/stop() are thread-safe and must be called immediately -
+        # putting only a CancelSignal would no-op if the worker is blocked in
+        # do_task(). The outcome is set beforehand so it's already in place
+        # once the worker thread's do_task() unblocks and finalizes.
+        default_reason = "Task failed for unknown reason"
         if failure:
-            default_reason = "Task failed for unknown reason"
+            with self._status_lock:
+                if current.outcome is None:
+                    current.set_exception(Exception(reason or default_reason))
             self._ctx.run_engine.abort(reason or default_reason)
-            add_span_attributes({"Task aborted": reason or default_reason})
         else:
+            with self._status_lock:
+                if current.outcome is None:
+                    current.set_result(None)
             self._ctx.run_engine.stop()
-            default_reason = "Cancellation successful: Task stopped without error"
-            add_span_attributes({"Task stopped": reason or default_reason})
-        return self._current.task_id
+        self._task_channel.put(CancelSignal(failure=failure, reason=reason))
+        add_span_attributes(
+            {"Task aborted" if failure else "Task stopped": reason or ""}
+        )
+        return current.task_id
 
     @start_as_current_span(TRACER, "task_id")
     def get_task_by_id(self, task_id: str) -> TrackableTask | None:
@@ -410,7 +438,7 @@ def resume(self):
         Command the worker to resume
         """
         LOGGER.info("Requesting to resume the worker")
-        self._ctx.run_engine.resume()
+        self._task_channel.put(ResumeSignal())
 
     @start_as_current_span(TRACER)
     def _cycle_with_error_handling(self) -> None:
@@ -436,11 +464,22 @@ def process_task():
                         LOGGER.info(
                             "Task ran successfully - returned: %s", result, extra=meta
                         )
-                        self._current.set_result(result)
+                        with self._status_lock:
+                            # cancel_active_task() may have set this concurrently.
+                            if self._current.outcome is None:
+                                self._current.set_result(result)
+                    except RunEngineInterrupted:
+                        # Raised by both a pause (outcome still None) and an
+                        # abort (outcome already TaskError) - only the latter
+                        # is a failure.
+                        if isinstance(self._current.outcome, TaskError):
+                            self._report_error(Exception(self._current.outcome.message))
                     except Exception as e:
                         LOGGER.error("Task failed", extra=meta)
-                        self._current.set_exception(e)
-                        self._report_error(e)
+                        with self._status_lock:
+                            if self._current.outcome is None:
+                                self._current.set_exception(e)
+                        raise
 
                 with plan_tag_filter_context(next_task.task.name, LOGGER):
                     if self._current_task_otel_context is not None:
@@ -458,28 +497,71 @@ def process_task():
                     else:
                         process_task()
 
+            elif isinstance(next_task, ResumeSignal):
+                pending_cancel = self._pending_cancel
+                if pending_cancel is not None:
+                    # A cancel queued after this resume takes priority, so
+                    # this resume never runs.
+                    self._apply_cancel(pending_cancel)
+                elif self._ctx.run_engine.state == "paused":
+                    if self._current is not None:
+                        try:
+                            result = self._ctx.run_engine.resume()
+                            self._current.set_result(result)
+                        except RunEngineInterrupted:
+                            # Plan paused again immediately - not a failure,
+                            # leave the outcome unset so it stays resumable.
+                            LOGGER.debug("RunEngine resume interrupted; ignoring")
+                    else:
+                        LOGGER.warning(
+                            "Received resume signal but no active task, ignoring"
+                        )
+                else:
+                    LOGGER.warning(
+                        "Received resume signal but RunEngine is not paused, ignoring"
+                    )
+
+            elif isinstance(next_task, CancelSignal):
+                self._apply_cancel(self._pending_cancel or next_task)
+
             elif isinstance(next_task, KillSignal):
                 # If we receive a kill signal we begin to shut the worker down.
                 # Note that the kill signal is explicitly not a type of task as we don't
                 # want it to be part of the worker's public API
+                self._pending_cancel = None
+                if self._current is not None and self._ctx.run_engine.state == "paused":
+                    self._apply_cancel(
+                        CancelSignal(
+                            failure=True,
+                            reason="Worker is stopping while the task was paused",
+                        )
+                    )
                 self._stopping.set()
                 add_span_attributes({"server shutting down": "true"})
             else:
                 raise KeyError(f"Unknown command: {next_task}")
         except Exception as err:
             self._report_error(err)
         finally:
-            if self._current_task_otel_context is not None:
+            if (
+                self._current_task_otel_context is not None
+                and self._ctx.run_engine.state not in ["panicked", "paused"]
+            ):
                 self._current_task_otel_context = None
 
         if self._current is not None:
-            self._current.is_complete = True
-            self._pending_tasks.pop(self._current.task_id)
-            self._completed_tasks[self._current.task_id] = self._current
-        self._report_status()
-        self._errors.clear()
-        self._warnings.clear()
-        self._completed_statuses.clear()
+            # Not done yet while paused - it may still be resumed or cancelled.
+            if self._ctx.run_engine.state != "paused":
+                self._current.is_complete = True
+                self._pending_tasks.pop(self._current.task_id)
+                self._completed_tasks[self._current.task_id] = self._current
+        if self._ctx.run_engine.state != "paused":
+            finished_task = self._current
+            self._current = None
+            self._report_status(finished_task)
+            self._errors.clear()
+            self._warnings.clear()
+            self._completed_statuses.clear()
 
     @property
     def worker_events(self) -> EventStream[WorkerEvent, int]:
@@ -521,34 +603,59 @@ def _on_state_change(
             old_state = WorkerState.UNKNOWN
         LOGGER.debug(f"Notifying state change {old_state} -> {new_state}")
         self._state = new_state
-        self._report_status()
+        self._report_status(self._current)
 
     def _report_error(self, err: Exception) -> None:
         LOGGER.error(err, exc_info=True)
         if self._current is not None:
             self._current.errors.append(str(err))
         self._errors.append(str(err))
 
+    def _apply_cancel(self, signal: "CancelSignal") -> None:
+        self._pending_cancel = None
+        default_reason = "Task failed for unknown reason"
+        if self._current is not None:
+            if signal.failure:
+                reason = signal.reason or default_reason
+                self._ctx.run_engine.abort(reason)
+                self._current.set_exception(Exception(reason))
+            else:
+                self._ctx.run_engine.stop()
+                self._current.set_result(None)
+
+        if signal.failure:
+            error_message = signal.reason or default_reason
+            add_span_attributes({"Task aborted": error_message})
+            LOGGER.error("Task failed: %s", error_message)
+            if self._current is not None:
+                self._report_error(Exception(error_message))
+        else:
+            add_span_attributes(
+                {
+                    "Task stopped": signal.reason
+                    or "Cancellation successful: Task stopped without error"
+                }
+            )
+
     @start_as_current_span(TRACER)
-    def _report_status(
-        self,
-    ) -> None:
+    def _report_status(self, current: TrackableTask | None) -> None:
         task_status: TaskStatus | None
         errors = self._errors
         warnings = self._warnings
-        if self._current is not None:
+        if current is not None:
             task_status = TaskStatus(
-                task_id=self._current.task_id,
-                task_complete=self._current.is_complete,
-                task_failed=bool(self._current.errors),
-                result=self._current.outcome,
+                task_id=current.task_id,
+                task_complete=current.is_complete,
+                task_failed=bool(current.errors)
+                or isinstance(current.outcome, TaskError),
+                result=current.outcome,
             )
-            correlation_id = self._current.task_id
+            correlation_id = current.task_id
             add_span_attributes(
                 {
-                    "task_id": self._current.task_id,
-                    "task_complete": self._current.is_complete,
-                    "task_failed": self._current.errors,
+                    "task_id": current.task_id,
+                    "task_complete": current.is_complete,
+                    "task_failed": current.errors,
                 }
             )
         else:
@@ -598,7 +705,7 @@ def _on_document(self, name: str, document: Mapping[str, Any]) -> None:
                 )
 
         else:
-            raise KeyError(
+            raise RuntimeError(
                 "Trying to emit a document despite the fact that the RunEngine is idle"
             )
 
@@ -683,6 +790,25 @@ class KillSignal:
     ...
 
 
+@dataclass
+class ResumeSignal:
+    """
+    Object put in the worker's task queue to tell it to resume if paused.
+    """
+
+    pass
+
+
+@dataclass
+class CancelSignal:
+    """
+    Object put in the worker's task queue to tell it to cancel the current task.
+    """
+
+    failure: bool
+    reason: str | None
+
+
 def run_worker_in_own_thread(
     worker: TaskWorker, executor: ThreadPoolExecutor | None = None
 ) -> Future:

tests/system_tests/test_blueapi_system.py

@@ -409,11 +409,14 @@ def test_delete_non_existent_task(rest_client: BlueapiRestClient):
         rest_client.clear_task("Not-exists")
 
 
-def test_put_worker_task(rest_client: BlueapiRestClient, small_task: TaskRequest):
-    created_task = rest_client.create_task(small_task)
+def test_put_worker_task(rest_client: BlueapiRestClient, long_task: TaskRequest):
+    # long_task, since a near-instant task could complete (clearing the
+    # active task) before get_active_task() below is called.
+    created_task = rest_client.create_task(long_task)
     rest_client.update_worker_task(WorkerTask(task_id=created_task.task_id))
     active_task = rest_client.get_active_task()
     assert active_task.task_id == created_task.task_id
+    rest_client.cancel_current_task(WorkerState.ABORTING)
     rest_client.clear_task(created_task.task_id)
 
 
@@ -806,14 +809,18 @@ def test_any_user_can_retrieve_active_task(
     task_id = (
         client_factory[AdminUser.admin]
         .create_and_start_task(
-            task_factory(AdminUser.admin, VALID_INSTRUMENT_SESSION[AdminUser.admin])
+            task_factory(
+                AdminUser.admin, VALID_INSTRUMENT_SESSION[AdminUser.admin], time=1
+            )
         )
         .task_id
     )
 
     for user in User:
         assert client_factory[user].get_active_task().task_id == task_id
 
+    client_factory[AdminUser.admin].abort()
+
 
 def test_non_admin_can_only_start_own_tasks(
     client_factory: dict[ValidUser, BlueapiClient],