From 108ef4f83f6e098b1d9e3dc0b15d865deaa6b534 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 17:20:02 +0000
Subject: [PATCH 01/60] wip

---
 src/blueapi/service/interface.py | 15 +++++++++++++
 src/blueapi/service/main.py      | 38 +++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index da262d33a..ae66177ec 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -1,5 +1,6 @@
 from collections.abc import Mapping
 from functools import cache
+from multiprocessing.connection import Connection
 from typing import Any
 
 from bluesky.callbacks.tiled_writer import TiledWriter
@@ -281,3 +282,17 @@ def get_python_env(
     """Retrieve information about the Python environment"""
     scratch = config().scratch
     return get_python_environment(config=scratch, name=name, source=source)
+
+
+def pipe_events(tx: Connection) -> int:
+
+    def handler(
+        worker_event: WorkerEvent,
+        cor_id: str | None,
+    ) -> None:
+        tx.send(worker_event)
+
+    task_worker = worker()
+    sub_id = task_worker.worker_events.subscribe(handler)
+
+    return sub_id
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 55b929b31..5ea458da0 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -2,6 +2,7 @@
 import urllib.parse
 from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
+from multiprocessing import Pipe
 from typing import Annotated
 
 import jwt
@@ -15,8 +16,10 @@
     HTTPException,
     Request,
     Response,
+    WebSocket,
     status,
 )
+from fastapi.concurrency import run_in_threadpool
 from fastapi.datastructures import Address
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse, StreamingResponse
@@ -38,7 +41,7 @@
     VersionHeaders,
 )
 from blueapi.worker import TrackableTask, WorkerState
-from blueapi.worker.event import TaskStatusEnum
+from blueapi.worker.event import TaskStatusEnum, WorkerEvent
 
 from .authorization import (
     OpaClient,
@@ -604,6 +607,39 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
     )
 
 
+@secure_router.websocket("/run_plan")
+async def run_plan(
+    ws: WebSocket,
+    runner: Annotated[WorkerDispatcher, Depends(_runner)],
+):
+    user = "alice"
+
+    # ack ws
+    await ws.accept()
+    # accept task request through socket
+    rq = await ws.receive_json()
+    # submit task to runner
+    task_request: TaskRequest = TaskRequest.model_validate(rq)
+    task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+    # add listener to runner
+    tx, rx = Pipe()
+    h = runner.run(interface.pipe_events, tx=tx)
+    # start task
+    task = WorkerTask(task_id=task_id)
+    runner.run(
+        interface.begin_task,
+        task=task,
+    )
+    # pipe events to ws
+    while True:
+        event: WorkerEvent = await run_in_threadpool(rx.recv)
+        await ws.send_json(event.model_dump(mode="json"))
+        if event.is_complete():
+            break
+    # ???
+    # profit
+
+
 @start_as_current_span(TRACER, "config")
 def start(config: ApplicationConfig):
     import uvicorn

From be1c4f1dda809ed2f34b4359c5fe4568c95f8240 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 17:37:08 +0000
Subject: [PATCH 02/60] client wip

---
 src/blueapi/cli/cli.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index 8c746fe92..fba636077 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -401,6 +401,29 @@ def on_event(event: AnyEvent) -> None:
         raise ClickException(f"task could not run: {ve}") from ve
 
 
+@controller.command(name="ws")
+@click.argument("name", type=str)
+@click.argument("parameters", type=ParametersType(), default={}, required=False)
+def run_blocking(
+    name: str,
+    parameters: TaskParameters,
+):
+    instrument_session = "cm33-3"
+
+    from websockets.sync.client import connect
+
+    task_req = TaskRequest(
+        name=name,
+        params=parameters,
+        instrument_session=instrument_session,
+    )
+
+    with connect("ws://localhost:8007/run_plan") as ws:
+        ws.send(task_req.model_dump_json())
+        while message := ws.recv():
+            print(message)
+
+
 @controller.command(name="state")
 @click.pass_obj
 @check_connection

From 344f8689d95a309008b4d08f6dcc0c1668597a64 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 17:48:31 +0000
Subject: [PATCH 03/60] use normal iter

