Update bluesky-httpserver to support RE watcher updates over the progress websocket

bluesky_httpserver/app.py

@@ -16,7 +16,7 @@
 from fastapi.openapi.utils import get_openapi
 
 from .authentication import Mode
-from .console_output import CollectPublishedConsoleOutput, ConsoleOutputStream, SystemInfoStream
+from .console_output import CollectPublishedConsoleOutput, ConsoleOutputStream, ProgressStream, SystemInfoStream
 from .core import PatchedStreamingResponse
 from .database.core import purge_expired
 from .resources import SERVER_RESOURCES as SR
@@ -351,6 +351,8 @@ async def purge_expired_sessions_and_api_keys():
         SR.console_output_loader.subscribe(SR.console_output_stream.add_message)
         SR.set_system_info_stream(SystemInfoStream(rm_ref=RM))
         SR.system_info_stream.start()
+        SR.set_progress_stream(ProgressStream(rm_ref=RM))
+        SR.progress_stream.start()
 
         # Import module with custom code
         module_names_str = os.getenv("QSERVER_CUSTOM_MODULES", None)
@@ -394,6 +396,7 @@ async def shutdown_event():
         await SR.console_output_loader.stop()
         await SR.console_output_stream.stop()
         await SR.system_info_stream.stop()
+        await SR.progress_stream.stop()
 
     @lru_cache(1)
     def override_get_authenticators():

bluesky_httpserver/console_output.py

@@ -330,3 +330,72 @@ def start(self):
     async def stop(self):
         await self._stop_background_task()
         await self._RM.system_info_monitor.disable_wait()
+
+
+class ProgressStream:
+    def __init__(self, *, rm_ref):
+        self._RM = rm_ref
+        self._queues = {}
+        self._background_task = None
+        self._background_task_running = False
+        self._background_task_stopped = asyncio.Event()
+        self._background_task_stopped.set()
+        self._queue_max_size = 1000
+
+    @property
+    def background_task_running(self):
+        return self._background_task_running
+
+    @property
+    def queues(self):
+        return self._queues
+
+    def add_queue(self, key):
+        """
+        Add a new queue to the dictionary of queues. The key is a reference to the socket for
+        for connection with the client.
+        """
+        queue = asyncio.Queue(maxsize=self._queue_max_size)
+        self._queues[key] = queue
+        return queue
+
+    def remove_queue(self, key):
+        """
+        Remove the queue identified by the key from the dictionary of queues.
+        """
+        if key in self._queues:
+            del self._queues[key]
+
+    def _start_background_task(self):
+        if not self._background_task_running:
+            self._background_task = asyncio.create_task(self._load_msgs_task())
+
+    async def _stop_background_task(self):
+        self._background_task_running = False
+        await self._background_task_stopped.wait()
+
+    async def _load_msgs_task(self):
+        self._background_task_stopped.clear()
+        self._background_task_running = True
+        while self._background_task_running:
+            try:
+                msg = await self._RM.progress_monitor.next_msg(timeout=0.5)
+
+                if isinstance(msg, dict) and "msg" in msg:
+                    msg_json = json.dumps(msg)
+                    for q in self._queues.values():
+                        # Protect from overflow. It's ok to discard old messages.
+                        if q.full():
+                            q.get_nowait()
+                        await q.put(msg_json)
+            except self._RM.RequestTimeoutError:
+                pass
+        self._background_task_stopped.set()
+
+    def start(self):
+        self._RM.progress_monitor.enable()
+        self._start_background_task()
+
+    async def stop(self):
+        await self._stop_background_task()
+        await self._RM.progress_monitor.disable_wait()

bluesky_httpserver/resources.py

@@ -60,5 +60,16 @@ def system_info_stream(self):
     def system_info_stream(self, _):
         raise RuntimeError("Attempting to set read-only property 'system_info_stream'")
 
+    def set_progress_stream(self, progress_stream):
+        self._progress_stream = progress_stream
+
+    @property
+    def progress_stream(self):
+        return self._progress_stream
+
+    @progress_stream.setter
+    def progress_stream(self, _):
+        raise RuntimeError("Attempting to set read-only property 'progress_stream'")
+
 
 SERVER_RESOURCES = _ServerResources()

bluesky_httpserver/routers/core_api.py

@@ -1215,3 +1215,29 @@ async def info_ws(websocket: WebSocket, scopes=["read:monitor"]):
         pass
     finally:
         SR.system_info_stream.remove_queue_info(websocket)
