feat: Run blocking plans via websockets

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "PyYAML>=6.0.2",
     "click>=8.2.0",
     "fastapi>=0.112.0",
-    "uvicorn",
+    "uvicorn>=0.52.1",
     "requests",
     "GitPython",
     "event-model==1.23.1",                           # https://github.com/DiamondLightSource/blueapi/issues/684

src/blueapi/cli/cli.py

@@ -321,6 +321,7 @@ def on_event(
 @controller.command(name="run")
 @click.argument("name", type=str)
 @click.argument("parameters", type=ParametersType(), default={}, required=False)
+@click.option("--ws", type=bool, is_flag=True, default=False)
 @click.option(
     "--foreground/--background", "--fg/--bg", type=bool, is_flag=True, default=True
 )
@@ -348,6 +349,7 @@ def run_plan(
     name: str,
     timeout: float | None,
     foreground: bool,
+    ws: bool,
     instrument_session: str,
     parameters: TaskParameters,
 ) -> None:
@@ -374,7 +376,13 @@ def on_event(event: AnyEvent) -> None:
                 elif isinstance(event, DataEvent):
                     callback(event.name, event.doc)
 
-            resp = client.run_task(task, on_event=on_event)
+            client.add_callback(on_event)
+
+            if ws:
+                resp = client.run_blocking(task)
+            else:
+                resp = client.run_task(task)
+
             match resp.result:
                 case TaskResult(result=None, type="NoneType"):
                     print("Plan succeeded")

src/blueapi/client/client.py

@@ -459,6 +459,27 @@ def get_active_task(self) -> WorkerTask:
 
         return self.active_task
 
+    @start_as_current_span(TRACER, "request")
+    def run_blocking(
+        self, request: TaskRequest, on_event: OnAnyEvent | None = None
+    ) -> TaskStatus:
+        for event in self._rest.run_blocking(request):
+            if on_event is not None:
+                on_event(event)
+            for cb in self._callbacks.values():
+                try:
+                    cb(event)
+                except Exception as e:
+                    log.error(f"Callback ({cb}) failed for event: {event}", exc_info=e)
+            if isinstance(event, WorkerEvent) and event.is_complete():
+                # task_status will always be present if event is complete
+                if event.task_status is None:  # pragma: no cover
+                    raise BlueskyRemoteControlError(
+                        "Server completed without task status"
+                    )
+                return event.task_status
+        raise BlueskyRemoteControlError("Connection closed before plan completed.")
+
     @start_as_current_span(TRACER, "task", "timeout")
     def run_task(
         self,

src/blueapi/client/rest.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from typing import Any, Literal, TypeVar
 
 import requests
@@ -12,10 +12,13 @@
 )
 from pydantic import BaseModel, TypeAdapter, ValidationError
 from pydantic_core import PydanticSerializationError
+from websockets.exceptions import InvalidStatus
+from websockets.sync.client import connect
 
 from blueapi import __version__
 from blueapi.client import client
 from blueapi.config import RestConfig
+from blueapi.core.bluesky_types import DataEvent
 from blueapi.service.authentication import JWTAuth, SessionManager
 from blueapi.service.model import (
     DeviceModel,
@@ -31,14 +34,26 @@
     TasksListResponse,
     WorkerTask,
 )
+from blueapi.service.protocol import (
+    ControlResponse,
+    InvalidArgs,
+    PlanNotFound,
+    ServerBusy,
+    Submit,
+    Unauthorized,
+    Update,
+)
 from blueapi.worker import TrackableTask, WorkerState
+from blueapi.worker.event import ProgressEvent, WorkerEvent
 
 T = TypeVar("T")
 
 TRACER = get_tracer("rest")
 
 LOGGER = logging.getLogger(__name__)
 
+USER_AGENT = f"blueapi cli {__version__}"
+
 
 class BlueskyRequestError(Exception):
     """An error response from the blueapi server."""
@@ -86,8 +101,8 @@ def __init__(self, target_type: type) -> None:
 
 class ParameterError(BaseModel):
     loc: list[str | int]
-    msg: str
-    type: str
+    msg: str | None
+    type: str | None
     input: Any
 
     def field(self):
@@ -307,14 +322,15 @@ def _request_and_deserialize(
     ) -> T:
         url = self._config.url.unicode_string().removesuffix("/") + suffix
         # Get the trace context to propagate to the REST API
-        carr = get_context_propagator()
+        headers = get_context_propagator()
+        headers["User-Agent"] = USER_AGENT
         try:
             response = self._pool.request(
                 method,
                 url,
                 json=data,
                 params=params,
-                headers=carr,
+                headers=headers,
                 auth=JWTAuth(self._session_manager),
             )
         except requests.exceptions.ConnectionError as ce:
@@ -340,6 +356,52 @@ def _request_and_deserialize(
         )
         return deserialized
 
+    def run_blocking(
+        self, req: TaskRequest
+    ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]:
+        url = self._config.ws_address.unicode_string().rstrip("/") + "/api/v2/run_plan"
+        headers = get_context_propagator()
+        if self._session_manager:
+            auth = self._session_manager.get_valid_access_token()
+            headers["Authorization"] = f"Bearer {auth}"
+        try:
+            with connect(
+                url,
+                additional_headers=headers,
+                user_agent_header=USER_AGENT,
+            ) as ws:
+                ws.send(Submit(task=req).model_dump_json())
+                for message in ws:
+                    event = ControlResponse.validate_json(message)
+                    match event:
+                        case Update(data=data):
+                            yield data
+                        case InvalidArgs(errors=errors):
+                            raise InvalidParametersError(
+                                [
+                                    ParameterError(
+                                        loc=e.loc, msg=e.msg, type=e.type, input=e.input
+                                    )
+                                    for e in errors
+                                ]
+                            )
+                        case PlanNotFound(plan_name=name):
+                            raise UnknownPlanError(message=name)
+                        case ServerBusy():
+                            raise BlueskyRemoteControlError(409, "Server is busy")
+                        case Unauthorized():
+                            raise UnauthorisedAccessError(
+                                403, "Not authorized to submit task"
+                            )
+        except InvalidStatus as istat:
+            match istat.response.status_code:
+                case 401 | 403:
+                    raise UnauthorisedAccessError() from None
+                case _:
+                    raise BlueskyRemoteControlError() from istat
+        except ConnectionRefusedError as cre:
+            raise ServiceUnavailableError() from cre
+
 
 # https://github.com/DiamondLightSource/blueapi/issues/1256 - remove before 2.0
 def __getattr__(name: str):

src/blueapi/config.py

@@ -22,6 +22,7 @@
     TypeAdapter,
     UrlConstraints,
     ValidationError,
+    WebsocketUrl,
     field_validator,
     model_validator,
 )