---
 src/blueapi/cli/cli.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index fba636077..d6271f719 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -420,7 +420,7 @@ def run_blocking(
 
     with connect("ws://localhost:8007/run_plan") as ws:
         ws.send(task_req.model_dump_json())
-        while message := ws.recv():
+        for message in ws:
             print(message)
 
 

From 7f55bccd6a9a87da53b408e2d76bbce3a28ff109 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 17:48:48 +0000
Subject: [PATCH 04/60] close ws

---
 src/blueapi/service/main.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 5ea458da0..ac424336a 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -636,8 +636,7 @@ async def run_plan(
         await ws.send_json(event.model_dump(mode="json"))
         if event.is_complete():
             break
-    # ???
-    # profit
+    await ws.close()
 
 
 @start_as_current_span(TRACER, "config")

From b1d09a7c389a33a7047884f67644caab9dff7739 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 18:41:02 +0000
Subject: [PATCH 05/60] add some trys

---
 src/blueapi/service/main.py | 39 ++++++++++++++++++++++++-------------
 1 file changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index ac424336a..af243e1ab 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -42,6 +42,7 @@
 )
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import TaskStatusEnum, WorkerEvent
+from blueapi.worker.worker_errors import WorkerBusyError
 
 from .authorization import (
     OpaClient,
@@ -619,24 +620,36 @@ async def run_plan(
     # accept task request through socket
     rq = await ws.receive_json()
     # submit task to runner
-    task_request: TaskRequest = TaskRequest.model_validate(rq)
-    task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+    try:
+        task_request: TaskRequest = TaskRequest.model_validate(rq)
+        task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+    except ValidationError:
+        await ws.close(code=1003, reason="invalid args")
+        return
+
     # add listener to runner
     tx, rx = Pipe()
     h = runner.run(interface.pipe_events, tx=tx)
     # start task
-    task = WorkerTask(task_id=task_id)
-    runner.run(
-        interface.begin_task,
-        task=task,
-    )
+    try:
+        task = WorkerTask(task_id=task_id)
+        runner.run(
+            interface.begin_task,
+            task=task,
+        )
+    except WorkerBusyError:
+        await ws.close(code=1013, reason="Worker busy")
+        return
     # pipe events to ws
-    while True:
-        event: WorkerEvent = await run_in_threadpool(rx.recv)
-        await ws.send_json(event.model_dump(mode="json"))
-        if event.is_complete():
-            break
-    await ws.close()
+    try:
+        while True:
+            event: WorkerEvent = await run_in_threadpool(rx.recv)
+            await ws.send_json(event.model_dump(mode="json"))
+            if event.is_complete():
+                break
+    finally:
+        await ws.close()
+        runner.run(interface.unpipe_events, h=h)
 
 
 @start_as_current_span(TRACER, "config")

From 176487cb48037f93ff1d46816a556fab95ca1f06 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Tue, 24 Feb 2026 18:42:16 +0000
Subject: [PATCH 06/60] unpipe

---
 src/blueapi/service/interface.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index ae66177ec..958ac032e 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Mapping
 from functools import cache
 from multiprocessing.connection import Connection
@@ -23,6 +24,7 @@
     WorkerTask,
 )
 from blueapi.utils.serialization import access_blob
+from blueapi.worker import task_worker
 from blueapi.worker.event import TaskStatusEnum, WorkerEvent, WorkerState
 from blueapi.worker.task import Task
 from blueapi.worker.task_worker import TaskWorker, TrackableTask
@@ -30,7 +32,7 @@
 """This module provides interface between web application and underlying Bluesky
 context and worker"""
 
-
+LOGGER = logging.getLogger(__name__)
 _CONFIG: ApplicationConfig = ApplicationConfig()
 
 
@@ -290,9 +292,14 @@ def handler(
         worker_event: WorkerEvent,
         cor_id: str | None,
     ) -> None:
+        LOGGER.info("Sending event")
         tx.send(worker_event)
 
     task_worker = worker()
     sub_id = task_worker.worker_events.subscribe(handler)
-
     return sub_id
+
+
+def unpipe_events(h: int) -> None:
+    task_worker = worker()
+    task_worker.worker_events.unsubscribe(h)

From 8d15f23b9c62f3e6bc863c7a04c3ca038c0a5d3f Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 3 Mar 2026 11:52:21 +0000
Subject: [PATCH 07/60] Move websocket handling into BlueapiRestClient

---
 src/blueapi/cli/cli.py       | 25 ++++++++++++++-----------
 src/blueapi/client/client.py |  4 ++++
 src/blueapi/client/rest.py   | 17 ++++++++++++++++-
 3 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index d6271f719..463841d4c 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -402,26 +402,29 @@ def on_event(event: AnyEvent) -> None:
 
 
 @controller.command(name="ws")
+@click.pass_obj
 @click.argument("name", type=str)
 @click.argument("parameters", type=ParametersType(), default={}, required=False)
+@click.option(
+    "-i",
+    "--instrument-session",
+    type=str,
+    help=textwrap.dedent("""
+        Instrument session associated with running the plan,
+        used to tell blueapi where to store any data and as a security check:
+        the session must be valid and active and you must be a member of it."""),
+    required=True,
+)
 def run_blocking(
-    name: str,
-    parameters: TaskParameters,
+    obj: dict, name: str, parameters: TaskParameters, instrument_session: str
 ):
-    instrument_session = "cm33-3"
-
-    from websockets.sync.client import connect
-
+    client = cast(BlueapiClient, obj["client"])
     task_req = TaskRequest(
         name=name,
         params=parameters,
         instrument_session=instrument_session,
     )
-
-    with connect("ws://localhost:8007/run_plan") as ws:
-        ws.send(task_req.model_dump_json())
-        for message in ws:
-            print(message)
+    client.run_blocking(task_req)
 
 
 @controller.command(name="state")
diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 20de15892..918227e57 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -459,6 +459,10 @@ def get_active_task(self) -> WorkerTask:
 
         return self.active_task
 
+    @start_as_current_span(TRACER, "request")
+    def run_blocking(self, request: TaskRequest):
+        self._rest.run_blocking(request)
+
     @start_as_current_span(TRACER, "task", "timeout")
     def run_task(
         self,
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 0bddb5c87..ea5133217 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -10,8 +10,9 @@
     get_tracer,
     start_as_current_span,
 )
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError, WebsocketUrl
 from pydantic_core import PydanticSerializationError
+from websockets.sync.client import connect
 
 from blueapi import __version__
 from blueapi.client import client
@@ -340,6 +341,20 @@ def _request_and_deserialize(
         )
         return deserialized
 
+    def run_blocking(self, req: TaskRequest):
+        url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
+        print(url)
+        with connect(url) as ws:
+            ws.send(req.model_dump_json())
+            for message in ws:
+                print(message)
+
+    def _ws_address(self) -> WebsocketUrl:
+        # url = WebsocketUrl.build(
+        #     scheme="ws", host=api.host, port=api.port, path=api.path
+        # )
+        return WebsocketUrl("ws://localhost:8000/")
+
 
 # https://github.com/DiamondLightSource/blueapi/issues/1256 - remove before 2.0
 def __getattr__(name: str):

From 07078e2f3050d424592e21e6a631ac9c465c3b81 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 3 Mar 2026 12:16:10 +0000
Subject: [PATCH 08/60] Send all events through websocket

---
 src/blueapi/service/interface.py | 13 ++++++++-----
 src/blueapi/service/main.py      | 10 +++++++---
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index 958ac032e..4589e646a 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -11,6 +11,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
@@ -24,8 +25,7 @@
     WorkerTask,
 )
 from blueapi.utils.serialization import access_blob
-from blueapi.worker import task_worker
-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
 
@@ -289,17 +289,20 @@ def get_python_env(
 def pipe_events(tx: Connection) -> int:
 
     def handler(
-        worker_event: WorkerEvent,
-        cor_id: str | None,
+        worker_event: WorkerEvent | DataEvent | ProgressEvent,
+        _cor_id: str | None,
     ) -> None:
-        LOGGER.info("Sending event")
         tx.send(worker_event)
 
     task_worker = worker()
     sub_id = task_worker.worker_events.subscribe(handler)
+    sub_id = task_worker.data_events.subscribe(handler)
+    sub_id = task_worker.progress_events.subscribe(handler)
     return sub_id
 
 
 def unpipe_events(h: int) -> None:
     task_worker = worker()
     task_worker.worker_events.unsubscribe(h)
+    task_worker.data_events.unsubscribe(h)
+    task_worker.progress_events.unsubscribe(h)
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index af243e1ab..229878e88 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -34,6 +34,7 @@
 from starlette.responses import JSONResponse
 
 from blueapi.config import ApplicationConfig, OIDCConfig, Tag
+from blueapi.core.bluesky_types import DataEvent
 from blueapi.service import interface
 from blueapi.service.authentication import Fedid, build_access_token_check
 from blueapi.service.middleware import (
@@ -41,7 +42,7 @@
     VersionHeaders,
 )
 from blueapi.worker import TrackableTask, WorkerState
-from blueapi.worker.event import TaskStatusEnum, WorkerEvent
+from blueapi.worker.event import ProgressEvent, TaskStatusEnum, WorkerEvent
 from blueapi.worker.worker_errors import WorkerBusyError
 
 from .authorization import (
@@ -75,6 +76,9 @@
 TRACER = get_tracer("interface")
 
 
+AnyEvent = WorkerEvent | DataEvent | ProgressEvent
+
+
 def _runner() -> WorkerDispatcher:
     """Intended to be used only with FastAPI Depends"""
     if RUNNER is None:
@@ -643,9 +647,9 @@ async def run_plan(
     # pipe events to ws
     try:
         while True:
-            event: WorkerEvent = await run_in_threadpool(rx.recv)
+            event: AnyEvent = await run_in_threadpool(rx.recv)
             await ws.send_json(event.model_dump(mode="json"))
-            if event.is_complete():
+            if isinstance(event, WorkerEvent) and event.is_complete():
                 break
     finally:
         await ws.close()

From 01384acbb29018d912eaf61f1a704cc7002c497b Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 4 Mar 2026 14:50:32 +0000
Subject: [PATCH 09/60] Split pipe subscribe handles

---
 src/blueapi/service/interface.py | 22 +++++++++++++---------
 src/blueapi/service/main.py      |  2 +-
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index 4589e646a..07847c3b5 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -286,7 +286,10 @@ def get_python_env(
     return get_python_environment(config=scratch, name=name, source=source)
 
 
-def pipe_events(tx: Connection) -> int:
+SubHandle = tuple[int, int, int]
+
+
+def pipe_events(tx: Connection) -> SubHandle:
 
     def handler(
         worker_event: WorkerEvent | DataEvent | ProgressEvent,
@@ -295,14 +298,15 @@ def handler(
         tx.send(worker_event)
 
     task_worker = worker()
-    sub_id = task_worker.worker_events.subscribe(handler)
-    sub_id = task_worker.data_events.subscribe(handler)
-    sub_id = task_worker.progress_events.subscribe(handler)
-    return sub_id
+    w_id = task_worker.worker_events.subscribe(handler)
+    d_id = task_worker.data_events.subscribe(handler)
+    p_id = task_worker.progress_events.subscribe(handler)
+    return (w_id, d_id, p_id)
 
 
-def unpipe_events(h: int) -> None:
+def unpipe_events(hnd: SubHandle) -> None:
     task_worker = worker()
-    task_worker.worker_events.unsubscribe(h)
-    task_worker.data_events.unsubscribe(h)
-    task_worker.progress_events.unsubscribe(h)
+    w, d, p = hnd
+    task_worker.worker_events.unsubscribe(w)
+    task_worker.data_events.unsubscribe(d)
+    task_worker.progress_events.unsubscribe(p)
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 229878e88..fc756bee7 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -653,7 +653,7 @@ async def run_plan(
                 break
     finally:
         await ws.close()
-        runner.run(interface.unpipe_events, h=h)
+        runner.run(interface.unpipe_events, hnd=h)
 
 
 @start_as_current_span(TRACER, "config")

From d23a5c925b65313e181c22b473ee288b77a8ac06 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 4 Mar 2026 16:10:43 +0000
Subject: [PATCH 10/60] Re-use run subcommand for websockets

---
 src/blueapi/cli/cli.py       | 10 +++++++++-
 src/blueapi/client/client.py | 19 +++++++++++++++++--
 src/blueapi/client/rest.py   |  5 +++--
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index 463841d4c..ca8ab39d5 100644
--- a/src/blueapi/cli/cli.py
+++ b/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")
diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 918227e57..11b5900af 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -460,8 +460,23 @@ def get_active_task(self) -> WorkerTask:
         return self.active_task
 
     @start_as_current_span(TRACER, "request")
-    def run_blocking(self, request: TaskRequest):
-        self._rest.run_blocking(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():
+                if event.task_status is None:
+                    raise BlueskyRemoteControlError(
+                        "Server completed without task status"
+                    )
+                return event.task_status
 
     @start_as_current_span(TRACER, "task", "timeout")
     def run_task(
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index ea5133217..30770c5b0 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -16,6 +16,7 @@
 
 from blueapi import __version__
 from blueapi.client import client
+from blueapi.client.event_bus import AnyEvent
 from blueapi.config import RestConfig
 from blueapi.service.authentication import JWTAuth, SessionManager
 from blueapi.service.model import (
@@ -343,11 +344,11 @@ def _request_and_deserialize(
 
     def run_blocking(self, req: TaskRequest):
         url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
-        print(url)
         with connect(url) as ws:
             ws.send(req.model_dump_json())
             for message in ws:
-                print(message)
+                event = TypeAdapter(AnyEvent).validate_json(message)
+                yield event
 
     def _ws_address(self) -> WebsocketUrl:
         # url = WebsocketUrl.build(

From 787a81a5e8d9b803f41e46de6d9855aa0ae76e99 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Wed, 4 Mar 2026 16:40:04 +0000
Subject: [PATCH 11/60] Raise for connection closing pre plan completed

---
 src/blueapi/client/client.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 11b5900af..3d6e3bf6a 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -477,6 +477,7 @@ def run_blocking(
                         "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(

From 166582f0aabddb32767092b904d3cb42c2b3143c Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Wed, 4 Mar 2026 16:41:45 +0000
Subject: [PATCH 12/60] Remove run blocking from cli

---
 src/blueapi/cli/cli.py | 26 --------------------------
 1 file changed, 26 deletions(-)

diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py
index ca8ab39d5..6368da0ac 100644
--- a/src/blueapi/cli/cli.py
+++ b/src/blueapi/cli/cli.py
@@ -409,32 +409,6 @@ def on_event(event: AnyEvent) -> None:
         raise ClickException(f"task could not run: {ve}") from ve
 
 
-@controller.command(name="ws")
-@click.pass_obj
-@click.argument("name", type=str)
-@click.argument("parameters", type=ParametersType(), default={}, required=False)
-@click.option(
-    "-i",
-    "--instrument-session",
-    type=str,
-    help=textwrap.dedent("""
-        Instrument session associated with running the plan,
-        used to tell blueapi where to store any data and as a security check:
-        the session must be valid and active and you must be a member of it."""),
-    required=True,
-)
-def run_blocking(
-    obj: dict, name: str, parameters: TaskParameters, instrument_session: str
-):
-    client = cast(BlueapiClient, obj["client"])
-    task_req = TaskRequest(
-        name=name,
-        params=parameters,
-        instrument_session=instrument_session,
-    )
-    client.run_blocking(task_req)
-
-
 @controller.command(name="state")
 @click.pass_obj
 @check_connection

From f1fa7e3d5f4393c0d1a36f03d1adadbf49552398 Mon Sep 17 00:00:00 2001
From: Abigail Emery <abigail.emery@diamond.ac.uk>
Date: Wed, 4 Mar 2026 17:24:11 +0000
Subject: [PATCH 13/60] Catch plan key error in run_plan

---
 src/blueapi/service/main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index fc756bee7..24cefa181 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -630,6 +630,9 @@ async def run_plan(
     except ValidationError:
         await ws.close(code=1003, reason="invalid args")
         return
+    except KeyError:
+        await ws.close(code=1003, reason="unknown plan")
+        return
 
     # add listener to runner
     tx, rx = Pipe()

From 0df0c437a2c971f58c5b1d98012a6cc0b7d2c17d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 6 Mar 2026 15:45:13 +0000
Subject: [PATCH 14/60] Refactor event pipe handling into context manager and
 iterable

---
 src/blueapi/service/interface.py | 41 ++++++++++++--------
 src/blueapi/service/main.py      | 48 +++++++++++------------
 src/blueapi/service/runner.py    | 66 ++++++++++++++++++++++++++------
 3 files changed, 103 insertions(+), 52 deletions(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index 07847c3b5..ec6633f4d 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -1,5 +1,6 @@
 import logging
 from collections.abc import Mapping
+from dataclasses import dataclass
 from functools import cache
 from multiprocessing.connection import Connection
 from typing import Any
@@ -286,27 +287,35 @@ def get_python_env(
     return get_python_environment(config=scratch, name=name, source=source)
 
 
-SubHandle = tuple[int, int, int]
+@dataclass
+class SubHandles:
+    worker: int
+    progress: int
+    data: int
 
 
-def pipe_events(tx: Connection) -> SubHandle:
+def pipe_events(tx: Connection) -> SubHandles:
+    tw = worker()
 
     def handler(
         worker_event: WorkerEvent | DataEvent | ProgressEvent,
         _cor_id: str | None,
     ) -> None:
-        tx.send(worker_event)
 
-    task_worker = worker()
-    w_id = task_worker.worker_events.subscribe(handler)
-    d_id = task_worker.data_events.subscribe(handler)
-    p_id = task_worker.progress_events.subscribe(handler)
-    return (w_id, d_id, p_id)
-
-
-def unpipe_events(hnd: SubHandle) -> None:
-    task_worker = worker()
-    w, d, p = hnd
-    task_worker.worker_events.unsubscribe(w)
-    task_worker.data_events.unsubscribe(d)
-    task_worker.progress_events.unsubscribe(p)
+        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):
+    tw = worker()
+    tw.worker_events.unsubscribe(hnd.worker)
+    tw.data_events.unsubscribe(hnd.data)
+    tw.progress_events.unsubscribe(hnd.progress)
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 24cefa181..2625e7b9c 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -2,7 +2,6 @@
 import urllib.parse
 from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
-from multiprocessing import Pipe
 from typing import Annotated
 
 import jwt
@@ -17,9 +16,9 @@
     Request,
     Response,
     WebSocket,
+    WebSocketDisconnect,
     status,
 )
-from fastapi.concurrency import run_in_threadpool
 from fastapi.datastructures import Address
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse, StreamingResponse
@@ -619,44 +618,45 @@ async def run_plan(
 ):
     user = "alice"
 
-    # ack ws
+    LOGGER.info("Starting WS plan")
     await ws.accept()
-    # accept task request through socket
     rq = await ws.receive_json()
-    # submit task to runner
+    LOGGER.info("Raw request: %s", rq)
     try:
         task_request: TaskRequest = TaskRequest.model_validate(rq)
+        LOGGER.info("Plan request: %s", task_request)
         task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
+        LOGGER.info("Task ID: %s", task_id)
     except ValidationError:
+        LOGGER.error("Args not valid", exc_info=True)
         await ws.close(code=1003, reason="invalid args")
         return
     except KeyError:
+        LOGGER.error("Plan not found", exc_info=True)
         await ws.close(code=1003, reason="unknown plan")
         return
 
-    # add listener to runner
-    tx, rx = Pipe()
-    h = runner.run(interface.pipe_events, tx=tx)
-    # start task
     try:
-        task = WorkerTask(task_id=task_id)
-        runner.run(
-            interface.begin_task,
-            task=task,
-        )
+        with runner.event_pipe() as events:
+            LOGGER.info("Created event pipe")
+            runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
+            async for evt in events:
+                LOGGER.debug("Event: %s", evt)
+                await ws.send_json(evt.model_dump(mode="json"))
+                if isinstance(evt, WorkerEvent) and evt.is_complete():
+                    LOGGER.info("End of stream")
+                    break
     except WorkerBusyError:
+        LOGGER.error("Worker was busy")
         await ws.close(code=1013, reason="Worker busy")
-        return
-    # pipe events to ws
-    try:
-        while True:
-            event: AnyEvent = await run_in_threadpool(rx.recv)
-            await ws.send_json(event.model_dump(mode="json"))
-            if isinstance(event, WorkerEvent) and event.is_complete():
-                break
-    finally:
+    except WebSocketDisconnect:
+        LOGGER.info("Client disconnected")
+        runner.run(
+            interface.cancel_active_task, failure=True, reason="Client disconnected"
+        )
+    else:
+        LOGGER.info("Plan complete")
         await ws.close()
-        runner.run(interface.unpipe_events, hnd=h)
 
 
 @start_as_current_span(TRACER, "config")
diff --git a/src/blueapi/service/runner.py b/src/blueapi/service/runner.py
index 2b5a5f37f..6b34c6ad1 100644
--- a/src/blueapi/service/runner.py
+++ b/src/blueapi/service/runner.py
@@ -1,10 +1,12 @@
+import asyncio
 import inspect
 import logging
 import signal
 import uuid
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
 from importlib import import_module
 from multiprocessing import Pool, set_start_method
+from multiprocessing.connection import Connection, Pipe
 from multiprocessing.pool import Pool as PoolClass
 from typing import Any, ParamSpec, TypeVar
 
@@ -15,11 +17,13 @@
 )
 from opentelemetry.context import attach
 from opentelemetry.propagate import get_global_textmap
-from pydantic import TypeAdapter
 
 from blueapi.config import ApplicationConfig
-from blueapi.service.interface import setup, teardown
+from blueapi.core.bluesky_types import DataEvent
+from blueapi.service import interface
+from blueapi.service.interface import SubHandles, setup, teardown
 from blueapi.service.model import EnvironmentResponse
+from blueapi.worker.event import ProgressEvent, WorkerEvent
 
 # The default multiprocessing start method is fork
 set_start_method("spawn", force=True)
@@ -145,11 +149,57 @@ def run(
             kwargs,
         )
 
+    def event_pipe(self):
+        return EventPipe(self)
+
     @property
     def state(self) -> EnvironmentResponse:
         return self._state
 
 
+class EventStream:
+    def __init__(self, rx: Connection):
+        self._rx = rx
+
+    def __aiter__(self) -> AsyncIterator[WorkerEvent | DataEvent | ProgressEvent]:
+        return self
+
+    async def __anext__(self) -> WorkerEvent | DataEvent | ProgressEvent:
+        data_available = asyncio.Event()
+        asyncio.get_event_loop().add_reader(self._rx.fileno(), data_available.set)
+        try:
+            while not self._rx.poll():
+                await data_available.wait()
+                data_available.clear()
+            return self._rx.recv()
+        except BrokenPipeError:
+            raise StopAsyncIteration() from None
+        finally:
+            asyncio.get_event_loop().remove_reader(self._rx.fileno())
+
+
+class EventPipe:
+    runner: WorkerDispatcher
+    handles: list[tuple[SubHandles, Connection]]
+
+    def __init__(self, runner: WorkerDispatcher):
+        self.runner = runner
+        self.handles = []
+
+    def __enter__(self) -> EventStream:
+        tx, rx = Pipe()
+        hnd = self.runner.run(interface.pipe_events, tx)
+        LOGGER.debug("Subscribing new event pipe: %s", hnd)
+        self.handles.append((hnd, tx))
+        return EventStream(rx)
+
+    def __exit__(self, *exc):
+        hnd, conn = self.handles.pop()
+        LOGGER.debug("Unsubscribing event pipe: %s", hnd)
+        conn.close()
+        self.runner.run(interface.unpipe_events, hnd)
+
+
 class InvalidRunnerStateError(Exception):
     def __init__(self, message):
         super().__init__(message)
@@ -173,15 +223,7 @@ def import_and_run_function(
     func: Callable[..., T] = _validate_function(
         mod.__dict__.get(function_name, None), function_name
     )
-    value = func(*args, **kwargs)
-    return _valid_return(value, expected_type)
-
-
-def _valid_return(value: Any, expected_type: type[T] | None = None) -> T:
-    if expected_type is None:
-        return value
-    else:
-        return TypeAdapter(expected_type).validate_python(value)
+    return func(*args, **kwargs)
 
 
 def _validate_function(func: Any, function_name: str) -> Callable:

From ecbf7aa7868c5908f5413fe5ff300cf6b78a4ca5 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 10:37:12 +0000
Subject: [PATCH 15/60] Testing auth tokens

---
 src/blueapi/client/rest.py  |  8 +++++++-
 src/blueapi/service/main.py | 28 ++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 30770c5b0..283e08b22 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -344,7 +344,13 @@ def _request_and_deserialize(
 
     def run_blocking(self, req: TaskRequest):
         url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
-        with connect(url) as ws:
+        with connect(
+            url,
+            additional_headers={
+                "Cookie": "Authorization=Bearer cook",
+                "Authorization": "Bearer head",
+            },
+        ) as ws:
             ws.send(req.model_dump_json())
             for message in ws:
                 event = TypeAdapter(AnyEvent).validate_json(message)
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 2625e7b9c..b1d000192 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -10,8 +10,10 @@
     APIRouter,
     BackgroundTasks,
     Body,
+    Cookie,
     Depends,
     FastAPI,
+    Header,
     HTTPException,
     Request,
     Response,
@@ -118,6 +120,7 @@ async def inner(app: FastAPI):
     return inner
 
 
+ws_router = APIRouter()
 open_router = APIRouter()
 secure_router = APIRouter(deprecated=True)
 secure_router_v1 = APIRouter(prefix="/api/v1")
@@ -135,13 +138,16 @@ def get_app(config: ApplicationConfig):
         openapi_tags=ApplicationConfig.TAG_METADATA,
     )
     dependencies = []
+    ws_dependencies = []
     if config.oidc:
         dependencies.append(Depends(build_access_token_check(config.oidc)))
+        ws_dependencies.append(Depends(init_ws_auth(config.oidc)))
         app.swagger_ui_init_oauth = {
             "clientId": "NOT_SUPPORTED",
         }
     app.include_router(open_router)
     app.include_router(secure_router_v1, dependencies=dependencies)
+    app.include_router(ws_router, dependencies=ws_dependencies)
     app.include_router(secure_router, dependencies=dependencies)
     app.add_exception_handler(KeyError, on_key_error_404)
     app.add_exception_handler(jwt.PyJWTError, on_token_error_401)
@@ -191,6 +197,24 @@ async def start_task_permission(
     await access_task_permission(opa, task.task_id, fedid, runner)
 
 
+def init_ws_auth(oidc_config: OIDCConfig):
+    LOGGER.info("Creating ws auth dependency")
+
+    async def inner(
+        ws: WebSocket,
+        auth_header: str | None = Header(alias="authorization", default=None),
+        auth_cookie: str | None = Cookie(default=None, alias="Authorization"),
+    ):
+        print(auth_header)
+        print(auth_cookie)
+        print(ws.headers)
+        print(ws.cookies)
+        await ws.accept()
+        LOGGER.info("Authenticating websocket")
+
+    return inner
+
+
 async def on_key_error_404(_: Request, __: Exception):
     return JSONResponse(
         status_code=status.HTTP_404_NOT_FOUND,
@@ -611,7 +635,7 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
     )
 
 
-@secure_router.websocket("/run_plan")
+@ws_router.websocket("/run_plan")
 async def run_plan(
     ws: WebSocket,
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
@@ -619,7 +643,7 @@ async def run_plan(
     user = "alice"
 
     LOGGER.info("Starting WS plan")
-    await ws.accept()
+    # await ws.accept()
     rq = await ws.receive_json()
     LOGGER.info("Raw request: %s", rq)
     try:

From d491229e30efeefb5619213aab4a63c25ad17fc0 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 14:03:49 +0000
Subject: [PATCH 16/60] Re-use existing auth dependency for websocket endpoint

---
 src/blueapi/service/authentication.py |  9 ++++---
 src/blueapi/service/main.py           | 37 ++++++---------------------
 2 files changed, 14 insertions(+), 32 deletions(-)

diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py
index 6761256de..a5f3206f9 100644
--- a/src/blueapi/service/authentication.py
+++ b/src/blueapi/service/authentication.py
@@ -17,6 +17,9 @@
 import requests
 from fastapi import Depends, HTTPException, Request
 from fastapi.security.utils import get_authorization_scheme_param
+from fastapi.requests import HTTPConnection
+from fastapi.security import OAuth2AuthorizationCodeBearer
+from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import TypeAdapter
 from requests.auth import AuthBase
 from starlette.status import HTTP_401_UNAUTHORIZED
@@ -278,7 +281,7 @@ 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"""
     # This is an abridged version of the same feature of
     # OAuth2AuthorizationCodeBearer from fastapi. Replicating here prevents
@@ -303,7 +306,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 +329,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)
 
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index b1d000192..f9377b653 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -10,10 +10,8 @@
     APIRouter,
     BackgroundTasks,
     Body,
-    Cookie,
     Depends,
     FastAPI,
-    Header,
     HTTPException,
     Request,
     Response,
@@ -37,7 +35,10 @@
 from blueapi.config import ApplicationConfig, OIDCConfig, Tag
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service import interface
-from blueapi.service.authentication import Fedid, build_access_token_check
+from blueapi.service.authentication import (
+    Fedid,
+    build_access_token_check,
+)
 from blueapi.service.middleware import (
     ObservabilityContextPropagator,
     VersionHeaders,
@@ -120,7 +121,6 @@ async def inner(app: FastAPI):
     return inner
 
 
-ws_router = APIRouter()
 open_router = APIRouter()
 secure_router = APIRouter(deprecated=True)
 secure_router_v1 = APIRouter(prefix="/api/v1")
@@ -138,16 +138,13 @@ def get_app(config: ApplicationConfig):
         openapi_tags=ApplicationConfig.TAG_METADATA,
     )
     dependencies = []
-    ws_dependencies = []
     if config.oidc:
         dependencies.append(Depends(build_access_token_check(config.oidc)))
-        ws_dependencies.append(Depends(init_ws_auth(config.oidc)))
         app.swagger_ui_init_oauth = {
             "clientId": "NOT_SUPPORTED",
         }
     app.include_router(open_router)
     app.include_router(secure_router_v1, dependencies=dependencies)
-    app.include_router(ws_router, dependencies=ws_dependencies)
     app.include_router(secure_router, dependencies=dependencies)
     app.add_exception_handler(KeyError, on_key_error_404)
     app.add_exception_handler(jwt.PyJWTError, on_token_error_401)
@@ -197,24 +194,6 @@ async def start_task_permission(
     await access_task_permission(opa, task.task_id, fedid, runner)
 
 
-def init_ws_auth(oidc_config: OIDCConfig):
-    LOGGER.info("Creating ws auth dependency")
-
-    async def inner(
-        ws: WebSocket,
-        auth_header: str | None = Header(alias="authorization", default=None),
-        auth_cookie: str | None = Cookie(default=None, alias="Authorization"),
-    ):
-        print(auth_header)
-        print(auth_cookie)
-        print(ws.headers)
-        print(ws.cookies)
-        await ws.accept()
-        LOGGER.info("Authenticating websocket")
-
-    return inner
-
-
 async def on_key_error_404(_: Request, __: Exception):
     return JSONResponse(
         status_code=status.HTTP_404_NOT_FOUND,
@@ -635,15 +614,15 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
     )
 
 
-@ws_router.websocket("/run_plan")
+@secure_router.websocket("/run_plan")
 async def run_plan(
     ws: WebSocket,
     runner: Annotated[WorkerDispatcher, Depends(_runner)],
 ):
-    user = "alice"
+    user = ws.state.decoded_access_token["fedid"]
 
-    LOGGER.info("Starting WS plan")
-    # await ws.accept()
+    LOGGER.info("Starting WS plan as %s", user)
+    await ws.accept()
     rq = await ws.receive_json()
     LOGGER.info("Raw request: %s", rq)
     try:

From 860710de5e6e232d70491f57db5b5a76f49f8240 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 14:04:34 +0000
Subject: [PATCH 17/60] Add user auth token in websocket client

---
 src/blueapi/client/rest.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 283e08b22..6e8e2d5a2 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -344,13 +344,11 @@ def _request_and_deserialize(
 
     def run_blocking(self, req: TaskRequest):
         url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
-        with connect(
-            url,
-            additional_headers={
-                "Cookie": "Authorization=Bearer cook",
-                "Authorization": "Bearer head",
-            },
-        ) as ws:
+        headers = {}
+        if self._session_manager:
+            auth = self._session_manager.get_valid_access_token()
+            headers["Authorization"] = f"Bearer {auth}"
+        with connect(url, additional_headers=headers) as ws:
             ws.send(req.model_dump_json())
             for message in ws:
                 event = TypeAdapter(AnyEvent).validate_json(message)

From 309472e8973a8f94e59292c14d08597723071b5d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 16:39:56 +0000
Subject: [PATCH 18/60] Read authorization from cookie as well as header

---
 src/blueapi/service/authentication.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py
index a5f3206f9..edbf495d2 100644
--- a/src/blueapi/service/authentication.py
+++ b/src/blueapi/service/authentication.py
@@ -15,10 +15,8 @@
 import httpx
 import jwt
 import requests
-from fastapi import Depends, HTTPException, Request
-from fastapi.security.utils import get_authorization_scheme_param
+from fastapi import Depends, HTTPException
 from fastapi.requests import HTTPConnection
-from fastapi.security import OAuth2AuthorizationCodeBearer
 from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import TypeAdapter
 from requests.auth import AuthBase
@@ -288,7 +286,8 @@ def unchecked_bearer_token(req: HTTPConnection) -> str | None:
     # 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)
+    auth_cookie = req.cookies.get("Authorization")
+    scheme, param = get_authorization_scheme_param(auth or auth_cookie)
     if scheme.casefold() != "bearer":
         return None
     return param.strip()

From 4da76dd080db43a9df5554ab5b172dba61a6497e Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 16:41:32 +0000
Subject: [PATCH 19/60] Add user agent to websocket request

---
 src/blueapi/client/rest.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 6e8e2d5a2..de2ff6600 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -348,7 +348,11 @@ def run_blocking(self, req: TaskRequest):
         if self._session_manager:
             auth = self._session_manager.get_valid_access_token()
             headers["Authorization"] = f"Bearer {auth}"
-        with connect(url, additional_headers=headers) as ws:
+        with connect(
+            url,
+            additional_headers=headers,
+            user_agent_header="blueapi cli",
+        ) as ws:
             ws.send(req.model_dump_json())
             for message in ws:
                 event = TypeAdapter(AnyEvent).validate_json(message)

From abf70542c9d13c4c73b0e5c0855a9ce60f15070d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 7 Apr 2026 11:19:49 +0100
Subject: [PATCH 20/60] Add user agent to all requests

---
 src/blueapi/client/rest.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index de2ff6600..879d28ba4 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -41,6 +41,8 @@
 
 LOGGER = logging.getLogger(__name__)
 
+USER_AGENT = f"blueapi cli {__version__}"
+
 
 class BlueskyRequestError(Exception):
     """An error response from the blueapi server."""
@@ -309,14 +311,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:
@@ -344,14 +347,14 @@ def _request_and_deserialize(
 
     def run_blocking(self, req: TaskRequest):
         url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
-        headers = {}
+        headers = get_context_propagator()
         if self._session_manager:
             auth = self._session_manager.get_valid_access_token()
             headers["Authorization"] = f"Bearer {auth}"
         with connect(
             url,
             additional_headers=headers,
-            user_agent_header="blueapi cli",
+            user_agent_header=USER_AGENT,
         ) as ws:
             ws.send(req.model_dump_json())
             for message in ws:

From 117bbf90b0d427ec9a316ef743736be25a7c5b87 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 29 Jun 2026 18:38:43 +0100
Subject: [PATCH 21/60] Use new fedid dependency for user name

---
 src/blueapi/service/main.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index f9377b653..d8206a6ba 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -616,11 +616,8 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
 
 @secure_router.websocket("/run_plan")
 async def run_plan(
-    ws: WebSocket,
-    runner: Annotated[WorkerDispatcher, Depends(_runner)],
+    ws: WebSocket, runner: Annotated[WorkerDispatcher, Depends(_runner)], user: Fedid
 ):
-    user = ws.state.decoded_access_token["fedid"]
-
     LOGGER.info("Starting WS plan as %s", user)
     await ws.accept()
     rq = await ws.receive_json()

From 42ff1ec8b84146f4afa6c0baee7ef953ec5a15e8 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 29 Jun 2026 18:46:41 +0100
Subject: [PATCH 22/60] Test auth from cookie

---
 .../unit_tests/service/test_authentication.py | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/tests/unit_tests/service/test_authentication.py b/tests/unit_tests/service/test_authentication.py
index 01bc426e2..e3ee86f9d 100644
--- a/tests/unit_tests/service/test_authentication.py
+++ b/tests/unit_tests/service/test_authentication.py
@@ -189,18 +189,25 @@ def test_tiled_auth_sync_auth_flow():
 
 
 @pytest.mark.parametrize(
-    "header,token",
+    "header,cookie,token",
     [
-        (None, None),
-        ("ApiKey foobar", None),
-        ("Bearer foobar", "foobar"),
-        ("Bearer   with_whitespace   ", "with_whitespace"),
-        ("Bearerfoobar", None),
+        (None, None, None),
+        ("", None, None),
+        ("ApiKey foobar", None, None),
+        ("Bearer foobar", None, "foobar"),
+        ("Bearer   with_whitespace   ", None, "with_whitespace"),
+        ("Bearerfoobar", None, None),
+        (None, "Bearer foobar", "foobar"),
+        ("", "Bearer foo", "foo"),
+        ("Bearer foo", "bearer bar", "foo"),
     ],
 )
-def test_unchecked_bearer_token(header: str | None, token: str | None):
+def test_unchecked_bearer_token(
+    header: str | None, cookie: str | None, token: str | None
+):
     req = Mock()
     req.headers.get.side_effect = lambda key: header if key == "Authorization" else None
+    req.cookies.get.side_effect = lambda key: cookie if key == "Authorization" else None
 
     assert unchecked_bearer_token(req) == token
 

From c093afe0ba62766ec6f82cda988c3f5beeb37048 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 29 Jun 2026 18:54:01 +0100
Subject: [PATCH 23/60] Fix CLI event handler test

---
 tests/unit_tests/cli/test_cli.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py
index 3e27cfc98..fe71825c1 100644
--- a/tests/unit_tests/cli/test_cli.py
+++ b/tests/unit_tests/cli/test_cli.py
@@ -7,7 +7,6 @@
 from pathlib import Path
 from textwrap import dedent
 from typing import Any, TypeVar
-from unittest import mock
 from unittest.mock import Mock, patch
 
 import pytest
@@ -385,9 +384,9 @@ def test_run_plan_feedback(
         main,
         ["controller", "run", "-i", "cm12345-1", "name"],
     )
+    bc.add_callback.assert_called_once()
     bc.run_task.assert_called_once_with(
         TaskRequest(name="name", params={}, instrument_session="cm12345-1"),
-        on_event=mock.ANY,
     )
     assert res.exit_code == 0
     assert res.stdout == message

From 373e08dd28627d4ea2cb44e5ea2013eda877161d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 29 Jun 2026 19:01:33 +0100
Subject: [PATCH 24/60] Reinstate _valid_return check

Should still investigate if this is still useful but not in this change
---
 src/blueapi/service/runner.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/runner.py b/src/blueapi/service/runner.py
index 6b34c6ad1..83153ed5f 100644
--- a/src/blueapi/service/runner.py
+++ b/src/blueapi/service/runner.py
@@ -17,6 +17,7 @@
 )
 from opentelemetry.context import attach
 from opentelemetry.propagate import get_global_textmap
+from pydantic import TypeAdapter
 
 from blueapi.config import ApplicationConfig
 from blueapi.core.bluesky_types import DataEvent
@@ -223,7 +224,15 @@ def import_and_run_function(
     func: Callable[..., T] = _validate_function(
         mod.__dict__.get(function_name, None), function_name
     )
-    return func(*args, **kwargs)
+    value = func(*args, **kwargs)
+    return _valid_return(value, expected_type)
+
+
+def _valid_return(value: Any, expected_type: type[T] | None = None) -> T:
+    if expected_type is None:
+        return value
+    else:
+        return TypeAdapter(expected_type).validate_python(value)
 
 
 def _validate_function(func: Any, function_name: str) -> Callable:

From 3118ea51e46be178ca1e277cc69b6d8b94062822 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 29 Jun 2026 19:16:22 +0100
Subject: [PATCH 25/60] Use versioned api for websockets

---
 src/blueapi/service/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index d8206a6ba..8f5f5a1e6 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -614,7 +614,7 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
     )
 
 
-@secure_router.websocket("/run_plan")
+@secure_router_v1.websocket("/run_plan")
 async def run_plan(
     ws: WebSocket, runner: Annotated[WorkerDispatcher, Depends(_runner)], user: Fedid
 ):

From d9631eadaec713e2ca77dc2c04c32f606a69c3b2 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 30 Jun 2026 16:04:10 +0100
Subject: [PATCH 26/60] Add type annotation to unpipe

---
 src/blueapi/service/interface.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index ec6633f4d..1bd6e7f4a 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -314,7 +314,7 @@ def handler(
     return SubHandles(worker=w, data=d, progress=p)
 
 
-def unpipe_events(hnd: SubHandles):
+def unpipe_events(hnd: SubHandles) -> None:
     tw = worker()
     tw.worker_events.unsubscribe(hnd.worker)
     tw.data_events.unsubscribe(hnd.data)

From d881acd4ae76d96f83fe3bd194580cdf6a62877e Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 30 Jun 2026 16:11:24 +0100
Subject: [PATCH 27/60] Move ws endpoint to v2 api

---
 src/blueapi/client/rest.py  | 2 +-
 src/blueapi/service/main.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 879d28ba4..e9d313e9a 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -346,7 +346,7 @@ def _request_and_deserialize(
         return deserialized
 
     def run_blocking(self, req: TaskRequest):
-        url = self._ws_address().unicode_string().removesuffix("/") + "/run_plan"
+        url = self._ws_address().unicode_string().removesuffix("/") + "/api/v2/run_plan"
         headers = get_context_propagator()
         if self._session_manager:
             auth = self._session_manager.get_valid_access_token()
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 8f5f5a1e6..bf81f9740 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -124,6 +124,7 @@ async def inner(app: FastAPI):
 open_router = APIRouter()
 secure_router = APIRouter(deprecated=True)
 secure_router_v1 = APIRouter(prefix="/api/v1")
+secure_router_v2 = APIRouter(prefix="/api/v2")
 
 
 def get_app(config: ApplicationConfig):
@@ -145,6 +146,7 @@ def get_app(config: ApplicationConfig):
         }
     app.include_router(open_router)
     app.include_router(secure_router_v1, dependencies=dependencies)
+    app.include_router(secure_router_v2, dependencies=dependencies)
     app.include_router(secure_router, dependencies=dependencies)
     app.add_exception_handler(KeyError, on_key_error_404)
     app.add_exception_handler(jwt.PyJWTError, on_token_error_401)
@@ -614,7 +616,7 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
     )
 
 
-@secure_router_v1.websocket("/run_plan")
+@secure_router_v2.websocket("/run_plan")
 async def run_plan(
     ws: WebSocket, runner: Annotated[WorkerDispatcher, Depends(_runner)], user: Fedid
 ):

From 4d03ec3ed78cbcc66ea6a9b5dd58431e204aad10 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 12 Mar 2026 16:39:56 +0000
Subject: [PATCH 28/60] Use Depends for header and cookie

---
 src/blueapi/service/authentication.py           | 11 ++++++-----
 tests/unit_tests/service/test_authentication.py |  6 +-----
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py
index edbf495d2..614f42b7c 100644
--- a/src/blueapi/service/authentication.py
+++ b/src/blueapi/service/authentication.py
@@ -15,7 +15,7 @@
 import httpx
 import jwt
 import requests
-from fastapi import Depends, HTTPException
+from fastapi import Cookie, Depends, Header, HTTPException
 from fastapi.requests import HTTPConnection
 from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import TypeAdapter
@@ -279,15 +279,16 @@ def sync_auth_flow(self, request):
         yield request
 
 
-def unchecked_bearer_token(req: HTTPConnection) -> str | None:
+def unchecked_bearer_token(
+    auth_header: str | None = Header(alias="Authorization", default=None),
+    auth_cookie: str | None = Cookie(alias="Authorization", default=None),
+) -> str | None:
     """Get bearer token value from authorization header"""
     # 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")
-    auth_cookie = req.cookies.get("Authorization")
-    scheme, param = get_authorization_scheme_param(auth or auth_cookie)
+    scheme, param = get_authorization_scheme_param(auth_header or auth_cookie)
     if scheme.casefold() != "bearer":
         return None
     return param.strip()
diff --git a/tests/unit_tests/service/test_authentication.py b/tests/unit_tests/service/test_authentication.py
index e3ee86f9d..cfc277ee2 100644
--- a/tests/unit_tests/service/test_authentication.py
+++ b/tests/unit_tests/service/test_authentication.py
@@ -205,11 +205,7 @@ def test_tiled_auth_sync_auth_flow():
 def test_unchecked_bearer_token(
     header: str | None, cookie: str | None, token: str | None
 ):
-    req = Mock()
-    req.headers.get.side_effect = lambda key: header if key == "Authorization" else None
-    req.cookies.get.side_effect = lambda key: cookie if key == "Authorization" else None
-
-    assert unchecked_bearer_token(req) == token
+    assert unchecked_bearer_token(header, cookie) == token
 
 
 def test_access_token():

From 7d1dd180651d7b6934926578d6813dfba833444e Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 17 Apr 2026 16:15:47 +0100
Subject: [PATCH 29/60] Add sub-protocol to ws communication

Allows error messages to be propagated correctly.
---
 src/blueapi/client/rest.py                | 52 ++++++++++---
 src/blueapi/service/main.py               | 38 ++++++---
 src/blueapi/service/protocol.py           | 94 +++++++++++++++++++++++
 tests/unit_tests/service/test_protocol.py | 70 +++++++++++++++++
 4 files changed, 232 insertions(+), 22 deletions(-)
 create mode 100644 src/blueapi/service/protocol.py
 create mode 100644 tests/unit_tests/service/test_protocol.py

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index e9d313e9a..6e5d86593 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -12,11 +12,11 @@
 )
 from pydantic import BaseModel, TypeAdapter, ValidationError, WebsocketUrl
 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.client.event_bus import AnyEvent
 from blueapi.config import RestConfig
 from blueapi.service.authentication import JWTAuth, SessionManager
 from blueapi.service.model import (
@@ -33,6 +33,13 @@
     TasksListResponse,
     WorkerTask,
 )
+from blueapi.service.protocol import (
+    ControlResponse,
+    InvalidArgs,
+    PlanNotFound,
+    Submit,
+    Update,
+)
 from blueapi.worker import TrackableTask, WorkerState
 
 T = TypeVar("T")
@@ -90,8 +97,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):
@@ -351,15 +358,36 @@ def run_blocking(self, req: TaskRequest):
         if self._session_manager:
             auth = self._session_manager.get_valid_access_token()
             headers["Authorization"] = f"Bearer {auth}"
-        with connect(
-            url,
-            additional_headers=headers,
-            user_agent_header=USER_AGENT,
-        ) as ws:
-            ws.send(req.model_dump_json())
-            for message in ws:
-                event = TypeAdapter(AnyEvent).validate_json(message)
-                yield event
+        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(name)
+                    yield event
+        except InvalidStatus as istat:
+            match istat.response.status_code:
+                case 401 | 403:
+                    raise UnauthorisedAccessError() from None
+            print(vars(istat))
+            return
 
     def _ws_address(self) -> WebsocketUrl:
         # url = WebsocketUrl.build(
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index bf81f9740..af111abc4 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -43,6 +43,13 @@
     ObservabilityContextPropagator,
     VersionHeaders,
 )
+from blueapi.service.protocol import (
+    ControlRequest,
+    InvalidArgs,
+    PlanNotFound,
+    ServerBusy,
+    Update,
+)
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import ProgressEvent, TaskStatusEnum, WorkerEvent
 from blueapi.worker.worker_errors import WorkerBusyError
@@ -622,20 +629,30 @@ async def run_plan(
 ):
     LOGGER.info("Starting WS plan as %s", user)
     await ws.accept()
-    rq = await ws.receive_json()
+    rq = await ws.receive_text()
     LOGGER.info("Raw request: %s", rq)
     try:
-        task_request: TaskRequest = TaskRequest.model_validate(rq)
-        LOGGER.info("Plan request: %s", task_request)
-        task_id: str = runner.run(interface.submit_task, task_request, {"user": user})
-        LOGGER.info("Task ID: %s", task_id)
+        task_request = ControlRequest.validate_json(rq)
     except ValidationError:
+        LOGGER.error("Failed to deserialize request", exc_info=True)
+        await ws.close(code=1007, reason="Invalid Request")
+        return
+    LOGGER.info("Plan request: %s", task_request)
+
+    try:
+        task_id: str = runner.run(
+            interface.submit_task, task_request.task, {"user": user}
+        )
+        LOGGER.info("Task ID: %s", task_id)
+    except ValidationError as ve:
         LOGGER.error("Args not valid", exc_info=True)
-        await ws.close(code=1003, reason="invalid args")
+        await ws.send_text(InvalidArgs.from_validation_error(ve).model_dump_json())
+        await ws.close(code=4002, reason="Invalid Args")
         return
-    except KeyError:
-        LOGGER.error("Plan not found", exc_info=True)
-        await ws.close(code=1003, reason="unknown plan")
+    except KeyError as ke:
+        LOGGER.error("Plan %r not found", ke.args[0])
+        await ws.send_text(PlanNotFound(plan_name=ke.args[0]).model_dump_json())
+        await ws.close(code=4001, reason="unknown plan")
         return
 
     try:
@@ -644,12 +661,13 @@ async def run_plan(
             runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
             async for evt in events:
                 LOGGER.debug("Event: %s", evt)
-                await ws.send_json(evt.model_dump(mode="json"))
+                await ws.send_json(Update(data=evt).model_dump(mode="json"))
                 if isinstance(evt, WorkerEvent) and evt.is_complete():
                     LOGGER.info("End of stream")
                     break
     except WorkerBusyError:
         LOGGER.error("Worker was busy")
+        await ws.send_json(ServerBusy())
         await ws.close(code=1013, reason="Worker busy")
     except WebSocketDisconnect:
         LOGGER.info("Client disconnected")
diff --git a/src/blueapi/service/protocol.py b/src/blueapi/service/protocol.py
new file mode 100644
index 000000000..67924318a
--- /dev/null
+++ b/src/blueapi/service/protocol.py
@@ -0,0 +1,94 @@
+"""
+The application level sub-protocol used to communicate between the server and
+client when running plans via websockets
+"""
+
+# Client to server
+# * Submit task
+# * Pause
+# * Resume
+# * Abort
+#
+# Server to client
+# * Plan not found
+# * Args not valid
+# * Server busy
+# * Event update
+
+from typing import Annotated, Any, Literal, Self
+
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+
+from blueapi.core.bluesky_types import DataEvent
+from blueapi.service.model import TaskRequest
+from blueapi.worker.event import ProgressEvent, WorkerEvent
+
+
+class ArgumentError(BaseModel):
+    loc: list[str | int]
+    msg: str | None
+    type: str | None
+    input: Any
+
+
+class Submit(BaseModel):
+    kind: Literal["submit"] = "submit"
+    task: TaskRequest
+
+
+class Pause(BaseModel):
+    kind: Literal["pause"] = "pause"
+
+
+class Resume(BaseModel):
+    kind: Literal["resume"] = "resume"
+
+
+class Abort(BaseModel):
+    kind: Literal["abort"] = "abort"
+    reason: str | None = None
+
+
+ControlRequest = TypeAdapter(
+    Annotated[Submit | Pause | Resume | Abort, Field(discriminator="kind")]
+)
+
+
+class PlanNotFound(BaseModel):
+    kind: Literal["plan_not_found"] = "plan_not_found"
+    plan_name: str
+
+
+class InvalidArgs(BaseModel):
+    kind: Literal["invalid_args"] = "invalid_args"
+    errors: list[ArgumentError]
+
+    @classmethod
+    def from_validation_error(cls, e: ValidationError) -> Self:
+        errors = [
+            ArgumentError(
+                loc=["body", "params", *err.get("loc", [])],
+                msg=err.get("msg", None),
+                type=err.get("type", None),
+                # Input is not listed as required but is useful to have if available
+                input=err.get("input", None),
+            )
+            for err in e.errors()
+        ]
+        return cls(errors=errors)
+
+
+class ServerBusy(BaseModel):
+    kind: Literal["busy"] = "busy"
+
+
+class Update(BaseModel):
+    kind: Literal["update"] = "update"
+    data: WorkerEvent | DataEvent | ProgressEvent
+
+
+ControlResponse = TypeAdapter(
+    Annotated[
+        PlanNotFound | InvalidArgs | ServerBusy | Update, Field(discriminator="kind")
+    ]
+)
diff --git a/tests/unit_tests/service/test_protocol.py b/tests/unit_tests/service/test_protocol.py
new file mode 100644
index 000000000..ef95bb8ce
--- /dev/null
+++ b/tests/unit_tests/service/test_protocol.py
@@ -0,0 +1,70 @@
+from typing import Any
+
+import pytest
+
+from blueapi.service.model import TaskRequest
+from blueapi.service.protocol import (
+    Abort,
+    ArgumentError,
+    ControlRequest,
+    ControlResponse,
+    InvalidArgs,
+    Pause,
+    Resume,
+    Submit,
+)
+
+
+@pytest.mark.parametrize(
+    "src,res",
+    [
+        (
+            """{
+                "kind": "submit",
+                "task": {
+                    "name": "foo",
+                    "instrument_session": "cm12345-1"
+                }
+            }""",
+            Submit(
+                task=TaskRequest(name="foo", params={}, instrument_session="cm12345-1")
+            ),
+        ),
+        ('{"kind": "pause"}', Pause()),
+        ('{"kind": "resume"}', Resume()),
+        ('{"kind": "abort"}', Abort()),
+    ],
+)
+def test_request_deserialization(src: str, res: Any):
+    req = ControlRequest.validate_json(src)
+    assert req == res
+
+
+@pytest.mark.parametrize(
+    "src,res",
+    [
+        (
+            """{
+                "kind": "invalid_args",
+                "errors":[{
+                    "loc":["body","params","spec"],
+                     "msg":"error_message",
+                     "type":"error_type",
+                     "input":"original input"
+                 }]}""",
+            InvalidArgs(
+                errors=[
+                    ArgumentError(
+                        loc=["body", "params", "spec"],
+                        msg="error_message",
+                        type="error_type",
+                        input="original input",
+                    )
+                ]
+            ),
+        ),
+    ],
+)
+def test_response_deserialization(src: str, res: Any):
+    req = ControlResponse.validate_json(src)
+    assert req == res

From ceba3f7e992d1e570558d4496a43d8990ec25b8d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 17 Apr 2026 16:16:27 +0100
Subject: [PATCH 30/60] Used configured host for websockets

---
 src/blueapi/client/rest.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 6e5d86593..2fa2a3c21 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -390,10 +390,13 @@ def run_blocking(self, req: TaskRequest):
             return
 
     def _ws_address(self) -> WebsocketUrl:
-        # url = WebsocketUrl.build(
-        #     scheme="ws", host=api.host, port=api.port, path=api.path
-        # )
-        return WebsocketUrl("ws://localhost:8000/")
+        api = self._config.url
+        if api.host is None:
+            raise ValueError("No host configured")
+        scheme = "ws" if api.scheme == "http" else "wss"
+        return WebsocketUrl.build(
+            scheme=scheme, host=api.host, port=api.port, path=api.path
+        )
 
 
 # https://github.com/DiamondLightSource/blueapi/issues/1256 - remove before 2.0

From 0c9aab05a455b01183cd740830d96e9dfc7e161c Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 20 Apr 2026 17:03:27 +0100
Subject: [PATCH 31/60] Add debug logging of all websocket traffic

---
 src/blueapi/service/main.py       |  3 +-
 src/blueapi/service/middleware.py | 46 +++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index af111abc4..bec8e3f69 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -42,6 +42,7 @@
 from blueapi.service.middleware import (
     ObservabilityContextPropagator,
     VersionHeaders,
+    WebsocketTracing,
 )
 from blueapi.service.protocol import (
     ControlRequest,
@@ -160,6 +161,7 @@ def get_app(config: ApplicationConfig):
 
     app.add_middleware(ObservabilityContextPropagator)
     app.add_middleware(VersionHeaders)
+    app.add_middleware(WebsocketTracing)
     app.middleware("http")(log_request_details)
     if config.api.cors:
         app.add_middleware(
@@ -630,7 +632,6 @@ async def run_plan(
     LOGGER.info("Starting WS plan as %s", user)
     await ws.accept()
     rq = await ws.receive_text()
-    LOGGER.info("Raw request: %s", rq)
     try:
         task_request = ControlRequest.validate_json(rq)
     except ValidationError:
diff --git a/src/blueapi/service/middleware.py b/src/blueapi/service/middleware.py
index b31fe0fb9..a8cec77d6 100644
--- a/src/blueapi/service/middleware.py
+++ b/src/blueapi/service/middleware.py
@@ -8,6 +8,7 @@
 from blueapi.config import ApplicationConfig
 
 OBS_LOGGER = logging.getLogger("blueapi.service.middleware.observability")
+WS_LOGGER = logging.getLogger("blueapi.service.middleware.websocket")
 
 CONTEXT_HEADER = ApplicationConfig.CONTEXT_HEADER.encode()
 VENDOR_CONTEXT_HEADER = ApplicationConfig.VENDOR_CONTEXT_HEADER.encode()
@@ -56,3 +57,48 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
             attach(get_global_textmap().extract(carrier))
 
         return await self.app(scope, receive, send)
+
+
+class WebsocketTracing:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        active = WS_LOGGER.isEnabledFor(logging.DEBUG)
+
+        if scope.get("type") != "websocket" or not active:
+            return await self.app(scope, receive, send)
+
+        async def local_send(msg: Message):
+            match msg.get("type"):
+                case "websocket.send":
+                    WS_LOGGER.debug("Sending: %r", msg.get("text"))
+                case "websocket.accept":
+                    WS_LOGGER.debug(
+                        "Accepting websocket - sending headers: %r", msg.get("headers")
+                    )
+                case "websocket.close":
+                    WS_LOGGER.debug(
+                        "Closing with code: %r, reason: %r",
+                        msg.get("code"),
+                        msg.get("reason"),
+                    )
+                case "websocket.http.response.start":
+                    WS_LOGGER.debug(
+                        "HTTP Response: status=%r, headers=%r",
+                        msg.get("status"),
+                        msg.get("headers"),
+                    )
+                case "websocket.http.response.body":
+                    WS_LOGGER.debug("HTTP Response Content: %r", msg.get("body"))
+                case _:
+                    WS_LOGGER.debug("Sending other: %r", msg)
+
+            await send(msg)
+
+        async def local_receive() -> Message:
+            message = await receive()
+            WS_LOGGER.debug("Received: %r", message)
+            return message
+
+        return await self.app(scope, local_receive, local_send)

From d795d6cc7b7bb30ec8fc5c93c197e6dc5c0bcea4 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 20 Apr 2026 17:26:46 +0100
Subject: [PATCH 32/60] Include connection info in logging

---
 src/blueapi/service/middleware.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/blueapi/service/middleware.py b/src/blueapi/service/middleware.py
index a8cec77d6..493479033 100644
--- a/src/blueapi/service/middleware.py
+++ b/src/blueapi/service/middleware.py
@@ -1,4 +1,5 @@
 import logging
+import uuid
 
 from opentelemetry.context import attach
 from opentelemetry.propagate import get_global_textmap
@@ -69,30 +70,42 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         if scope.get("type") != "websocket" or not active:
             return await self.app(scope, receive, send)
 
+        conn_id = uuid.uuid4()
+        client: tuple[str, int] = scope.get("client", ("unknown", 0))
+        extra = {"conn": conn_id, "client": client}
+
+        WS_LOGGER.debug("%r", scope, extra=extra)
+
         async def local_send(msg: Message):
             match msg.get("type"):
                 case "websocket.send":
-                    WS_LOGGER.debug("Sending: %r", msg.get("text"))
+                    WS_LOGGER.debug("Sending: %r", msg.get("text"), extra=extra)
                 case "websocket.accept":
                     WS_LOGGER.debug(
-                        "Accepting websocket - sending headers: %r", msg.get("headers")
+                        "Accepting websocket - sending headers: %r",
+                        msg.get("headers"),
+                        extra=extra,
                     )
                 case "websocket.close":
                     WS_LOGGER.debug(
                         "Closing with code: %r, reason: %r",
                         msg.get("code"),
                         msg.get("reason"),
+                        extra=extra,
                     )
                 case "websocket.http.response.start":
                     WS_LOGGER.debug(
                         "HTTP Response: status=%r, headers=%r",
                         msg.get("status"),
                         msg.get("headers"),
+                        extra=extra,
                     )
                 case "websocket.http.response.body":
-                    WS_LOGGER.debug("HTTP Response Content: %r", msg.get("body"))
+                    WS_LOGGER.debug(
+                        "HTTP Response Content: %r", msg.get("body"), extra=extra
+                    )
                 case _:
-                    WS_LOGGER.debug("Sending other: %r", msg)
+                    WS_LOGGER.debug("Sending other: %r", msg, extra=extra)
 
             await send(msg)
 

From 87aeb4ce3451d7d45005c8cd7d60ea2f08e8a3da Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 20 Apr 2026 17:27:12 +0100
Subject: [PATCH 33/60] Split receive logging by type

---
 src/blueapi/service/middleware.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/middleware.py b/src/blueapi/service/middleware.py
index 493479033..a1177998f 100644
--- a/src/blueapi/service/middleware.py
+++ b/src/blueapi/service/middleware.py
@@ -111,7 +111,11 @@ async def local_send(msg: Message):
 
         async def local_receive() -> Message:
             message = await receive()
-            WS_LOGGER.debug("Received: %r", message)
+            match message.get("type"):
+                case "websocket.receive":
+                    WS_LOGGER.debug("Received: %r", message.get("text"))
+                case "websocket.connect":
+                    WS_LOGGER.debug("New connection from %s:%d", *client)
             return message
 
         return await self.app(scope, local_receive, local_send)

From 918347fe618bfebb682c99de54f5a92e824c1133 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 29 Apr 2026 11:19:41 +0100
Subject: [PATCH 34/60] Correct typing in rest run_blocking

---
 src/blueapi/client/rest.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 2fa2a3c21..75e900524 100644
--- a/src/blueapi/client/rest.py
+++ b/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
@@ -18,6 +18,7 @@
 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,
@@ -41,6 +42,7 @@
     Update,
 )
 from blueapi.worker import TrackableTask, WorkerState
+from blueapi.worker.event import ProgressEvent, WorkerEvent
 
 T = TypeVar("T")
 
@@ -352,7 +354,9 @@ def _request_and_deserialize(
         )
         return deserialized
 
-    def run_blocking(self, req: TaskRequest):
+    def run_blocking(
+        self, req: TaskRequest
+    ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]:
         url = self._ws_address().unicode_string().removesuffix("/") + "/api/v2/run_plan"
         headers = get_context_propagator()
         if self._session_manager:
@@ -381,7 +385,6 @@ def run_blocking(self, req: TaskRequest):
                             )
                         case PlanNotFound(plan_name=name):
                             raise UnknownPlanError(name)
-                    yield event
         except InvalidStatus as istat:
             match istat.response.status_code:
                 case 401 | 403:

From 2cf666a4fff30dd905e80f1f793a945e43bd770d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 1 Jul 2026 16:58:22 +0100
Subject: [PATCH 35/60] Use rstrip instead of removesuffix to remove multiple
 trailing slashes

---
 src/blueapi/client/rest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 75e900524..4d12cf4a7 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -357,7 +357,7 @@ def _request_and_deserialize(
     def run_blocking(
         self, req: TaskRequest
     ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]:
-        url = self._ws_address().unicode_string().removesuffix("/") + "/api/v2/run_plan"
+        url = self._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()

From 7651f0203d36de05efdd79b50012f029718895a8 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 1 Jul 2026 17:47:30 +0100
Subject: [PATCH 36/60] Redact auth tokens in websocket logging

---
 src/blueapi/service/middleware.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/middleware.py b/src/blueapi/service/middleware.py
index a1177998f..cb94747c5 100644
--- a/src/blueapi/service/middleware.py
+++ b/src/blueapi/service/middleware.py
@@ -1,5 +1,6 @@
 import logging
 import uuid
+from collections.abc import Iterable
 
 from opentelemetry.context import attach
 from opentelemetry.propagate import get_global_textmap
@@ -60,6 +61,17 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         return await self.app(scope, receive, send)
 
 
+Header = tuple[bytes, bytes]
+
+
+def _redact_headers(headers: list[Header] | None) -> Iterable[Header]:
+    for key, value in headers or []:
+        if key == b"authorization":
+            if (space := value.find(b" ")) >= 0:
+                value = value[:space] + b" [REDACTED]"
+        yield (key, value)
+
+
 class WebsocketTracing:
     def __init__(self, app: ASGIApp):
         self.app = app
@@ -74,7 +86,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         client: tuple[str, int] = scope.get("client", ("unknown", 0))
         extra = {"conn": conn_id, "client": client}
 
-        WS_LOGGER.debug("%r", scope, extra=extra)
+        WS_LOGGER.debug(
+            "New Connection from %r",
+            {**scope, "headers": list(_redact_headers(scope.get("headers")))},
+            extra=extra,
+        )
 
         async def local_send(msg: Message):
             match msg.get("type"):

From 63f5c26e27dead34c239e83fffb997db4f029dad Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 1 Jul 2026 18:01:00 +0100
Subject: [PATCH 37/60] Use send_text instead of send json

Saves converting messages is two steps
---
 src/blueapi/service/main.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index bec8e3f69..53c28c812 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -662,13 +662,13 @@ async def run_plan(
             runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
             async for evt in events:
                 LOGGER.debug("Event: %s", evt)
-                await ws.send_json(Update(data=evt).model_dump(mode="json"))
+                await ws.send_text(Update(data=evt).model_dump_json())
                 if isinstance(evt, WorkerEvent) and evt.is_complete():
                     LOGGER.info("End of stream")
                     break
     except WorkerBusyError:
         LOGGER.error("Worker was busy")
-        await ws.send_json(ServerBusy())
+        await ws.send_text(ServerBusy().model_dump_json())
         await ws.close(code=1013, reason="Worker busy")
     except WebSocketDisconnect:
         LOGGER.info("Client disconnected")

From dacd184711a8ea8044d990f03ea452e04cdb630c Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 1 Jul 2026 18:12:09 +0100
Subject: [PATCH 38/60] Improve error handling

---
 src/blueapi/client/rest.py  | 2 +-
 src/blueapi/service/main.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 4d12cf4a7..7cb9cd830 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -384,7 +384,7 @@ def run_blocking(
                                 ]
                             )
                         case PlanNotFound(plan_name=name):
-                            raise UnknownPlanError(name)
+                            raise UnknownPlanError(message=name)
         except InvalidStatus as istat:
             match istat.response.status_code:
                 case 401 | 403:
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 53c28c812..9bdd65828 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -646,7 +646,7 @@ async def run_plan(
         )
         LOGGER.info("Task ID: %s", task_id)
     except ValidationError as ve:
-        LOGGER.error("Args not valid", exc_info=True)
+        LOGGER.info("Plan args not valid: %s - %s", task_request, ve)
         await ws.send_text(InvalidArgs.from_validation_error(ve).model_dump_json())
         await ws.close(code=4002, reason="Invalid Args")
         return

From d3dedcea298fd39fb481a68fe8209cc6f518223d Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 2 Jul 2026 09:39:23 +0100
Subject: [PATCH 39/60] Check for active task before running task

This is trying to mask a race condition that should be handled elsewhere.
If a task is started between the check and submission the task will still
be started. This adds a tiledwriter to the run_engine that will receive
messages from the exisiting plan and then fall over when it doesn't recognise
UUIDs. It might do for now though.
---
 src/blueapi/client/rest.py       | 3 +++
 src/blueapi/service/interface.py | 2 +-
 src/blueapi/service/main.py      | 6 +++++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 7cb9cd830..c45fc0874 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -38,6 +38,7 @@
     ControlResponse,
     InvalidArgs,
     PlanNotFound,
+    ServerBusy,
     Submit,
     Update,
 )
@@ -385,6 +386,8 @@ def run_blocking(
                             )
                         case PlanNotFound(plan_name=name):
                             raise UnknownPlanError(message=name)
+                        case ServerBusy():
+                            raise BlueskyRemoteControlError(409, "Server is busy")
         except InvalidStatus as istat:
             match istat.response.status_code:
                 case 401 | 403:
diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py
index 1bd6e7f4a..0944af59c 100644
--- a/src/blueapi/service/interface.py
+++ b/src/blueapi/service/interface.py
@@ -230,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
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 9bdd65828..6bbacccc3 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -658,7 +658,11 @@ async def run_plan(
 
     try:
         with runner.event_pipe() as events:
-            LOGGER.info("Created event pipe")
+            active_task = runner.run(interface.get_active_task)
+            if active_task is not None and not active_task.is_complete:
+                await ws.send_text(ServerBusy().model_dump_json())
+                await ws.close(code=1013, reason="Worker busy")
+                return
             runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
             async for evt in events:
                 LOGGER.debug("Event: %s", evt)

From 9a5e46103b3cf761b588378cc70ab86178e57a15 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 3 Jul 2026 10:33:22 +0100
Subject: [PATCH 40/60] Filter events to only relevant ones

---
 src/blueapi/service/main.py | 2 ++
 src/blueapi/worker/event.py | 6 ++++++
 2 files changed, 8 insertions(+)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 6bbacccc3..6b34483b3 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -665,6 +665,8 @@ async def run_plan(
                 return
             runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
             async for evt in events:
+                if evt.task_id != task_id:
+                    continue
                 LOGGER.debug("Event: %s", evt)
                 await ws.send_text(Update(data=evt).model_dump_json())
                 if isinstance(evt, WorkerEvent) and evt.is_complete():
diff --git a/src/blueapi/worker/event.py b/src/blueapi/worker/event.py
index 880c3da21..25aae20f9 100644
--- a/src/blueapi/worker/event.py
+++ b/src/blueapi/worker/event.py
@@ -173,3 +173,9 @@ def is_error(self) -> bool:
 
     def is_complete(self) -> bool:
         return self.task_status is not None and self.task_status.task_complete
+
+    @property
+    def task_id(self) -> str | None:
+        if task := self.task_status:
+            return task.task_id
+        return None

From e8d53d9cc855afcbb02790fb282578a70402173b Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 3 Jul 2026 10:58:56 +0100
Subject: [PATCH 41/60] Logging adjustments

---
 src/blueapi/service/main.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index 6b34483b3..f0bfd8ac9 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -635,7 +635,7 @@ async def run_plan(
     try:
         task_request = ControlRequest.validate_json(rq)
     except ValidationError:
-        LOGGER.error("Failed to deserialize request", exc_info=True)
+        LOGGER.info("Failed to deserialize request: %r", rq, exc_info=True)
         await ws.close(code=1007, reason="Invalid Request")
         return
     LOGGER.info("Plan request: %s", task_request)
@@ -651,7 +651,7 @@ async def run_plan(
         await ws.close(code=4002, reason="Invalid Args")
         return
     except KeyError as ke:
-        LOGGER.error("Plan %r not found", ke.args[0])
+        LOGGER.info("Plan %r not recognised", ke.args[0])
         await ws.send_text(PlanNotFound(plan_name=ke.args[0]).model_dump_json())
         await ws.close(code=4001, reason="unknown plan")
         return
@@ -670,7 +670,7 @@ async def run_plan(
                 LOGGER.debug("Event: %s", evt)
                 await ws.send_text(Update(data=evt).model_dump_json())
                 if isinstance(evt, WorkerEvent) and evt.is_complete():
-                    LOGGER.info("End of stream")
+                    LOGGER.debug("End of stream")
                     break
     except WorkerBusyError:
         LOGGER.error("Worker was busy")

From fbd145e46c6dba8fcc465778208ed8d69b387100 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 3 Jul 2026 11:11:45 +0100
Subject: [PATCH 42/60] Exclude all private methods from
 blueapi_rest_client_get_methods

_ws_address doesn't count as a get method
---
 tests/system_tests/test_blueapi_system.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py
index 8e2f5e528..c09a272e4 100644
--- a/tests/system_tests/test_blueapi_system.py
+++ b/tests/system_tests/test_blueapi_system.py
@@ -232,7 +232,7 @@ def blueapi_rest_client_get_methods() -> list[str]:
     return [
         name
         for name, method in BlueapiRestClient.__dict__.items()
-        if not name.startswith("__")
+        if not name.startswith("_")
         and callable(method)
         and len(params := inspect.signature(method).parameters) == 1
         and "self" in params

From 313d183fdbed5c15af088411aadcf5cb6f0884b5 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 3 Jul 2026 17:32:58 +0100
Subject: [PATCH 43/60] Move ws_address method to config and deal with empty
 paths

---
 helm/blueapi/config_schema.json |  2 +-
 helm/blueapi/values.schema.json |  2 +-
 src/blueapi/client/rest.py      | 13 ++-----------
 src/blueapi/config.py           | 19 ++++++++++++++++++-
 tests/unit_tests/test_config.py | 15 ++++++++++++++-
 5 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json
index 3b9d03138..765648614 100644
--- a/helm/blueapi/config_schema.json
+++ b/helm/blueapi/config_schema.json
@@ -348,7 +348,7 @@
             "additionalProperties": false,
             "properties": {
                 "url": {
-                    "default": "http://localhost:8000/",
+                    "default": "http://localhost:8000",
                     "format": "uri",
                     "maxLength": 2083,
                     "minLength": 1,
diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json
index c16ca03dc..5ab65ff8b 100644
--- a/helm/blueapi/values.schema.json
+++ b/helm/blueapi/values.schema.json
@@ -782,7 +782,7 @@
                 },
                 "url": {
                     "title": "Url",
-                    "default": "http://localhost:8000/",
+                    "default": "http://localhost:8000",
                     "type": "string",
                     "format": "uri",
                     "maxLength": 2083,
diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index c45fc0874..b97040070 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -10,7 +10,7 @@
     get_tracer,
     start_as_current_span,
 )
-from pydantic import BaseModel, TypeAdapter, ValidationError, WebsocketUrl
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from pydantic_core import PydanticSerializationError
 from websockets.exceptions import InvalidStatus
 from websockets.sync.client import connect
@@ -358,7 +358,7 @@ def _request_and_deserialize(
     def run_blocking(
         self, req: TaskRequest
     ) -> Iterable[DataEvent | WorkerEvent | ProgressEvent]:
-        url = self._ws_address().unicode_string().rstrip("/") + "/api/v2/run_plan"
+        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()
@@ -395,15 +395,6 @@ def run_blocking(
             print(vars(istat))
             return
 
-    def _ws_address(self) -> WebsocketUrl:
-        api = self._config.url
-        if api.host is None:
-            raise ValueError("No host configured")
-        scheme = "ws" if api.scheme == "http" else "wss"
-        return WebsocketUrl.build(
-            scheme=scheme, host=api.host, port=api.port, path=api.path
-        )
-
 
 # https://github.com/DiamondLightSource/blueapi/issues/1256 - remove before 2.0
 def __getattr__(name: str):
diff --git a/src/blueapi/config.py b/src/blueapi/config.py
index a181f4c34..cd468931d 100644
--- a/src/blueapi/config.py
+++ b/src/blueapi/config.py
@@ -22,6 +22,7 @@
     TypeAdapter,
     UrlConstraints,
     ValidationError,
+    WebsocketUrl,
     field_validator,
     model_validator,
 )
@@ -167,9 +168,25 @@ class CORSConfig(BlueapiBaseModel):
 
 
 class RestConfig(BlueapiBaseModel):
-    url: HttpUrl = HttpUrl("http://localhost:8000")
+    url: Annotated[
+        HttpUrl,
+        UrlConstraints(preserve_empty_path=True),
+        Field("http://localhost:8000", validate_default=True),
+    ]
     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"
+        return WebsocketUrl.build(
+            scheme=scheme, host=api.host, port=api.port, path=api.path
+        )
+
 
 class ScratchRepository(BlueapiBaseModel):
     name: str = Field(
diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py
index ed00587a1..8ccdbb369 100644
--- a/tests/unit_tests/test_config.py
+++ b/tests/unit_tests/test_config.py
@@ -12,13 +12,14 @@
 import responses
 import yaml
 from bluesky_stomp.models import BasicAuthentication
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, WebsocketUrl
 
 from blueapi.config import (
     CONFIG_SCHEMA_LOCATION,
     ApplicationConfig,
     ConfigLoader,
     OIDCConfig,
+    RestConfig,
     generate_config_schema,
 )
 from blueapi.utils import InvalidConfigError
@@ -571,3 +572,15 @@ def test_issuer_url_preferred_over_well_known_url(caplog):
             oidc._well_known_url == "issuer_url" + "/.well-known/openid-configuration"
         )
     assert "well_known_url and issuer are both set" in caplog.text
+
+
+@pytest.mark.parametrize(
+    "url,ws_url",
+    [
+        ("http://localhost:8000", "ws://localhost:8000"),
+        ("https://localhost:8000", "wss://localhost:8000"),
+    ],
+)
+def test_ws_address(url: str, ws_url: str):
+    conf = RestConfig.model_validate({"url": url})
+    assert conf.ws_address == WebsocketUrl(ws_url)

From 337c9075229a9c5b93e1441a3037bef970498945 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 3 Jul 2026 17:33:22 +0100
Subject: [PATCH 44/60] Add tests for websocket tracing

---
 src/blueapi/service/middleware.py           |   6 +-
 tests/unit_tests/service/test_middleware.py | 195 +++++++++++++++++++-
 2 files changed, 198 insertions(+), 3 deletions(-)

diff --git a/src/blueapi/service/middleware.py b/src/blueapi/service/middleware.py
index cb94747c5..5adffc368 100644
--- a/src/blueapi/service/middleware.py
+++ b/src/blueapi/service/middleware.py
@@ -129,9 +129,11 @@ async def local_receive() -> Message:
             message = await receive()
             match message.get("type"):
                 case "websocket.receive":
-                    WS_LOGGER.debug("Received: %r", message.get("text"))
+                    WS_LOGGER.debug("Received: %r", message.get("text"), extra=extra)
                 case "websocket.connect":
-                    WS_LOGGER.debug("New connection from %s:%d", *client)
+                    WS_LOGGER.debug("New connection from %s:%d", *client, extra=extra)
+                case _:
+                    WS_LOGGER.debug("Received other: %r", message, extra=extra)
             return message
 
         return await self.app(scope, local_receive, local_send)
diff --git a/tests/unit_tests/service/test_middleware.py b/tests/unit_tests/service/test_middleware.py
index 5f6dbeac6..1e0a96a93 100644
--- a/tests/unit_tests/service/test_middleware.py
+++ b/tests/unit_tests/service/test_middleware.py
@@ -1,4 +1,5 @@
-from unittest.mock import AsyncMock, MagicMock, Mock, patch
+from typing import Any
+from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
 
 import pytest
 from starlette.types import ASGIApp
@@ -11,6 +12,8 @@
     VERSION,
     ObservabilityContextPropagator,
     VersionHeaders,
+    WebsocketTracing,
+    _redact_headers,
 )
 
 
@@ -106,3 +109,193 @@ async def test_obs_context_passes_vendor_context(app: Mock, protocol: str):
             ApplicationConfig.VENDOR_CONTEXT_HEADER: "vendor_context",
         }
     )
+
+
+def test_redact_headers():
+    assert list(_redact_headers([(b"authorization", b"Bearer foobar")])) == [
+        (b"authorization", b"Bearer [REDACTED]")
+    ]
+    assert list(_redact_headers([(b"other-header", b"Not affected")])) == [
+        (b"other-header", b"Not affected")
+    ]
+
+
+@pytest.fixture
+def asgi() -> AsyncMock:
+    return AsyncMock(name="asgi-app", spec=ASGIApp)
+
+
+@pytest.fixture
+def ws_tracer(asgi: Mock) -> WebsocketTracing:
+    return WebsocketTracing(asgi)
+
+
+@pytest.fixture
+def send() -> AsyncMock:
+    return AsyncMock()
+
+
+@pytest.fixture
+def receive() -> AsyncMock:
+    return AsyncMock()
+
+
+# logger patch that defaults to enabled for all levels
+def patch_ws_logger(func):
+    return patch(
+        "blueapi.service.middleware.WS_LOGGER",
+        isEnabledFor=Mock(name="custom", return_value=True),
+    )(func)
+
+
+@patch_ws_logger
+async def test_websocket_tracing_does_nothing_when_not_debug(
+    log: Mock,
+    asgi: AsyncMock,
+    ws_tracer: WebsocketTracing,
+    send: AsyncMock,
+    receive: AsyncMock,
+):
+    scope = {"type": "websocket"}
+    log.isEnabledFor.return_value = False
+    await ws_tracer(scope, receive, send)
+
+    asgi.assert_called_once_with(scope, receive, send)
+
+
+@patch_ws_logger
+async def test_websocket_tracing_does_nothing_for_http(
+    log: Mock,
+    asgi: AsyncMock,
+    ws_tracer: WebsocketTracing,
+    send: AsyncMock,
+    receive: AsyncMock,
+):
+    scope = {"type": "http"}
+    await ws_tracer(scope, receive, send)
+
+    asgi.assert_called_once_with(scope, receive, send)
+
+
+@patch_ws_logger
+async def test_websocket_tracing_logs_new_connection(
+    log: Mock, ws_tracer: WebsocketTracing, send: AsyncMock, receive: AsyncMock
+):
+    scope = {"type": "websocket", "headers": [(b"authorization", b"bearer foobar")]}
+    await ws_tracer(scope, receive, send)
+    log.debug.assert_called_once_with(
+        "New Connection from %r",
+        {"type": "websocket", "headers": [(b"authorization", b"bearer [REDACTED]")]},
+        extra=ANY,
+    )
+
+
+@pytest.mark.parametrize(
+    "type,other,log_args",
+    [
+        (
+            "websocket.send",
+            {"text": "demo"},
+            ("Sending: %r", "demo"),
+        ),
+        (
+            "websocket.accept",
+            {"headers": [(b"bapi-version", b"1.2.3")]},
+            (
+                "Accepting websocket - sending headers: %r",
+                [(b"bapi-version", b"1.2.3")],
+            ),
+        ),
+        (
+            "websocket.close",
+            {"code": 1234, "reason": "error_code"},
+            ("Closing with code: %r, reason: %r", 1234, "error_code"),
+        ),
+        (
+            "websocket.http.response.start",
+            {"status": "ws-status", "headers": [(b"bapi-version", b"1.2.3")]},
+            (
+                "HTTP Response: status=%r, headers=%r",
+                "ws-status",
+                [(b"bapi-version", b"1.2.3")],
+            ),
+        ),
+        (
+            "websocket.http.response.body",
+            {"body": "response content"},
+            (
+                "HTTP Response Content: %r",
+                "response content",
+            ),
+        ),
+        (
+            "unknown.msg.type",
+            {"other": "data"},
+            ("Sending other: %r", {"type": "unknown.msg.type", "other": "data"}),
+        ),
+    ],
+)
+@patch_ws_logger
+async def test_websocket_tracing_local_send(
+    log: Mock,
+    asgi: AsyncMock,
+    ws_tracer: WebsocketTracing,
+    send: AsyncMock,
+    type: str,
+    other: dict[str, Any],
+    log_args: tuple[tuple[str, ...], dict[str, Any]],
+):
+    await ws_tracer({"type": "websocket"}, AsyncMock(), send)
+
+    _, _, local_send = asgi.call_args[0]
+
+    message = {"type": type, **other}
+    await local_send(message)
+    log.debug.assert_called_with(*log_args, extra=ANY)
+
+    # Original send method should be called with original message
+    send.assert_called_once_with(message)
+
+
+@pytest.mark.parametrize(
+    "type,other,log_args",
+    [
+        (
+            "websocket.receive",
+            {"text": "demo"},
+            ("Received: %r", "demo"),
+        ),
+        (
+            "websocket.connect",
+            {},
+            ("New connection from %s:%d", "unknown", 0),
+        ),
+        ("unknown.msg", {}, ("Received other: %r", {"type": "unknown.msg"})),
+    ],
+)
+@patch_ws_logger
+async def test_websocket_tracing_local_receive(
+    log: Mock,
+    asgi: AsyncMock,
+    ws_tracer: WebsocketTracing,
+    receive: AsyncMock,
+    type: str,
+    other: dict[str, Any],
+    log_args: tuple[tuple[str, ...], dict[str, Any]],
+):
+    await ws_tracer({"type": "websocket"}, receive, AsyncMock())
+
+    _, local_recv, _ = asgi.call_args[0]
+
+    message = {"type": type, **other}
+    receive.return_value = message
+
+    received = await local_recv()
+
+    # original receive called to get message
+    receive.assert_called_once_with()
+
+    log.debug.assert_called_with(*log_args, extra=ANY)
+
+    # We should not be modifying anything
+    assert received == message

From 6e0bfbe598f807a9ac5e2f30ffe333a205d770cc Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 21 Jul 2026 12:16:11 +0100
Subject: [PATCH 45/60] Handle pydantic's messing up of URL paths

Handle the path mangling in code instead of via URL constraints and
Field annotations.

Stops the type checking falling over
---
 helm/blueapi/config_schema.json |  2 +-
 helm/blueapi/values.schema.json |  2 +-
 src/blueapi/config.py           | 12 ++++++------
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json
index 765648614..3b9d03138 100644
--- a/helm/blueapi/config_schema.json
+++ b/helm/blueapi/config_schema.json
@@ -348,7 +348,7 @@
             "additionalProperties": false,
             "properties": {
                 "url": {
-                    "default": "http://localhost:8000",
+                    "default": "http://localhost:8000/",
                     "format": "uri",
                     "maxLength": 2083,
                     "minLength": 1,
diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json
index 5ab65ff8b..c16ca03dc 100644
--- a/helm/blueapi/values.schema.json
+++ b/helm/blueapi/values.schema.json
@@ -782,7 +782,7 @@
                 },
                 "url": {
                     "title": "Url",
-                    "default": "http://localhost:8000",
+                    "default": "http://localhost:8000/",
                     "type": "string",
                     "format": "uri",
                     "maxLength": 2083,
diff --git a/src/blueapi/config.py b/src/blueapi/config.py
index cd468931d..e1b27b435 100644
--- a/src/blueapi/config.py
+++ b/src/blueapi/config.py
@@ -168,11 +168,7 @@ class CORSConfig(BlueapiBaseModel):
 
 
 class RestConfig(BlueapiBaseModel):
-    url: Annotated[
-        HttpUrl,
-        UrlConstraints(preserve_empty_path=True),
-        Field("http://localhost:8000", validate_default=True),
-    ]
+    url: HttpUrl = HttpUrl("http://localhost:8000")
     cors: CORSConfig | None = None
 
     @property
@@ -183,8 +179,12 @@ def ws_address(self) -> WebsocketUrl:
             # 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=api.path
+            scheme=scheme, host=api.host, port=api.port, path=path
         )
 
 

From bac37dcad6069031866f6c5418f1ad913434bc7f Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Tue, 21 Jul 2026 12:17:19 +0100
Subject: [PATCH 46/60] Extract cookies/headers from connection manually

Using pydantic's dependency handling means they end up in the openapi schema
even when auth is not being used.
---
 src/blueapi/service/authentication.py           | 11 ++++++-----
 tests/unit_tests/service/test_authentication.py |  3 ++-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py
index 614f42b7c..9177ad47f 100644
--- a/src/blueapi/service/authentication.py
+++ b/src/blueapi/service/authentication.py
@@ -15,7 +15,7 @@
 import httpx
 import jwt
 import requests
-from fastapi import Cookie, Depends, Header, HTTPException
+from fastapi import Depends, HTTPException
 from fastapi.requests import HTTPConnection
 from fastapi.security.utils import get_authorization_scheme_param
 from pydantic import TypeAdapter
@@ -279,11 +279,12 @@ def sync_auth_flow(self, request):
         yield request
 
 
-def unchecked_bearer_token(
-    auth_header: str | None = Header(alias="Authorization", default=None),
-    auth_cookie: str | None = Cookie(alias="Authorization", default=None),
-) -> 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
diff --git a/tests/unit_tests/service/test_authentication.py b/tests/unit_tests/service/test_authentication.py
index cfc277ee2..a76375812 100644
--- a/tests/unit_tests/service/test_authentication.py
+++ b/tests/unit_tests/service/test_authentication.py
@@ -205,7 +205,8 @@ def test_tiled_auth_sync_auth_flow():
 def test_unchecked_bearer_token(
     header: str | None, cookie: str | None, token: str | None
 ):
-    assert unchecked_bearer_token(header, cookie) == token
+    req = Mock(headers={"Authorization": header}, cookies={"Authorization": cookie})
+    assert unchecked_bearer_token(req) == token
 
 
 def test_access_token():

From 91be3d22b5bb835297b5affeede343a0c74d0647 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Thu, 23 Jul 2026 12:11:44 +0100
Subject: [PATCH 47/60] WorkerEvent property tests

---
 tests/unit_tests/worker/test_task_worker.py | 40 +++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py
index 588c37d8e..5e6553d2f 100644
--- a/tests/unit_tests/worker/test_task_worker.py
+++ b/tests/unit_tests/worker/test_task_worker.py
@@ -922,3 +922,43 @@ def test_task_result_serialization(plan_result, task_result, type_name):
     res = TaskResult.from_result(plan_result)
     assert res.result == task_result
     assert res.type == type_name
+
+
+@pytest.mark.parametrize(
+    "status,complete",
+    [
+        (None, False),
+        (
+            TaskStatus(
+                task_id="foo", result=None, task_complete=False, task_failed=False
+            ),
+            False,
+        ),
+        (
+            TaskStatus(
+                task_id="foo", result=None, task_complete=True, task_failed=False
+            ),
+            True,
+        ),
+    ],
+)
+def test_worker_event_complete(status: TaskStatus | None, complete: bool):
+    event = WorkerEvent(
+        state=WorkerState.IDLE, task_status=status
+    )  # state is not used in check
+    assert event.is_complete() == complete
+
+
+def test_worker_event_task_id():
+    event = WorkerEvent(
+        state=WorkerState.RUNNING,
+        task_status=TaskStatus(
+            task_id="foo", result=None, task_complete=False, task_failed=False
+        ),
+    )
+    assert event.task_id == "foo"
+
+
+def test_worker_event_no_task_id():
+    event = WorkerEvent(state=WorkerState.IDLE, task_status=None)
+    assert event.task_id is None

From 6a34fd2e3fd4830b2a3966b6aa3465bb5417e9bf Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 24 Jul 2026 18:47:57 +0100
Subject: [PATCH 48/60] More coverage hunting

---
 src/blueapi/service/runner.py           |  2 +-
 tests/unit_tests/cli/test_cli.py        | 15 +++++
 tests/unit_tests/service/test_runner.py | 81 +++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/service/runner.py b/src/blueapi/service/runner.py
index 83153ed5f..ab17c159d 100644
--- a/src/blueapi/service/runner.py
+++ b/src/blueapi/service/runner.py
@@ -173,7 +173,7 @@ async def __anext__(self) -> WorkerEvent | DataEvent | ProgressEvent:
                 await data_available.wait()
                 data_available.clear()
             return self._rx.recv()
-        except BrokenPipeError:
+        except EOFError:
             raise StopAsyncIteration() from None
         finally:
             asyncio.get_event_loop().remove_reader(self._rx.fileno())
diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py
index fe71825c1..632f28d65 100644
--- a/tests/unit_tests/cli/test_cli.py
+++ b/tests/unit_tests/cli/test_cli.py
@@ -1477,3 +1477,18 @@ def test_host_overrides_config(runner: CliRunner):
     )
     assert response.call_count == 1
     assert res.exit_code == 0
+
+
+@patch("blueapi.cli.cli.BlueapiClient")
+def test_run_ws_runs_blocking_plan(mock_client: Mock, runner: CliRunner):
+    bc = mock_client.from_config()
+    res = runner.invoke(
+        main,
+        ["controller", "run", "-i", "cm12345-1", "--ws", "name"],
+    )
+    bc.add_callback.assert_called_once()
+    bc.run_task.assert_not_called()
+    bc.run_blocking.assert_called_once_with(
+        TaskRequest(name="name", params={}, instrument_session="cm12345-1"),
+    )
+    assert res.exit_code == 0
diff --git a/tests/unit_tests/service/test_runner.py b/tests/unit_tests/service/test_runner.py
index b41b3fd4c..0e27d2901 100644
--- a/tests/unit_tests/service/test_runner.py
+++ b/tests/unit_tests/service/test_runner.py
@@ -1,5 +1,6 @@
 import uuid
 from collections.abc import Callable
+from multiprocessing import Pipe
 from multiprocessing.pool import Pool as PoolClass
 from typing import Any, Generic, TypeVar
 from unittest.mock import MagicMock, Mock, NonCallableMock, patch
@@ -14,12 +15,14 @@
 from blueapi.service import interface
 from blueapi.service.model import EnvironmentResponse
 from blueapi.service.runner import (
+    EventStream,
     InvalidRunnerStateError,
     RpcError,
     WorkerDispatcher,
     _safe_exception_message,
     import_and_run_function,
 )
+from blueapi.worker.event import TaskStatus, WorkerEvent, WorkerState
 
 
 @pytest.fixture
@@ -300,3 +303,81 @@ def test_run_span_ok(
 ):
     with asserting_span_exporter(exporter, "run", "function", "args", "kwargs"):
         started_runner.run(interface.get_plans)
+
+
+@patch("blueapi.service.runner.Pipe")
+def test_event_pipe(mock_pipe: Mock):
+    tx = Mock()
+    rx = Mock()
+
+    mock_pipe.return_value = (tx, rx)
+
+    dispatcher = Mock()
+    dispatcher.run.side_effect = lambda mth, *a: {
+        interface.pipe_events: 42,
+        interface.unpipe_events: None,
+    }[mth]
+    evt_pipe = WorkerDispatcher.event_pipe(dispatcher)
+
+    with evt_pipe:
+        dispatcher.run.assert_called_with(interface.pipe_events, tx)
+
+    dispatcher.run.assert_called_with(interface.unpipe_events, 42)
+
+    assert len(dispatcher.run.mock_calls) == 2
+
+
+async def test_event_stream():
+    tx, rx = Pipe()
+    stream = EventStream(rx)
+
+    evt = WorkerEvent(
+        state=WorkerState.RUNNING,
+        task_status=TaskStatus(
+            task_id="foo", task_complete=False, task_failed=False, result=None
+        ),
+    )
+
+    tx.send(evt)
+
+    assert (await anext(stream)) == evt
+
+
+async def test_end_of_event_stream():
+    tx, rx = Pipe()
+    stream = EventStream(rx)
+
+    tx.close()
+
+    with pytest.raises(StopAsyncIteration):
+        await anext(stream)
+
+
+async def test_aiter_event_stream_is_self():
+    stream = EventStream(Mock())
+    assert aiter(stream) is stream
+
+
+async def test_event_stream_no_event():
+    tx, rx = Pipe()
+    stream = EventStream(rx)
+
+    evt = WorkerEvent(
+        state=WorkerState.RUNNING,
+        task_status=TaskStatus(
+            task_id="foo", task_complete=False, task_failed=False, result=None
+        ),
+    )
+
+    read = anext(stream)
+    # start the coroutine
+    _ = read.send(None)
+
+    print("Task not done")
+
+    tx.send(evt)
+
+    with pytest.raises(StopAsyncIteration):
+        read.send(None)
+    # assert await read == evt
+    # assert not read.done()

From 0e64d548082eaa04296e511b00be9506ab70d656 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 27 Jul 2026 12:50:54 +0100
Subject: [PATCH 49/60] Add interface pipe tests

---
 tests/unit_tests/service/test_interface.py | 45 ++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tests/unit_tests/service/test_interface.py b/tests/unit_tests/service/test_interface.py
index ac756fe85..d8c0de24b 100644
--- a/tests/unit_tests/service/test_interface.py
+++ b/tests/unit_tests/service/test_interface.py
@@ -2,6 +2,7 @@
 import uuid
 from dataclasses import dataclass
 from inspect import isawaitable
+from multiprocessing.connection import Connection as PipeConnection
 from typing import Any
 from unittest.mock import ANY, MagicMock, Mock, patch
 
@@ -683,3 +684,47 @@ async def test_update_scan_num_side_effect_sets_scan_file_in_re_md(
     assert isawaitable(scan_id) and await scan_id
 
     assert ctx.run_engine.md["scan_file"] == "p46-11"
+
+
+@patch("blueapi.service.interface.worker")
+def test_pipe_events(mock_worker: Mock):
+    worker = mock_worker()
+    tx = Mock(spec=PipeConnection)
+
+    interface.pipe_events(tx)
+
+    worker.worker_events.subscribe.assert_called_once()
+    worker.data_events.subscribe.assert_called_once()
+    worker.progress_events.subscribe.assert_called_once()
+
+    handler = worker.worker_events.subscribe.call_args[0][0]
+
+    evt = Mock()
+    handler(evt, "ignored correlation id")
+    tx.send.assert_called_once_with(evt)
+
+
+@patch("blueapi.service.interface.worker")
+def test_pipe_events_ignores_broken_pipe(mock_worker: Mock):
+    worker = mock_worker()
+    tx = Mock(spec=PipeConnection)
+
+    interface.pipe_events(tx)
+
+    worker.worker_events.subscribe.assert_called_once()
+    handler = worker.worker_events.subscribe.call_args[0][0]
+
+    tx.send.side_effect = BrokenPipeError()
+    # ensure that exceptions are not raised
+    handler(Mock(), "ignored correlation id")
+
+
+@patch("blueapi.service.interface.worker")
+def test_unpipe_events(mock_worker: Mock):
+    worker = mock_worker()
+    handles = interface.SubHandles(worker=1, progress=2, data=3)
+    interface.unpipe_events(handles)
+
+    worker.worker_events.unsubscribe.assert_called_once_with(1)
+    worker.progress_events.unsubscribe.assert_called_once_with(2)
+    worker.data_events.unsubscribe.assert_called_once_with(3)

From 596ca48fd4c0598724756faf56e2b9580557fb0f Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 27 Jul 2026 12:51:37 +0100
Subject: [PATCH 50/60] Fix runner tests

---
 tests/unit_tests/service/test_runner.py | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/tests/unit_tests/service/test_runner.py b/tests/unit_tests/service/test_runner.py
index 0e27d2901..8380784e0 100644
--- a/tests/unit_tests/service/test_runner.py
+++ b/tests/unit_tests/service/test_runner.py
@@ -1,6 +1,7 @@
 import uuid
 from collections.abc import Callable
 from multiprocessing import Pipe
+from multiprocessing.connection import Connection
 from multiprocessing.pool import Pool as PoolClass
 from typing import Any, Generic, TypeVar
 from unittest.mock import MagicMock, Mock, NonCallableMock, patch
@@ -358,9 +359,15 @@ async def test_aiter_event_stream_is_self():
     assert aiter(stream) is stream
 
 
-async def test_event_stream_no_event():
+async def test_event_stream_wait_for_event():
     tx, rx = Pipe()
-    stream = EventStream(rx)
+
+    mrx = Mock(spec=Connection)
+    mrx.poll.side_effect = [False, False, True]
+    mrx.fileno.side_effect = rx.fileno
+    mrx.recv.side_effect = rx.recv
+
+    stream = EventStream(mrx)
 
     evt = WorkerEvent(
         state=WorkerState.RUNNING,
@@ -369,15 +376,9 @@ async def test_event_stream_no_event():
         ),
     )
 
-    read = anext(stream)
-    # start the coroutine
-    _ = read.send(None)
-
-    print("Task not done")
-
     tx.send(evt)
 
-    with pytest.raises(StopAsyncIteration):
-        read.send(None)
-    # assert await read == evt
-    # assert not read.done()
+    read = await anext(stream)
+    assert read == evt
+
+    assert len(mrx.poll.mock_calls) == 3

From 655a8fd4739a55fc25db1c51fd6f7072155877a1 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 27 Jul 2026 20:35:27 +0100
Subject: [PATCH 51/60] More tests

---
 src/blueapi/client/client.py              |  3 +-
 tests/unit_tests/client/test_client.py    | 38 +++++++++++++++++++++++
 tests/unit_tests/service/test_protocol.py | 34 ++++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py
index 3d6e3bf6a..b71da8aa6 100644
--- a/src/blueapi/client/client.py
+++ b/src/blueapi/client/client.py
@@ -472,7 +472,8 @@ def run_blocking(
                 except Exception as e:
                     log.error(f"Callback ({cb}) failed for event: {event}", exc_info=e)
             if isinstance(event, WorkerEvent) and event.is_complete():
-                if event.task_status is None:
+                # 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"
                     )
diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py
index eaf96a3b2..17a3c694e 100644
--- a/tests/unit_tests/client/test_client.py
+++ b/tests/unit_tests/client/test_client.py
@@ -1000,3 +1000,41 @@ def test_client_login_no_oidc(
     client.login()
 
     mock_session_manager.assert_not_called()
+
+
+@pytest.mark.parametrize("event", [COMPLETE_EVENT, FAILED_EVENT])
+def test_run_blocking(event: WorkerEvent, client: BlueapiClient, mock_rest: Mock):
+    mock_rest.run_blocking.side_effect = lambda req: [event]
+    res = client.run_blocking(TaskRequest(name="foo", instrument_session="cm12345-1"))
+    assert res == event.task_status
+
+
+@pytest.mark.parametrize("event", [COMPLETE_EVENT, FAILED_EVENT])
+def test_run_blocking_callbacks(
+    event: WorkerEvent, client: BlueapiClient, mock_rest: Mock
+):
+    callback = Mock()
+    mock_rest.run_blocking.side_effect = lambda req: [event]
+    client.run_blocking(Mock(), on_event=callback)
+
+    callback.assert_called_once_with(event)
+
+
+@pytest.mark.parametrize("event", [COMPLETE_EVENT, FAILED_EVENT])
+def test_run_blocking_client_callbacks(
+    event: WorkerEvent, client: BlueapiClient, mock_rest: Mock
+):
+    callback = Mock()
+    mock_rest.run_blocking.side_effect = lambda req: [event]
+    client.add_callback(callback)
+    client.run_blocking(Mock())
+
+    callback.assert_called_once_with(event)
+
+
+def test_run_blocking_error_if_cut_short(client: BlueapiClient, mock_rest: Mock):
+    mock_rest.run_blocking.side_effect = lambda req: []
+    with pytest.raises(
+        BlueskyRemoteControlError, match="Connection closed before plan completed"
+    ):
+        client.run_blocking(Mock())
diff --git a/tests/unit_tests/service/test_protocol.py b/tests/unit_tests/service/test_protocol.py
index ef95bb8ce..9178162e1 100644
--- a/tests/unit_tests/service/test_protocol.py
+++ b/tests/unit_tests/service/test_protocol.py
@@ -1,6 +1,8 @@
 from typing import Any
 
 import pytest
+from pydantic import ValidationError
+from pydantic_core import InitErrorDetails
 
 from blueapi.service.model import TaskRequest
 from blueapi.service.protocol import (
@@ -68,3 +70,35 @@ def test_request_deserialization(src: str, res: Any):
 def test_response_deserialization(src: str, res: Any):
     req = ControlResponse.validate_json(src)
     assert req == res
+
+
+def test_from_empty_validation_error():
+    err = InvalidArgs.from_validation_error(
+        ValidationError("Error validating request", [])
+    )
+    assert err == InvalidArgs(errors=[])
+
+
+def test_from_validation_error():
+    err = InvalidArgs.from_validation_error(
+        ValidationError.from_exception_data(
+            title="Error validating request",
+            line_errors=[
+                InitErrorDetails(
+                    loc=("foo", "bar"),
+                    type="missing",
+                    input={"foo": {"no": "bar"}},
+                )
+            ],
+        ),
+    )
+    assert err == InvalidArgs(
+        errors=[
+            ArgumentError(
+                loc=["body", "params", "foo", "bar"],
+                msg="Field required",
+                type="missing",
+                input={"foo": {"no": "bar"}},
+            )
+        ]
+    )

From b2936d4bbf0255b6d5eac23cf4b27f800c612a8e Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Mon, 27 Jul 2026 20:36:01 +0100
Subject: [PATCH 52/60] Start to test main method

---
 tests/unit_tests/service/test_main.py | 29 +++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/tests/unit_tests/service/test_main.py b/tests/unit_tests/service/test_main.py
index 93df36266..63bc73f19 100644
--- a/tests/unit_tests/service/test_main.py
+++ b/tests/unit_tests/service/test_main.py
@@ -1,18 +1,22 @@
+import contextlib
 from unittest import mock
-from unittest.mock import Mock, call, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
 
 import pytest
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, WebSocket
 from fastapi.testclient import TestClient
 
 from blueapi import __version__
 from blueapi.config import ApplicationConfig
+from blueapi.service import interface
 from blueapi.service.main import (
     get_passthrough_headers,
     lifespan,
     log_request_details,
+    run_plan,
 )
 from blueapi.service.middleware import VersionHeaders
+from blueapi.service.runner import WorkerDispatcher
 
 
 async def test_add_version_header():
@@ -116,3 +120,24 @@ async def test_lifespan(setup: Mock, teardown: Mock):
         teardown.assert_not_called()
 
     teardown.assert_called_once()
+
+
+async def test_websocket_run_plan():
+    ws = Mock(spec=WebSocket)
+    runner = Mock(spec=WorkerDispatcher)
+    events = MagicMock()
+    events.__aiter__.return_value = MagicMock(__anext__=Mock(side_effect=[1, 2, 3]))
+    runner.event_pipe.return_value = contextlib.nullcontext(events)
+    runner.run.side_effect = lambda mth, *a, **kw: {
+        interface.submit_task: "task_uid"
+    }.get(mth)
+
+    ws.receive_text = AsyncMock(
+        return_value="""{
+            "kind": "submit",
+            "task": {"name": "foo", "params": {}, "instrument_session": "cm12345-1"}
+            }"""
+    )
+
+    await run_plan(ws, runner, user="abc12345")
+    ws.close.assert_called_once_with()

From c5d1cfe307ecb09d3cc09551bdf1d1889c19ff99 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 29 Jul 2026 15:03:47 +0100
Subject: [PATCH 53/60] Add rest client run_blocking tests

---
 src/blueapi/client/rest.py             |  4 +-
 tests/unit_tests/client/test_client.py | 11 +++
 tests/unit_tests/client/test_rest.py   | 95 ++++++++++++++++++++++++++
 3 files changed, 108 insertions(+), 2 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index b97040070..3ab842488 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -392,8 +392,8 @@ def run_blocking(
             match istat.response.status_code:
                 case 401 | 403:
                     raise UnauthorisedAccessError() from None
-            print(vars(istat))
-            return
+                case _:
+                    raise BlueskyRemoteControlError() from istat
 
 
 # https://github.com/DiamondLightSource/blueapi/issues/1256 - remove before 2.0
diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py
index 17a3c694e..7c3608348 100644
--- a/tests/unit_tests/client/test_client.py
+++ b/tests/unit_tests/client/test_client.py
@@ -1038,3 +1038,14 @@ def test_run_blocking_error_if_cut_short(client: BlueapiClient, mock_rest: Mock)
         BlueskyRemoteControlError, match="Connection closed before plan completed"
     ):
         client.run_blocking(Mock())
+
+
+def test_run_blocking_ignores_callback_error(client: BlueapiClient, mock_rest: Mock):
+    mock_rest.run_blocking.side_effect = lambda req: [COMPLETE_EVENT]
+
+    def broken_callback(_: AnyEvent):
+        raise Exception("This callback is broken")
+
+    client.add_callback(broken_callback)
+    res = client.run_blocking(Mock())
+    assert res == COMPLETE_EVENT.task_status
diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py
index c3386ad55..439f2a00c 100644
--- a/tests/unit_tests/client/test_rest.py
+++ b/tests/unit_tests/client/test_rest.py
@@ -10,10 +10,12 @@
 from packaging.version import Version
 from pydantic_core import PydanticSerializationError
 from responses import DELETE, GET, PUT, matchers
+from websockets import Headers, InvalidStatus, Response
 
 from blueapi import __version__
 from blueapi.client.client import DeviceRef
 from blueapi.client.rest import (
+    USER_AGENT,
     BlueapiRestClient,
     BlueskyRemoteControlError,
     BlueskyRequestError,
@@ -431,3 +433,96 @@ def test_get_missing_plan(rest: BlueapiRestClient):
     responses.add(GET, "http://localhost:8000/plans/foo", status=404)
     with pytest.raises(UnknownPlanError):
         rest.get_plan("foo")
+
+
+@patch("blueapi.client.rest.connect")
+def test_run_blocking(mock_connect: Mock, rest: BlueapiRestClient):
+    ws = MagicMock()
+    ws.__enter__.return_value.__iter__.return_value = iter(
+        ['{"kind": "update", "data": {"name": "start", "doc":{}, "task_id":"t_uid"}}']
+    )
+    mock_connect.return_value = ws
+    conn = rest.run_blocking(
+        TaskRequest(name="foo", params={"one": "two"}, instrument_session="cm12345-1")
+    )
+    next(iter(conn))
+    mock_connect.assert_called_once_with(
+        "ws://localhost:8000/api/v2/run_plan",
+        additional_headers={},
+        user_agent_header=USER_AGENT,
+    )
+
+
+@pytest.mark.parametrize(
+    "event,error,message",
+    [
+        (
+            """{
+                "kind":"invalid_args",
+                "errors": [{
+                    "loc": ["bar"],
+                    "msg": "Field required",
+                    "type": "missing",
+                    "input": {}
+                }]
+            }""",
+            InvalidParametersError,
+            "ParameterError",
+        ),
+        ('{"kind": "busy"}', BlueskyRemoteControlError, "Server is busy"),
+        ('{"kind": "plan_not_found", "plan_name": "foo"}', UnknownPlanError, "foo"),
+    ],
+)
+@patch("blueapi.client.rest.connect")
+def test_run_blocking_errors(
+    mock_connect: Mock,
+    rest: BlueapiRestClient,
+    event: str,
+    error: type[Exception],
+    message: str,
+):
+    ws = MagicMock()
+    ws.__enter__.return_value.__iter__.return_value = iter([event])
+    mock_connect.return_value = ws
+    conn = rest.run_blocking(
+        TaskRequest(name="foo", params={"one": "two"}, instrument_session="cm12345-1")
+    )
+    with pytest.raises(error, match=message):
+        next(iter(conn))
+    mock_connect.assert_called_once_with(
+        "ws://localhost:8000/api/v2/run_plan",
+        additional_headers={},
+        user_agent_header=USER_AGENT,
+    )
+
+
+@pytest.mark.parametrize("status_code", [401, 403])
+@patch("blueapi.client.rest.connect")
+def test_run_blocking_ws_failures(
+    mock_connect: Mock, rest: BlueapiRestClient, status_code: int
+):
+    mock_connect.side_effect = InvalidStatus(
+        response=Response(
+            status_code=status_code, reason_phrase="test_error", headers=Headers()
+        )
+    )
+    conn = rest.run_blocking(
+        TaskRequest(name="foo", params={"one": "two"}, instrument_session="cm12345-1")
+    )
+    with pytest.raises(UnauthorisedAccessError):
+        next(iter(conn))
+
+
+@patch("blueapi.client.rest.connect")
+def test_run_blocking_unknown_error(mock_connect: Mock, rest: BlueapiRestClient):
+    mock_connect.side_effect = InvalidStatus(
+        response=Response(
+            status_code=1234, reason_phrase="test_error", headers=Headers()
+        )
+    )
+
+    conn = rest.run_blocking(
+        TaskRequest(name="foo", params={"one": "two"}, instrument_session="cm12345-1")
+    )
+    with pytest.raises(BlueskyRemoteControlError):
+        next(iter(conn))

From 5e88c02122b5b9beac5a3e689e652e7b0410191e Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 29 Jul 2026 16:07:03 +0100
Subject: [PATCH 54/60] test rest ws auth

---
 tests/conftest.py                    | 42 ++++++++++++++++++----------
 tests/unit_tests/client/test_rest.py | 30 ++++++++++++++++++--
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 78af27578..6f517b3bf 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -66,8 +66,13 @@ def oidc_config() -> OIDCConfig:
 
 
 @pytest.fixture
-def config_with_auth(tmp_path: Path) -> str:
-    config = ApplicationConfig(auth_token_path=tmp_path / CACHE_FILE)
+def token_cache_file(tmp_path: Path) -> Path:
+    return tmp_path / CACHE_FILE
+
+
+@pytest.fixture
+def config_with_auth(token_cache_file: Path, tmp_path: Path) -> str:
+    config = ApplicationConfig(auth_token_path=token_cache_file)
     config_path = tmp_path / "auth_config.yaml"
     with open(config_path, mode="w") as valid_auth_config_file:
         valid_auth_config_file.write(yaml.dump(config.model_dump()))
@@ -134,7 +139,7 @@ def _make_token(
 
 @pytest.fixture
 def cached_valid_refresh(
-    tmp_path: Path, expired_token: dict[str, Any], oidc_config: OIDCConfig
+    token_cache_file: Path, expired_token: dict[str, Any], oidc_config: OIDCConfig
 ) -> Path:
     cache = Cache(
         oidc_config=oidc_config,
@@ -142,14 +147,14 @@ def cached_valid_refresh(
         refresh_token=expired_token["refresh_token"],
         id_token=expired_token["id_token"],
     )
-    with open(cache_path := tmp_path / CACHE_FILE, "xb") as cache_file:
+    with open(token_cache_file, "xb") as cache_file:
         cache_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8")))
-    return cache_path
+    return token_cache_file
 
 
 @pytest.fixture
 def cached_expired_refresh(
-    tmp_path: Path, expired_token: dict[str, Any], oidc_config: OIDCConfig
+    token_cache_file: Path, expired_token: dict[str, Any], oidc_config: OIDCConfig
 ) -> Path:
     cache = Cache(
         oidc_config=oidc_config,
@@ -157,14 +162,16 @@ def cached_expired_refresh(
         refresh_token="expired_refresh",
         id_token=expired_token["id_token"],
     )
-    with open(cache_path := tmp_path / CACHE_FILE, "xb") as cache_file:
+    with open(token_cache_file, "xb") as cache_file:
         cache_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8")))
-    return cache_path
+    return token_cache_file
 
 
 @pytest.fixture
 def cached_valid_token(
-    tmp_path: Path, valid_token_with_jwt: dict[str, Any], oidc_config: OIDCConfig
+    token_cache_file: Path,
+    valid_token_with_jwt: dict[str, Any],
+    oidc_config: OIDCConfig,
 ) -> Path:
     cache = Cache(
         oidc_config=oidc_config,
@@ -172,14 +179,21 @@ def cached_valid_token(
         refresh_token=valid_token_with_jwt["refresh_token"],
         id_token=valid_token_with_jwt["id_token"],
     )
-    with open(cache_path := tmp_path / CACHE_FILE, "xb") as cache_file:
+    with open(token_cache_file, "xb") as cache_file:
         cache_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8")))
-    return cache_path
+    return token_cache_file
+
+
+@pytest.fixture
+def cached_valid_token_value(
+    valid_token_with_jwt: dict[str, Any],
+) -> str:
+    return valid_token_with_jwt["access_token"]
 
 
 @pytest.fixture
 def cache_with_invalid_audience(
-    tmp_path: Path,
+    token_cache_file: Path,
     oidc_config: OIDCConfig,
     valid_token_with_jwt_invalid_audience: dict[str, Any],
 ) -> Path:
@@ -189,9 +203,9 @@ def cache_with_invalid_audience(
         refresh_token=valid_token_with_jwt_invalid_audience["refresh_token"],
         id_token=valid_token_with_jwt_invalid_audience["id_token"],
     )
-    with open(cache_path := tmp_path / CACHE_FILE, "xb") as cache_file:
+    with open(token_cache_file, "xb") as cache_file:
         cache_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8")))
-    return cache_path
+    return token_cache_file
 
 
 @pytest.fixture
diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py
index 439f2a00c..0fd3bd0cd 100644
--- a/tests/unit_tests/client/test_rest.py
+++ b/tests/unit_tests/client/test_rest.py
@@ -49,11 +49,13 @@ def rest() -> BlueapiRestClient:
 
 
 @pytest.fixture
-def rest_with_auth(oidc_config: OIDCConfig, tmp_path) -> BlueapiRestClient:
+def rest_with_auth(
+    oidc_config: OIDCConfig, token_cache_file: Path
+) -> BlueapiRestClient:
     return BlueapiRestClient(
         session_manager=SessionManager(
             server_config=oidc_config,
-            cache_manager=SessionCacheManager(tmp_path / "blueapi_cache"),
+            cache_manager=SessionCacheManager(token_cache_file),
         )
     )
 
@@ -453,6 +455,30 @@ def test_run_blocking(mock_connect: Mock, rest: BlueapiRestClient):
     )
 
 
+@patch("blueapi.client.rest.connect")
+def test_run_blocking_auth(
+    mock_connect: Mock,
+    rest_with_auth: BlueapiRestClient,
+    mock_authn_server: responses.RequestsMock,
+    cached_valid_token: Path,  # creates the cache file
+    cached_valid_token_value: str,  # the value from the file
+):
+    ws = MagicMock()
+    ws.__enter__.return_value.__iter__.return_value = iter(
+        ['{"kind": "update", "data": {"name": "start", "doc":{}, "task_id":"t_uid"}}']
+    )
+    mock_connect.return_value = ws
+    conn = rest_with_auth.run_blocking(
+        TaskRequest(name="foo", params={"one": "two"}, instrument_session="cm12345-1")
+    )
+    next(iter(conn))
+    mock_connect.assert_called_once_with(
+        "ws://localhost:8000/api/v2/run_plan",
+        additional_headers={"Authorization": f"Bearer {cached_valid_token_value}"},
+        user_agent_header=USER_AGENT,
+    )
+
+
 @pytest.mark.parametrize(
     "event,error,message",
     [

From 5c88341ea13b6ec870cc5a57bf641bae736de742 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Wed, 29 Jul 2026 18:51:55 +0100
Subject: [PATCH 55/60] Add tests for main run_plan handler

---
 src/blueapi/service/main.py               |   6 +-
 tests/unit_tests/service/test_main.py     |  29 +--
 tests/unit_tests/service/test_rest_api.py | 272 +++++++++++++++++++++-
 3 files changed, 271 insertions(+), 36 deletions(-)

diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index f0bfd8ac9..a5ac96fb6 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -653,16 +653,14 @@ async def run_plan(
     except KeyError as ke:
         LOGGER.info("Plan %r not recognised", ke.args[0])
         await ws.send_text(PlanNotFound(plan_name=ke.args[0]).model_dump_json())
-        await ws.close(code=4001, reason="unknown plan")
+        await ws.close(code=4001, reason="Unknown Plan")
         return
 
     try:
         with runner.event_pipe() as events:
             active_task = runner.run(interface.get_active_task)
             if active_task is not None and not active_task.is_complete:
-                await ws.send_text(ServerBusy().model_dump_json())
-                await ws.close(code=1013, reason="Worker busy")
-                return
+                raise WorkerBusyError("Task already running")
             runner.run(interface.begin_task, task=WorkerTask(task_id=task_id))
             async for evt in events:
                 if evt.task_id != task_id:
diff --git a/tests/unit_tests/service/test_main.py b/tests/unit_tests/service/test_main.py
index 63bc73f19..93df36266 100644
--- a/tests/unit_tests/service/test_main.py
+++ b/tests/unit_tests/service/test_main.py
@@ -1,22 +1,18 @@
-import contextlib
 from unittest import mock
-from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
+from unittest.mock import Mock, call, patch
 
 import pytest
-from fastapi import FastAPI, Request, WebSocket
+from fastapi import FastAPI, Request
 from fastapi.testclient import TestClient
 
 from blueapi import __version__
 from blueapi.config import ApplicationConfig
-from blueapi.service import interface
 from blueapi.service.main import (
     get_passthrough_headers,
     lifespan,
     log_request_details,
-    run_plan,
 )
 from blueapi.service.middleware import VersionHeaders
-from blueapi.service.runner import WorkerDispatcher
 
 
 async def test_add_version_header():
@@ -120,24 +116,3 @@ async def test_lifespan(setup: Mock, teardown: Mock):
         teardown.assert_not_called()
 
     teardown.assert_called_once()
-
-
-async def test_websocket_run_plan():
-    ws = Mock(spec=WebSocket)
-    runner = Mock(spec=WorkerDispatcher)
-    events = MagicMock()
-    events.__aiter__.return_value = MagicMock(__anext__=Mock(side_effect=[1, 2, 3]))
-    runner.event_pipe.return_value = contextlib.nullcontext(events)
-    runner.run.side_effect = lambda mth, *a, **kw: {
-        interface.submit_task: "task_uid"
-    }.get(mth)
-
-    ws.receive_text = AsyncMock(
-        return_value="""{
-            "kind": "submit",
-            "task": {"name": "foo", "params": {}, "instrument_session": "cm12345-1"}
-            }"""
-    )
-
-    await run_plan(ws, runner, user="abc12345")
-    ws.close.assert_called_once_with()
diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py
index dedd9d7ff..2f91ab686 100644
--- a/tests/unit_tests/service/test_rest_api.py
+++ b/tests/unit_tests/service/test_rest_api.py
@@ -1,18 +1,20 @@
+import contextlib
 import uuid
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 from unittest.mock import MagicMock, Mock, patch
 
 import jwt
 import pytest
 from bluesky._vendor.super_state_machine.errors import TransitionError
 from bluesky.protocols import Stoppable
-from fastapi import HTTPException, status
+from fastapi import FastAPI, HTTPException, WebSocketDisconnect, status
 from fastapi.testclient import TestClient
 from httpx2 import Headers
 from pydantic import BaseModel, ValidationError
 from pydantic_core import InitErrorDetails
+from starlette.types import Message, Receive, Scope, Send
 
 from blueapi.config import (
     ApplicationConfig,
@@ -43,7 +45,7 @@
     WorkerTask,
 )
 from blueapi.service.runner import WorkerDispatcher
-from blueapi.worker.event import WorkerState
+from blueapi.worker.event import TaskStatus, WorkerEvent, WorkerState
 from blueapi.worker.task import Task
 from blueapi.worker.task_worker import TrackableTask
 
@@ -57,7 +59,7 @@ class MockCountModel(BaseModel): ...
 
 @pytest.fixture
 def mock_runner() -> Mock:
-    return Mock(spec=WorkerDispatcher)
+    return MagicMock(spec=WorkerDispatcher)
 
 
 @pytest.fixture
@@ -995,3 +997,263 @@ def test_logout_when_oidc_config_invalid(
 
     response = client_with_auth.get("/logout")
     assert response.status_code == status.HTTP_205_RESET_CONTENT
+
+
+async def test_websocket_run_plan(mock_runner: Mock, client: TestClient):
+    mock_runner.run.side_effect = lambda req, *a, **kw: {
+        interface.get_active_task: None,
+        interface.submit_task: "task_id",
+        interface.begin_task: None,
+    }.get(req)
+    mock_runner.event_pipe.return_value = contextlib.nullcontext(
+        _aiter(
+            WorkerEvent(
+                state=WorkerState.RUNNING,
+                task_status=TaskStatus(
+                    task_id="task_id",
+                    result=None,
+                    task_complete=False,
+                    task_failed=False,
+                ),
+            ),
+            WorkerEvent(
+                state=WorkerState.IDLE,
+                task_status=TaskStatus(
+                    task_id="task_id",
+                    result=None,
+                    task_complete=True,
+                    task_failed=False,
+                ),
+            ),
+        )
+    )
+
+    with client.websocket_connect("/api/v2/run_plan") as ws_client:
+        ws_client.send_json(
+            {
+                "kind": "submit",
+                "task": {
+                    "name": "foo",
+                    "params": {"one": "two"},
+                    "instrument_session": "cm12345-1",
+                },
+            }
+        )
+        assert ws_client.receive_json() == {
+            "kind": "update",
+            "data": {
+                "state": "RUNNING",
+                "task_status": {
+                    "task_id": "task_id",
+                    "result": None,
+                    "task_complete": False,
+                    "task_failed": False,
+                },
+                "errors": [],
+                "warnings": [],
+            },
+        }
+        assert ws_client.receive_json() == {
+            "kind": "update",
+            "data": {
+                "state": "IDLE",
+                "task_status": {
+                    "task_id": "task_id",
+                    "result": None,
+                    "task_complete": True,
+                    "task_failed": False,
+                },
+                "errors": [],
+                "warnings": [],
+            },
+        }
+        with pytest.raises(WebSocketDisconnect) as discon:
+            ws_client.receive_text()
+
+        # Check it's a 'normal' end of stream disconnect
+        assert discon.value.code == 1000
+        assert discon.value.reason == ""
+
+
+@pytest.mark.parametrize("req", ["not a json object", "[]", '{"invalid": "keys"}'])
+def test_websocket_run_plan_invalid_request(
+    req: str, mock_runner: Mock, client: TestClient
+):
+    with client.websocket_connect("/api/v2/run_plan") as ws_client:
+        ws_client.send_text(req)
+        with pytest.raises(WebSocketDisconnect) as disco:
+            ws_client.receive_text()
+        assert disco.value.code == 1007
+        assert disco.value.reason == "Invalid Request"
+
+
+@pytest.mark.parametrize(
+    "exc,err_message,code,reason",
+    [
+        (
+            KeyError("foo"),
+            {
+                "kind": "plan_not_found",
+                "plan_name": "foo",
+            },
+            4001,
+            "Unknown Plan",
+        ),
+        (
+            ValidationError("Not valid", []),
+            {"kind": "invalid_args", "errors": []},
+            4002,
+            "Invalid Args",
+        ),
+    ],
+)
+def test_websocket_run_plan_submit_error(
+    exc, err_message, code, reason, mock_runner: Mock, client: TestClient
+):
+    mock_runner.run.side_effect = exc
+    with client.websocket_connect("/api/v2/run_plan") as ws_client:
+        ws_client.send_json(
+            {
+                "kind": "submit",
+                "task": {"name": "foo", "instrument_session": "cm12345-1"},
+            }
+        )
+        assert ws_client.receive_json() == err_message
+        with pytest.raises(WebSocketDisconnect) as disco:
+            ws_client.receive_text()
+        assert disco.value.code == code
+        assert disco.value.reason == reason
+
+
+def test_websocket_run_plan_server_busy(mock_runner: Mock, client: TestClient):
+    mock_runner.run.side_effect = [
+        "task_id",  # submit_task
+        Mock(name="active_task", is_complete=False),
+    ]
+    with client.websocket_connect("/api/v2/run_plan") as ws_client:
+        ws_client.send_json(
+            {"kind": "submit", "task": {"name": "foo", "instrument_session": "cm123-1"}}
+        )
+        assert ws_client.receive_json() == {"kind": "busy"}
+        with pytest.raises(WebSocketDisconnect) as disco:
+            ws_client.receive_text()
+        assert disco.value.code == 1013
+        assert disco.value.reason == "Worker busy"
+    pass
+
+
+def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClient):
+    mock_runner.run.side_effect = lambda req, *a, **kw: {
+        interface.get_active_task: None,
+        interface.submit_task: "task_id",
+        interface.begin_task: None,
+    }.get(req)
+    mock_runner.event_pipe.return_value = contextlib.nullcontext(
+        _aiter(
+            WorkerEvent(
+                state=WorkerState.RUNNING,
+                task_status=TaskStatus(
+                    task_id="other_task_id",
+                    result=None,
+                    task_complete=False,
+                    task_failed=False,
+                ),
+            ),
+            WorkerEvent(
+                state=WorkerState.IDLE,
+                task_status=TaskStatus(
+                    task_id="task_id",
+                    result=None,
+                    task_complete=True,
+                    task_failed=False,
+                ),
+            ),
+        )
+    )
+
+    with client.websocket_connect("/api/v2/run_plan") as ws_client:
+        ws_client.send_json(
+            {
+                "kind": "submit",
+                "task": {
+                    "name": "foo",
+                    "params": {"one": "two"},
+                    "instrument_session": "cm12345-1",
+                },
+            }
+        )
+        # first event is not sent
+        assert ws_client.receive_json() == {
+            "kind": "update",
+            "data": {
+                "state": "IDLE",
+                "task_status": {
+                    "task_id": "task_id",
+                    "result": None,
+                    "task_complete": True,
+                    "task_failed": False,
+                },
+                "errors": [],
+                "warnings": [],
+            },
+        }
+        with pytest.raises(WebSocketDisconnect) as discon:
+            ws_client.receive_text()
+
+        # Check it's a 'normal' end of stream disconnect
+        assert discon.value.code == 1000
+        assert discon.value.reason == ""
+    pass
+
+
+def test_websocket_run_plan_client_disconnect_cancels(
+    mock_runner: Mock, client: TestClient
+):
+    mock_runner.run.side_effect = ["task_id", None, None, None]
+    mock_runner.event_pipe.return_value = contextlib.nullcontext(
+        _aiter(
+            WorkerEvent(
+                state=WorkerState.IDLE,
+                task_status=TaskStatus(
+                    task_id="task_id",
+                    result=None,
+                    task_complete=False,
+                    task_failed=False,
+                ),
+            ),
+        )
+    )
+
+    class Disconnector:
+        def __init__(self, app):
+            self.app = app
+
+        async def __call__(self, scope: Scope, receive: Receive, send: Send):
+            async def local_send(message: Message):
+                if message.get("type") == "websocket.send":
+                    # Simulate the connection being closed
+                    raise OSError()
+                await send(message)
+
+            return await self.app(scope, receive, local_send)
+
+    cast(FastAPI, client.app).add_middleware(Disconnector)
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(
+            {
+                "kind": "submit",
+                "task": {
+                    "name": "foo",
+                    "params": {"one": "two"},
+                    "instrument_session": "cm12345-1",
+                },
+            }
+        )
+    mock_runner.run.assert_called_with(
+        interface.cancel_active_task, failure=True, reason="Client disconnected"
+    )
+
+
+async def _aiter(*values: Any) -> AsyncIterator:
+    for value in values:
+        yield value

From 6379f8b0d0b9260cee2acbf5023b2deb7907a4e2 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 31 Jul 2026 11:24:51 +0100
Subject: [PATCH 56/60] Pin uvicorn to 0.49 until websockets are fixed

---
 pyproject.toml | 2 +-
 uv.lock        | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e3a131e5c..aed5101bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,7 @@ dependencies = [
     "PyYAML>=6.0.2",
     "click>=8.2.0",
     "fastapi>=0.112.0",
-    "uvicorn",
+    "uvicorn==0.49", # Latest version that works with websockets https://github.com/Kludex/uvicorn/issues/3040
     "requests",
     "GitPython",
     "event-model==1.23.1",                           # https://github.com/DiamondLightSource/blueapi/issues/684
diff --git a/uv.lock b/uv.lock
index 5cda625b2..e044a86f6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -531,7 +531,7 @@ requires-dist = [
     { name = "stomp-py" },
     { name = "tiled", extras = ["client"], specifier = ">=0.2.4" },
     { name = "tomlkit" },
-    { name = "uvicorn" },
+    { name = "uvicorn", specifier = "==0.49" },
 ]
 
 [package.metadata.requires-dev]
@@ -6097,19 +6097,20 @@ wheels = [
 
 [[package]]
 name = "uvicorn"
-version = "0.51.0"
+version = "0.49.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "click" },
     { name = "h11" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a2/65/b7c6c443ccc58678c91e1e973bbe2a878591538655d6e1d47f24ba1c51f3/uvicorn-0.51.0.tar.gz", hash = "sha256:f6f4b69b657c312f516dd2d268ab9ae6f254b11e4bac504f37b2ab58b24dd0b0", size = 94412, upload-time = "2026-07-08T10:59:05.962Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/c4/1f/fa18009dea8469069cca78a4e877a008ab78f08b064bfc9ab891579077ff/uvicorn-0.49.0.tar.gz", hash = "sha256:ebf4271aa580d9de97f93192d4595176df6e91f9aae919ca73e4fc07df1e66a3", size = 91284, upload-time = "2026-06-03T22:01:30.448Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/45/ec/dbb7e5a6b91f86bfb9eb7d2988a2730907b6a729875b949c7f022e8b88fa/uvicorn-0.51.0-py3-none-any.whl", hash = "sha256:5d38af6cd620f2ae3849fb44fd4879e0890aa1febe8d47eb355fb45d93fe6a5b", size = 73219, upload-time = "2026-07-08T10:59:04.44Z" },
+    { url = "https://files.pythonhosted.org/packages/88/fa/e1388bbcf24ef3274f45c0c1c7b501fd14971037c1b6ee23610553307497/uvicorn-0.49.0-py3-none-any.whl", hash = "sha256:ba3d14c3ee7e41c6c654c46c9eb489d33213cdd30aa1696eab1374337c13f68f", size = 71376, upload-time = "2026-06-03T22:01:29.037Z" },
 ]
 
 [package.optional-dependencies]
 standard = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
     { name = "httptools" },
     { name = "python-dotenv" },
     { name = "pyyaml" },

From c714ea4e72deb9d4bc644d71603ff2646eee693c Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 31 Jul 2026 16:27:16 +0100
Subject: [PATCH 57/60] Add system test to run a plan via websocket

---
 tests/system_tests/test_blueapi_system.py | 40 ++++++++++++++++++-----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py
index c09a272e4..85ffcc885 100644
--- a/tests/system_tests/test_blueapi_system.py
+++ b/tests/system_tests/test_blueapi_system.py
@@ -4,6 +4,7 @@
 from collections.abc import Generator
 from contextlib import nullcontext
 from enum import StrEnum, auto
+from itertools import count
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -117,6 +118,27 @@ def instrument_session(request) -> str:
     return getattr(request, "param", VALID_INSTRUMENT_SESSION[User.alice])
 
 
+@pytest.fixture(scope="module")
+def local_numtracker():
+    """
+    Local version of what we expect the numtracker to provide
+
+    Allows different combinations of tests to be run without having to hard
+    code the expected scan numbers.
+    """
+    return count(CURRENT_NUMTRACKER_NUM + 1)
+
+
+@pytest.fixture
+def scan_id(local_numtracker) -> int:
+    """
+    The next value we expect to use for a scan_id
+
+    Should be used in any test that calls out to numtracker.
+    """
+    return next(local_numtracker)
+
+
 def task_factory(
     user: ValidUser, instrument_session: str | None, time: float = 0.0
 ) -> TaskRequest:
@@ -532,20 +554,17 @@ def test_delete_current_environment(client: BlueapiClient):
 
 
 @pytest.mark.parametrize(
-    "task,scan_id,user",
+    "task,user",
     [
         (
             TaskRequest(
                 name="count",
                 params={
-                    "detectors": [
-                        "det",
-                    ],
+                    "detectors": ["det"],
                     "num": 5,
                 },
                 instrument_session=VALID_INSTRUMENT_SESSION[User.alice],
             ),
-            CURRENT_NUMTRACKER_NUM + 1,
             User.alice,
         ),
         (
@@ -574,13 +593,17 @@ def test_delete_current_environment(client: BlueapiClient):
                 },
                 instrument_session=VALID_INSTRUMENT_SESSION[User.bob],
             ),
-            CURRENT_NUMTRACKER_NUM + 2,
             User.bob,
         ),
     ],
 )
+@pytest.mark.parametrize("run_method", ["run_task", "run_blocking"])
 def test_plan_runs(
-    client_with_stomp: BlueapiClient, task: TaskRequest, scan_id: int, user: ValidUser
+    client_with_stomp: BlueapiClient,
+    task: TaskRequest,
+    scan_id: int,
+    user: ValidUser,
+    run_method: str,
 ):
     resource = Queue(maxsize=1)
     start = Queue(maxsize=1)
@@ -592,7 +615,8 @@ def on_event(event: AnyEvent) -> None:
             if event.name == "stream_resource":
                 resource.put_nowait(event.doc)
 
-    final_event = client_with_stomp.run_task(task, on_event)
+    runner = getattr(client_with_stomp, run_method)
+    final_event = runner(task, on_event)
     assert isinstance(final_event.result, TaskResult)
     assert final_event.task_complete
     assert not final_event.task_failed

From 61e70dbbf7283c5db1a0b583db95bf0648273ee7 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 31 Jul 2026 17:28:49 +0100
Subject: [PATCH 58/60] Add authz to ws

---
 src/blueapi/client/rest.py           |  5 ++++
 src/blueapi/service/authorization.py |  5 ++--
 src/blueapi/service/main.py          | 34 +++++++++++++++++++++-------
 src/blueapi/service/protocol.py      | 11 ++++++++-
 4 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 3ab842488..7276bd26b 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -40,6 +40,7 @@
     PlanNotFound,
     ServerBusy,
     Submit,
+    Unauthorized,
     Update,
 )
 from blueapi.worker import TrackableTask, WorkerState
@@ -388,6 +389,10 @@ def run_blocking(
                             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:
diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py
index f9008138a..c324015a4 100644
--- a/src/blueapi/service/authorization.py
+++ b/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)):
diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py
index a5ac96fb6..c3e3137b0 100644
--- a/src/blueapi/service/main.py
+++ b/src/blueapi/service/main.py
@@ -31,10 +31,11 @@
 from opentelemetry.trace import get_tracer_provider
 from pydantic import ValidationError
 from starlette.responses import JSONResponse
+from starlette.status import WS_1007_INVALID_FRAME_PAYLOAD_DATA, WS_1013_TRY_AGAIN_LATER
 
 from blueapi.config import ApplicationConfig, OIDCConfig, Tag
 from blueapi.core.bluesky_types import DataEvent
-from blueapi.service import interface
+from blueapi.service import interface, protocol
 from blueapi.service.authentication import (
     Fedid,
     build_access_token_check,
@@ -45,10 +46,11 @@
     WebsocketTracing,
 )
 from blueapi.service.protocol import (
-    ControlRequest,
     InvalidArgs,
     PlanNotFound,
     ServerBusy,
+    Submit,
+    Unauthorized,
     Update,
 )
 from blueapi.worker import TrackableTask, WorkerState
@@ -627,19 +629,35 @@ def logout(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> Response:
 
 @secure_router_v2.websocket("/run_plan")
 async def run_plan(
-    ws: WebSocket, runner: Annotated[WorkerDispatcher, Depends(_runner)], user: Fedid
+    ws: WebSocket,
+    runner: Annotated[WorkerDispatcher, Depends(_runner)],
+    user: Fedid,
+    opa: Annotated[OpaUserClient | None, Depends(opa)],
 ):
     LOGGER.info("Starting WS plan as %s", user)
     await ws.accept()
     rq = await ws.receive_text()
     try:
-        task_request = ControlRequest.validate_json(rq)
+        task_request = Submit.model_validate_json(rq)
     except ValidationError:
         LOGGER.info("Failed to deserialize request: %r", rq, exc_info=True)
-        await ws.close(code=1007, reason="Invalid Request")
+        await ws.close(
+            code=WS_1007_INVALID_FRAME_PAYLOAD_DATA, reason="Invalid Request"
+        )
         return
     LOGGER.info("Plan request: %s", task_request)
 
+    if opa:
+        try:
+            await opa.can_submit_task(task_request.task)
+        except Exception as e:
+            LOGGER.info(
+                "User %s does not have permission to run task", user, exc_info=e
+            )
+            await ws.send_text(Unauthorized().model_dump_json())
+            await ws.close(code=protocol.AUTHZ_ERROR, reason="Unauthorized")
+            return
+
     try:
         task_id: str = runner.run(
             interface.submit_task, task_request.task, {"user": user}
@@ -648,12 +666,12 @@ async def run_plan(
     except ValidationError as ve:
         LOGGER.info("Plan args not valid: %s - %s", task_request, ve)
         await ws.send_text(InvalidArgs.from_validation_error(ve).model_dump_json())
-        await ws.close(code=4002, reason="Invalid Args")
+        await ws.close(code=protocol.INVALID_ARGS, reason="Invalid Args")
         return
     except KeyError as ke:
         LOGGER.info("Plan %r not recognised", ke.args[0])
         await ws.send_text(PlanNotFound(plan_name=ke.args[0]).model_dump_json())
-        await ws.close(code=4001, reason="Unknown Plan")
+        await ws.close(code=protocol.UNKNOWN_PLAN, reason="Unknown Plan")
         return
 
     try:
@@ -673,7 +691,7 @@ async def run_plan(
     except WorkerBusyError:
         LOGGER.error("Worker was busy")
         await ws.send_text(ServerBusy().model_dump_json())
-        await ws.close(code=1013, reason="Worker busy")
+        await ws.close(code=WS_1013_TRY_AGAIN_LATER, reason="Worker busy")
     except WebSocketDisconnect:
         LOGGER.info("Client disconnected")
         runner.run(
diff --git a/src/blueapi/service/protocol.py b/src/blueapi/service/protocol.py
index 67924318a..8af581913 100644
--- a/src/blueapi/service/protocol.py
+++ b/src/blueapi/service/protocol.py
@@ -23,6 +23,10 @@
 from blueapi.service.model import TaskRequest
 from blueapi.worker.event import ProgressEvent, WorkerEvent
 
+UNKNOWN_PLAN = 4001
+INVALID_ARGS = 4002
+AUTHZ_ERROR = 4003
+
 
 class ArgumentError(BaseModel):
     loc: list[str | int]
@@ -82,6 +86,10 @@ class ServerBusy(BaseModel):
     kind: Literal["busy"] = "busy"
 
 
+class Unauthorized(BaseModel):
+    kind: Literal["unauthorized"] = "unauthorized"
+
+
 class Update(BaseModel):
     kind: Literal["update"] = "update"
     data: WorkerEvent | DataEvent | ProgressEvent
@@ -89,6 +97,7 @@ class Update(BaseModel):
 
 ControlResponse = TypeAdapter(
     Annotated[
-        PlanNotFound | InvalidArgs | ServerBusy | Update, Field(discriminator="kind")
+        PlanNotFound | InvalidArgs | ServerBusy | Unauthorized | Update,
+        Field(discriminator="kind"),
     ]
 )

From 28b129f4707a7d3c0f415f1fb7fc22056c478d98 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 31 Jul 2026 18:01:26 +0100
Subject: [PATCH 59/60] Handle server being missing

---
 src/blueapi/client/rest.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py
index 7276bd26b..2ae87df65 100644
--- a/src/blueapi/client/rest.py
+++ b/src/blueapi/client/rest.py
@@ -399,6 +399,8 @@ def run_blocking(
                     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

From 36c8cdad7726b72bc4fdfbae552b462fc475b127 Mon Sep 17 00:00:00 2001
From: Peter Holloway <peter.holloway@diamond.ac.uk>
Date: Fri, 31 Jul 2026 18:24:51 +0100
Subject: [PATCH 60/60] Add websocket auth tests

---
 tests/system_tests/test_blueapi_system.py |  94 +++++++---------
 tests/unit_tests/service/test_rest_api.py | 131 ++++++++++++----------
 2 files changed, 112 insertions(+), 113 deletions(-)

diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py
index 85ffcc885..80f708651 100644
--- a/tests/system_tests/test_blueapi_system.py
+++ b/tests/system_tests/test_blueapi_system.py
@@ -2,7 +2,6 @@
 import time
 from asyncio import Queue
 from collections.abc import Generator
-from contextlib import nullcontext
 from enum import StrEnum, auto
 from itertools import count
 from pathlib import Path
@@ -693,56 +692,40 @@ def test_task_submission_after_invalid_task(client_with_stomp: BlueapiClient):
 
 
 @pytest.mark.parametrize(
-    "instrument_session,user,expectation",
+    "instrument_session,user",
     [
-        (
-            # bob cannot submit a task that alice is on
-            VALID_INSTRUMENT_SESSION[User.alice],
-            User.bob,
-            pytest.raises(
-                UnauthorisedAccessError, match="Not authorized to submit task"
-            ),
-        ),
-        (
-            # alice cannot submit a task that bob is on
-            VALID_INSTRUMENT_SESSION[User.bob],
-            User.alice,
-            pytest.raises(
-                UnauthorisedAccessError, match="Not authorized to submit task"
-            ),
-        ),
-        (
-            # alice can submit a task that alice is on
-            VALID_INSTRUMENT_SESSION[User.alice],
-            User.alice,
-            nullcontext(),
-        ),
-        (
-            # bob can submit a task that bob is on
-            VALID_INSTRUMENT_SESSION[User.bob],
-            User.bob,
-            nullcontext(),
-        ),
-        (
-            # admin can submit a task that bob is on
-            VALID_INSTRUMENT_SESSION[User.bob],
-            User.admin,
-            nullcontext(),
-        ),
-        (
-            # admin can submit a task that alice is on
-            VALID_INSTRUMENT_SESSION[User.alice],
-            User.admin,
-            nullcontext(),
-        ),
-        (
-            # admin still needs to put a valid instrument_session
-            INVALID_INSTRUMENT_SESSION,
-            User.admin,
-            pytest.raises(
-                UnauthorisedAccessError, match="Not authorized to submit task"
-            ),
-        ),
+        # bob cannot submit a task that alice is on
+        (VALID_INSTRUMENT_SESSION[User.alice], User.bob),
+        # alice cannot submit a task that bob is on
+        (VALID_INSTRUMENT_SESSION[User.bob], User.alice),
+        # admin still needs to put a valid instrument_session
+        (INVALID_INSTRUMENT_SESSION, User.admin),
+    ],
+)
+@pytest.mark.parametrize("method_name", ["create_task", "run_blocking"])
+def test_run_task_without_authz(
+    client: BlueapiClient,
+    small_task: TaskRequest,
+    user: ValidUser,
+    instrument_session: str,
+    method_name: str,
+):
+    method = getattr(client, method_name)
+    with pytest.raises(UnauthorisedAccessError, match="Not authorized to submit task"):
+        method(small_task)
+
+
+@pytest.mark.parametrize(
+    "instrument_session,user",
+    [
+        # alice can submit a task that alice is on
+        (VALID_INSTRUMENT_SESSION[User.alice], User.alice),
+        # bob can submit a task that bob is on
+        (VALID_INSTRUMENT_SESSION[User.bob], User.bob),
+        # admin can submit a task that bob is on
+        (VALID_INSTRUMENT_SESSION[User.bob], User.admin),
+        # admin can submit a task that alice is on
+        (VALID_INSTRUMENT_SESSION[User.alice], User.admin),
     ],
 )
 def test_create_task_authorization(
@@ -750,10 +733,8 @@ def test_create_task_authorization(
     small_task: TaskRequest,
     user: ValidUser,
     instrument_session: str,
-    expectation,
 ):
-    with expectation:
-        client.create_task(small_task)
+    client.create_task(small_task)
 
 
 def test_non_admin_can_only_get_own_tasks(
@@ -914,3 +895,10 @@ def test_admin_can_abort_any_task(
             task_factory(user, VALID_INSTRUMENT_SESSION[user], time=1)
         )
         client_factory[AdminUser.admin].abort()
+
+
+def test_run_blocking_requires_auth(
+    client_without_auth: BlueapiClient, small_task: TaskRequest
+):
+    with pytest.raises(UnauthorisedAccessError):
+        client_without_auth.run_blocking(small_task)
diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py
index 2f91ab686..a10d5ab3c 100644
--- a/tests/unit_tests/service/test_rest_api.py
+++ b/tests/unit_tests/service/test_rest_api.py
@@ -14,6 +14,7 @@
 from httpx2 import Headers
 from pydantic import BaseModel, ValidationError
 from pydantic_core import InitErrorDetails
+from starlette.testclient import WebSocketDenialResponse
 from starlette.types import Message, Receive, Scope, Send
 
 from blueapi.config import (
@@ -57,6 +58,16 @@ class MockCountModel(BaseModel): ...
 FAKE_INSTRUMENT_SESSION = "cm12345-1"
 
 
+SUBMIT_REQUEST = {
+    "kind": "submit",
+    "task": {
+        "name": "foo",
+        "params": {"one": "two"},
+        "instrument_session": "cm12345-1",
+    },
+}
+
+
 @pytest.fixture
 def mock_runner() -> Mock:
     return MagicMock(spec=WorkerDispatcher)
@@ -1028,18 +1039,9 @@ async def test_websocket_run_plan(mock_runner: Mock, client: TestClient):
         )
     )
 
-    with client.websocket_connect("/api/v2/run_plan") as ws_client:
-        ws_client.send_json(
-            {
-                "kind": "submit",
-                "task": {
-                    "name": "foo",
-                    "params": {"one": "two"},
-                    "instrument_session": "cm12345-1",
-                },
-            }
-        )
-        assert ws_client.receive_json() == {
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(SUBMIT_REQUEST)
+        assert ws.receive_json() == {
             "kind": "update",
             "data": {
                 "state": "RUNNING",
@@ -1053,7 +1055,7 @@ async def test_websocket_run_plan(mock_runner: Mock, client: TestClient):
                 "warnings": [],
             },
         }
-        assert ws_client.receive_json() == {
+        assert ws.receive_json() == {
             "kind": "update",
             "data": {
                 "state": "IDLE",
@@ -1068,7 +1070,7 @@ async def test_websocket_run_plan(mock_runner: Mock, client: TestClient):
             },
         }
         with pytest.raises(WebSocketDisconnect) as discon:
-            ws_client.receive_text()
+            ws.receive_text()
 
         # Check it's a 'normal' end of stream disconnect
         assert discon.value.code == 1000
@@ -1079,10 +1081,10 @@ async def test_websocket_run_plan(mock_runner: Mock, client: TestClient):
 def test_websocket_run_plan_invalid_request(
     req: str, mock_runner: Mock, client: TestClient
 ):
-    with client.websocket_connect("/api/v2/run_plan") as ws_client:
-        ws_client.send_text(req)
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_text(req)
         with pytest.raises(WebSocketDisconnect) as disco:
-            ws_client.receive_text()
+            ws.receive_text()
         assert disco.value.code == 1007
         assert disco.value.reason == "Invalid Request"
 
@@ -1092,10 +1094,7 @@ def test_websocket_run_plan_invalid_request(
     [
         (
             KeyError("foo"),
-            {
-                "kind": "plan_not_found",
-                "plan_name": "foo",
-            },
+            {"kind": "plan_not_found", "plan_name": "foo"},
             4001,
             "Unknown Plan",
         ),
@@ -1108,19 +1107,19 @@ def test_websocket_run_plan_invalid_request(
     ],
 )
 def test_websocket_run_plan_submit_error(
-    exc, err_message, code, reason, mock_runner: Mock, client: TestClient
+    exc: Exception,
+    err_message: str,
+    code: int,
+    reason,
+    mock_runner: Mock,
+    client: TestClient,
 ):
     mock_runner.run.side_effect = exc
-    with client.websocket_connect("/api/v2/run_plan") as ws_client:
-        ws_client.send_json(
-            {
-                "kind": "submit",
-                "task": {"name": "foo", "instrument_session": "cm12345-1"},
-            }
-        )
-        assert ws_client.receive_json() == err_message
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(SUBMIT_REQUEST)
+        assert ws.receive_json() == err_message
         with pytest.raises(WebSocketDisconnect) as disco:
-            ws_client.receive_text()
+            ws.receive_text()
         assert disco.value.code == code
         assert disco.value.reason == reason
 
@@ -1130,13 +1129,11 @@ def test_websocket_run_plan_server_busy(mock_runner: Mock, client: TestClient):
         "task_id",  # submit_task
         Mock(name="active_task", is_complete=False),
     ]
-    with client.websocket_connect("/api/v2/run_plan") as ws_client:
-        ws_client.send_json(
-            {"kind": "submit", "task": {"name": "foo", "instrument_session": "cm123-1"}}
-        )
-        assert ws_client.receive_json() == {"kind": "busy"}
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(SUBMIT_REQUEST)
+        assert ws.receive_json() == {"kind": "busy"}
         with pytest.raises(WebSocketDisconnect) as disco:
-            ws_client.receive_text()
+            ws.receive_text()
         assert disco.value.code == 1013
         assert disco.value.reason == "Worker busy"
     pass
@@ -1171,19 +1168,10 @@ def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClie
         )
     )
 
-    with client.websocket_connect("/api/v2/run_plan") as ws_client:
-        ws_client.send_json(
-            {
-                "kind": "submit",
-                "task": {
-                    "name": "foo",
-                    "params": {"one": "two"},
-                    "instrument_session": "cm12345-1",
-                },
-            }
-        )
+    with client.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(SUBMIT_REQUEST)
         # first event is not sent
-        assert ws_client.receive_json() == {
+        assert ws.receive_json() == {
             "kind": "update",
             "data": {
                 "state": "IDLE",
@@ -1198,7 +1186,7 @@ def test_websocket_run_plan_unrelated_events(mock_runner: Mock, client: TestClie
             },
         }
         with pytest.raises(WebSocketDisconnect) as discon:
-            ws_client.receive_text()
+            ws.receive_text()
 
         # Check it's a 'normal' end of stream disconnect
         assert discon.value.code == 1000
@@ -1239,21 +1227,44 @@ async def local_send(message: Message):
 
     cast(FastAPI, client.app).add_middleware(Disconnector)
     with client.websocket_connect("/api/v2/run_plan") as ws:
-        ws.send_json(
-            {
-                "kind": "submit",
-                "task": {
-                    "name": "foo",
-                    "params": {"one": "two"},
-                    "instrument_session": "cm12345-1",
-                },
-            }
-        )
+        ws.send_json(SUBMIT_REQUEST)
     mock_runner.run.assert_called_with(
         interface.cancel_active_task, failure=True, reason="Client disconnected"
     )
 
 
+@pytest.mark.parametrize("token", ["Bearer invalid", None])
+def test_websocket_run_plan_needs_auth_token(
+    client_with_auth: TestClient, token: str | None
+):
+    del client_with_auth.headers["Authorization"]
+    if token:
+        client_with_auth.headers["Authorization"] = token
+    with pytest.raises(WebSocketDenialResponse) as wsdr:
+        with client_with_auth.websocket_connect("/api/v2/run_plan"):
+            pass
+    assert wsdr.value.status_code == 401
+
+
+def test_websocket_run_plan_needs_permission(
+    mock_runner: Mock,
+    client_with_opa: TestClient,
+    mock_opa_client: Mock,
+    access_token: str,
+):
+    client_with_opa.headers["Authorization"] = f"Bearer {access_token}"
+    mock_opa_client.can_submit_task.side_effect = HTTPException(status_code=403)
+    mock_runner.run.side_effect = RuntimeError("Task should not be submitted")
+    with client_with_opa.websocket_connect("/api/v2/run_plan") as ws:
+        ws.send_json(SUBMIT_REQUEST)
+        assert ws.receive_json() == {"kind": "unauthorized"}
+        with pytest.raises(WebSocketDisconnect) as discon:
+            ws.receive_text()
+
+        assert discon.value.code == 4003
+        assert discon.value.reason == "Unauthorized"
+
+
 async def _aiter(*values: Any) -> AsyncIterator:
     for value in values:
         yield value