+
+
+@router.websocket("/progress/ws")
+async def progress_ws(websocket: WebSocket, scopes=["read:monitor"]):
+    principal = get_current_principal_websocket(websocket=websocket, scopes=scopes)
+    if not principal:
+        await websocket.close(code=4001, reason="Invalid token")
+        return
+
+    await websocket.accept()
+    q = SR.progress_stream.add_queue(websocket)
+    wsmon = WebSocketMonitor(websocket)
+    wsmon.start()
+    try:
+        while wsmon.is_alive:
+            try:
+                msg = await asyncio.wait_for(q.get(), timeout=1)
+                await websocket.send_text(msg)
+            except asyncio.TimeoutError:
+                pass
+            except RuntimeError:  # 'send' after the client is disconnected
+                pass
+    except WebSocketDisconnect:
+        pass
+    finally:
+        SR.progress_stream.remove_queue(websocket)

bluesky_httpserver/tests/test_progress_socket.py

@@ -0,0 +1,211 @@
+import json
+import pprint
+import threading
+import time as ttime
+
+import pytest
+from bluesky_queueserver.manager.tests.common import (  # noqa F401
+    append_code_to_last_startup_file,
+    copy_default_profile_collection,
+    re_manager_cmd,
+)
+from websockets.sync.client import connect
+
+from bluesky_httpserver.tests.conftest import (  # noqa F401
+    API_KEY_FOR_TESTS,
+    SERVER_ADDRESS,
+    SERVER_PORT,
+    fastapi_server_fs,
+    request_to_json,
+    set_qserver_zmq_encoding,
+    wait_for_environment_to_be_closed,
+    wait_for_environment_to_be_created,
+    wait_for_queue_execution_to_complete,
+)
+
+# Startup code that defines a slow (delayed) motor and a plan that moves it. Moving a delayed
+# motor makes the RE 'waiting_hook' emit incremental watcher updates (name/current/fraction/...),
+# as opposed to the instantaneous demo motors which only produce completion messages.
+_progress_test_startup_code = """
+from ophyd.sim import SynAxis as _SynAxisProgressTest
+from bluesky.plan_stubs import mv as _mv_progress_test
+
+motor_slow = _SynAxisProgressTest(name="motor_slow", delay=0.6)
+
+
+def progress_test_plan(npts: int = 5):
+    for _i in range(npts):
+        yield from _mv_progress_test(motor_slow, _i + 1)
+"""
+
+
+class _ReceiveProgressSocket(threading.Thread):
+    """
+    Catch streaming progress updates by connecting to the ``/progress/ws`` socket and
+    save messages to the buffer.
+    """
+
+    def __init__(self, *, endpoint="/progress/ws", api_key=API_KEY_FOR_TESTS, **kwargs):
+        super().__init__(**kwargs)
+        self.received_data_buffer = []
+        self._exit = False
+        self._api_key = api_key
+        self._endpoint = endpoint
+
+    def run(self):
+        websocket_uri = f"ws://{SERVER_ADDRESS}:{SERVER_PORT}/api{self._endpoint}"
+        additional_headers = {"Authorization": f"ApiKey {self._api_key}"}
+        try:
+            with connect(websocket_uri, additional_headers=additional_headers) as websocket:
+                while not self._exit:
+                    try:
+                        msg_json = websocket.recv(timeout=0.1, decode=False)
+                        try:
+                            msg = json.loads(msg_json)
+                            self.received_data_buffer.append(msg)
+                        except json.JSONDecodeError:
+                            pass
+                    except TimeoutError:
+                        pass
+        except Exception as ex:
+            print(f"Failed to connect to server: {ex}")
+
+    def stop(self):
+        self._exit = True
+
+    def __del__(self):
+        self.stop()
+
+
+@pytest.mark.parametrize("zmq_port", (None, 60619))
+def test_http_server_progress_socket_1(
+    monkeypatch, re_manager_cmd, fastapi_server_fs, zmq_port  # noqa F811
+):
+    """
+    Test for the ``/progress/ws`` websocket. Runs a plan that moves a motor so that the
+    RE Manager publishes ``waiting_hook`` (watcher) progress updates over 0MQ, and verifies
+    that the updates are streamed to the connected websocket client.
+    """
+    # Start HTTP Server
+    if zmq_port is not None:
+        monkeypatch.setenv("QSERVER_ZMQ_INFO_ADDRESS", f"tcp://localhost:{zmq_port}")
+    fastapi_server_fs()
+
+    # Start RE Manager with progress publishing enabled
+    params = ["--zmq-publish-progress", "ON"]
+    if zmq_port is not None:
+        params.extend(["--zmq-info-addr", f"tcp://*:{zmq_port}"])
+    re_manager_cmd(params)
+
+    rps = _ReceiveProgressSocket()
+    rps.start()
+    ttime.sleep(1)  # Wait until the client connects to the socket
+
+    resp1 = request_to_json("post", "/environment/open")
+    assert resp1["success"] is True, pprint.pformat(resp1)
+    assert wait_for_environment_to_be_created(timeout=10)
+
+    # A 'scan' moves 'motor', which makes the RE wait on status objects and publish progress updates
+    plan = {"name": "scan", "args": [["det"], "motor", -1, 1, 10], "item_type": "plan"}
+    resp2 = request_to_json("post", "/queue/item/add", json={"item": plan})
+    assert resp2["success"] is True, pprint.pformat(resp2)
+
+    resp3 = request_to_json("post", "/queue/start")
+    assert resp3["success"] is True, pprint.pformat(resp3)
+
+    assert wait_for_queue_execution_to_complete(timeout=30)
+
+    resp4 = request_to_json("post", "/environment/close")
+    assert resp4["success"] is True, pprint.pformat(resp4)
+    assert wait_for_environment_to_be_closed(timeout=10)
+
+    # Wait until capture is complete
+    ttime.sleep(2)
+    rps.stop()
+    rps.join()
+
+    buffer = rps.received_data_buffer
+    assert len(buffer) > 0, "No progress updates were received"
+    for msg in buffer:
+        assert "time" in msg, msg
+        assert isinstance(msg["time"], float), msg
+        assert "msg" in msg, msg
+        assert isinstance(msg["msg"], dict), msg
+
+    # The RE sends a completion message each time it finishes waiting on status objects.
+    completion_msgs = [_ for _ in buffer if _["msg"].get("completed") is True]
+    assert len(completion_msgs) > 0, pprint.pformat(buffer)
+
+
+def test_http_server_progress_socket_2(
+    tmp_path, monkeypatch, re_manager_cmd, fastapi_server_fs  # noqa F811
+):
+    """
+    Test that incremental (watcher) progress updates are streamed over ``/progress/ws``.
+    Runs a plan that moves a delayed motor, which makes the RE Manager publish per-status
+    progress updates (with ``name``/``current``/``fraction``) and not just completion messages.
+    """
+    # Prepare a startup profile with a delayed motor and a plan that moves it
+    pc_path = copy_default_profile_collection(tmp_path)
+    append_code_to_last_startup_file(pc_path, _progress_test_startup_code)
+
+    fastapi_server_fs()
+
+    # 'ENVIRONMENT_OPEN' regenerates the list of existing plans/devices so 'progress_test_plan' is known
+    params = [
+        "--zmq-publish-progress",
+        "ON",
+        "--startup-dir",
+        pc_path,
+        "--update-existing-plans-devices",
+        "ENVIRONMENT_OPEN",
+    ]
+    re_manager_cmd(params)
+
+    rps = _ReceiveProgressSocket()
+    rps.start()
+    ttime.sleep(1)  # Wait until the client connects to the socket
+
+    resp1 = request_to_json("post", "/environment/open")
+    assert resp1["success"] is True, pprint.pformat(resp1)
+    assert wait_for_environment_to_be_created(timeout=10)
+
+    plan = {"name": "progress_test_plan", "kwargs": {"npts": 5}, "item_type": "plan"}
+    resp2 = request_to_json("post", "/queue/item/add", json={"item": plan})
+    assert resp2["success"] is True, pprint.pformat(resp2)
+
+    resp3 = request_to_json("post", "/queue/start")
+    assert resp3["success"] is True, pprint.pformat(resp3)
+
+    assert wait_for_queue_execution_to_complete(timeout=30)
+
+    resp4 = request_to_json("post", "/environment/close")
+    assert resp4["success"] is True, pprint.pformat(resp4)
+    assert wait_for_environment_to_be_closed(timeout=10)
+
+    # Wait until capture is complete
+    ttime.sleep(2)
+    rps.stop()
+    rps.join()
+
+    buffer = rps.received_data_buffer
+    assert len(buffer) > 0, "No progress updates were received"
+    for msg in buffer:
+        assert "time" in msg, msg
+        assert isinstance(msg["time"], float), msg
+        assert "msg" in msg, msg
+        assert isinstance(msg["msg"], dict), msg
+
+    # Incremental watcher updates carry the moved device name (not just {"completed": True})
+    incremental_msgs = [_ for _ in buffer if _["msg"].get("name") == "motor_slow"]
+    assert len(incremental_msgs) > 0, pprint.pformat(buffer)
+    for msg in incremental_msgs:
+        m = msg["msg"]
+        assert "current" in m, msg
+        assert "fraction" in m, msg
+        assert "done" in m, msg
+
+    # Completion messages are still sent when each wait finishes
+    completion_msgs = [_ for _ in buffer if _["msg"].get("completed") is True]
+    assert len(completion_msgs) > 0, pprint.pformat(buffer)
+