@@ -170,6 +171,22 @@ class RestConfig(BlueapiBaseModel):
     url: HttpUrl = HttpUrl("http://localhost:8000")
     cors: CORSConfig | None = None
 
+    @property
+    def ws_address(self) -> WebsocketUrl:
+        api = self.url
+        if api.host is None:
+            # type hints say it could be None but not possible to construct
+            # HttpUrl without host
+            raise ValueError("No host configured")  # pragma: no cover
+        scheme = "ws" if api.scheme == "http" else "wss"
+
+        # HttpUrl adds "/" to the start of paths, even if none was specified so
+        # remove existing leading '/' to prevent duplication
+        path = (api.path or "").removeprefix("/")
+        return WebsocketUrl.build(
+            scheme=scheme, host=api.host, port=api.port, path=path
+        )
+
 
 class ScratchRepository(BlueapiBaseModel):
     name: str = Field(

src/blueapi/service/authentication.py

@@ -15,7 +15,8 @@
 import httpx
 import jwt
 import requests
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, HTTPException
+from fastapi.requests import HTTPConnection
 from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import TypeAdapter
 from requests.auth import AuthBase
@@ -278,14 +279,17 @@ def sync_auth_flow(self, request):
         yield request
 
 
-def unchecked_bearer_token(req: Request) -> str | None:
+def unchecked_bearer_token(req: HTTPConnection) -> str | None:
     """Get bearer token value from authorization header"""
+
+    auth_header = req.headers.get("Authorization")
+    auth_cookie = req.cookies.get("Authorization")
+
     # This is an abridged version of the same feature of
     # OAuth2AuthorizationCodeBearer from fastapi. Replicating here prevents
     # passing unused configuration and means the schema does not include auth
     # details for servers that do not support it.
-    auth = req.headers.get("Authorization")
-    scheme, param = get_authorization_scheme_param(auth)
+    scheme, param = get_authorization_scheme_param(auth_header or auth_cookie)
     if scheme.casefold() != "bearer":
         return None
     return param.strip()
@@ -303,7 +307,7 @@ def build_access_token_check(config: OIDCConfig):
     """
     jwkclient = jwt.PyJWKClient(config.jwks_uri)
 
-    def validate_bearer_token(request: Request, token: UncheckedBearerToken):
+    def validate_bearer_token(request: HTTPConnection, token: UncheckedBearerToken):
         """Check that a bearer token is valid and inject into request state"""
         if not token:
             raise HTTPException(
@@ -326,7 +330,7 @@ def validate_bearer_token(request: Request, token: UncheckedBearerToken):
     return validate_bearer_token
 
 
-def access_token(request: Request) -> Mapping[str, Any] | None:
+def access_token(request: HTTPConnection) -> Mapping[str, Any] | None:
     """Get the decoded and verified access token of the user making the request"""
     return getattr(request.state, "decoded_access_token", None)
 

src/blueapi/service/authorization.py

@@ -4,7 +4,8 @@
 from typing import Annotated, Any, Self, cast
 
 from aiohttp import ClientSession
-from fastapi import Depends, HTTPException, Request
+from fastapi import Depends, HTTPException
+from fastapi.requests import HTTPConnection
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
 from blueapi.config import OIDCConfig, OpaConfig, ServiceAccount
@@ -114,7 +115,7 @@ async def validate_tiled_config(
 
 
 async def opa(
-    request: Request, token: str | None = Depends(unchecked_bearer_token)
+    request: HTTPConnection, token: str | None = Depends(unchecked_bearer_token)
 ) -> OpaUserClient | None:
 
     if opa := cast(OpaClient | None, getattr(request.app.state, "authz", None)):

src/blueapi/service/interface.py

@@ -1,5 +1,8 @@
+import logging
 from collections.abc import Mapping
+from dataclasses import dataclass
 from functools import cache
+from multiprocessing.connection import Connection
 from typing import Any
 
 from bluesky.callbacks.tiled_writer import TiledWriter
@@ -9,6 +12,7 @@
 
 from blueapi.cli.scratch import get_python_environment
 from blueapi.config import ApplicationConfig, OIDCConfig, ServiceAccount, StompConfig
+from blueapi.core.bluesky_types import DataEvent
 from blueapi.core.context import BlueskyContext
 from blueapi.core.event import EventStream
 from blueapi.log import set_up_logging
@@ -22,14 +26,14 @@
     WorkerTask,
 )
 from blueapi.utils.serialization import access_blob
-from blueapi.worker.event import TaskStatusEnum, WorkerEvent, WorkerState
+from blueapi.worker.event import ProgressEvent, TaskStatusEnum, WorkerEvent, WorkerState
 from blueapi.worker.task import Task
 from blueapi.worker.task_worker import TaskWorker, TrackableTask
 
 """This module provides interface between web application and underlying Bluesky
 context and worker"""
 
-
+LOGGER = logging.getLogger(__name__)
 _CONFIG: ApplicationConfig = ApplicationConfig()
 
 
@@ -226,7 +230,7 @@ def remove_callback_when_task_finished(
     if task.task_id is not None:
         try:
             active_worker.begin_task(task.task_id)
-        except KeyError:
+        except:
             for channel, token in subscribers:
                 channel.unsubscribe(token)
             raise
@@ -281,3 +285,37 @@ def get_python_env(
     """Retrieve information about the Python environment"""
     scratch = config().scratch
     return get_python_environment(config=scratch, name=name, source=source)
+
+
+@dataclass
+class SubHandles:
+    worker: int
+    progress: int
+    data: int
+
+
+def pipe_events(tx: Connection) -> SubHandles:
+    tw = worker()
+
+    def handler(
+        worker_event: WorkerEvent | DataEvent | ProgressEvent,
+        _cor_id: str | None,
+    ) -> None:
+
+        try:
+            tx.send(worker_event)
+        except BrokenPipeError:
+            LOGGER.warning("Sending event to broken pipe")
+            pass
+
+    w = tw.worker_events.subscribe(handler)
+    d = tw.data_events.subscribe(handler)
+    p = tw.progress_events.subscribe(handler)
+    return SubHandles(worker=w, data=d, progress=p)
+
+
+def unpipe_events(hnd: SubHandles) -> None:
+    tw = worker()
+    tw.worker_events.unsubscribe(hnd.worker)
+    tw.data_events.unsubscribe(hnd.data)
+    tw.progress_events.unsubscribe(hnd.progress)