WIP returning results from plans

Author: tpoliaw

src/blueapi/core/context.py

@@ -120,7 +120,7 @@ class BlueskyContext:
     configuration: InitVar[ApplicationConfig | None] = None
 
     run_engine: RunEngine = field(
-        default_factory=lambda: RunEngine(context_managers=[])
+        default_factory=lambda: RunEngine(context_managers=[], call_returns_result=True)
     )
     tiled_conf: TiledConfig | None = field(default=None, init=False, repr=False)
     numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False)

src/blueapi/worker/event.py

@@ -1,5 +1,6 @@
 from collections.abc import Mapping
 from enum import Enum
+from typing import Any
 
 from bluesky.run_engine import RunEngineStateMachine
 from pydantic import Field
@@ -109,6 +110,7 @@ class TaskStatus(BlueapiBaseModel):
     task_id: str
     task_complete: bool
     task_failed: bool
+    result: Any = None
 
 
 class WorkerEvent(BlueapiBaseModel):

src/blueapi/worker/task.py

@@ -38,7 +38,10 @@ def do_task(self, ctx: BlueskyContext) -> None:
         func = ctx.plan_functions[self.name]
         prepared_params = self.prepare_params(ctx)
         ctx.run_engine.md.update(self.metadata)
-        ctx.run_engine(func(**prepared_params))
+        result = ctx.run_engine(func(**prepared_params))
+        if isinstance(result, tuple):
+            return None
+        return result.plan_result
 
 
 def _lookup_params(ctx: BlueskyContext, task: Task) -> BaseModel:

src/blueapi/worker/task_worker.py

@@ -69,6 +69,7 @@ class TrackableTask(BlueapiBaseModel):
     is_complete: bool = False
     is_pending: bool = True
     errors: list[str] = Field(default_factory=list)
+    result: Any | None = None
 
 
 class TaskWorker:
@@ -423,11 +424,12 @@ def _cycle(self) -> None:
             next_task: TrackableTask | KillSignal = self._task_channel.get()
             if isinstance(next_task, TrackableTask):
 
-                def process_task():
+                def process_task() -> Any:
                     LOGGER.info(f"Got new task: {next_task}")
                     self._current = next_task
                     self._current.is_pending = False
-                    self._current.task.do_task(self._ctx)
+                    result = self._current.task.do_task(self._ctx)
+                    self._current.result = result
 
                 with plan_tag_filter_context(next_task.task.name, LOGGER):
                     if self._current_task_otel_context is not None:
@@ -528,6 +530,7 @@ def _report_status(
                 task_id=self._current.task_id,
                 task_complete=self._current.is_complete,
                 task_failed=bool(self._current.errors),
+                result=self._current.result,
             )
             correlation_id = self._current.task_id
             add_span_attributes(