fix(waveform): track async subscriptions per scan and device

bec_widgets/widgets/plots/waveform/waveform.py

@@ -148,7 +148,8 @@ class Waveform(PlotBase):
         self._async_curves = []
         # Async streams already subscribed in the current scan. Keyed by
         # (device, stream) (stream is None for the old per-device endpoint); the value
-        # is the list of signal entries that share that single subscription.
+        # holds the signal entries that share that single subscription ("signals") and the
+        # exact "endpoint" and "cb_info" it was registered with.
         self._async_streams_setup: dict = {}
         self._history_curves = []
         self._slice_index = None
@@ -1656,26 +1657,41 @@ class Waveform(PlotBase):
             self.update_scan_info_from_source(self.scan_item)
             self._slice_index = None  # Reset the slice index
             self._update_curve_visibility()
-            self._mode = self._categorise_device_curves()
-
-            # First trigger to sync and async data. Async curve subscriptions
-            # were already set up by _categorise_device_curves above (it calls
-            # _setup_async_curve per async curve for live scans); repeating it
-            # here cleared each curve's data a second time and re-attempted
-            # the stream registration for every new scan.
-            if self._mode == "sync":
-                self.sync_signal_update.emit()
-                logger.info("Scan status: Sync mode")
-            elif self._mode == "async":
-                self.async_signal_update.emit()
-                logger.info("Scan status: Async mode")
-            else:
-                self.sync_signal_update.emit()
-                self.async_signal_update.emit()
-                logger.info("Scan status: Mixed mode")
-                logger.warning("Mixed mode - integrity of x axis cannot be guaranteed.")
+            self._start_scan_updates()
+        elif self.scan_item is None:
+            # The first status of this scan arrived before its ScanItem existed, so the curves
+            # could not be categorised and no live async subscription was set up. Do it as
+            # soon as the ScanItem is available, although the scan ID did not change.
+            self.scan_item = self.queue.scan_storage.find_scan_by_ID(self.scan_id)
+            if self.scan_item is not None:
+                self.update_scan_info_from_source(self.scan_item)
+                self._start_scan_updates()
         self.setup_dap_for_scan()
 
+    def _start_scan_updates(self):
+        """
+        Categorise the device curves for the current scan, which also sets up the live async
+        subscriptions, and trigger the first sync and async updates.
+        """
+        self._mode = self._categorise_device_curves()
+
+        # First trigger to sync and async data. Async curve subscriptions
+        # were already set up by _categorise_device_curves above (it calls
+        # _setup_async_curve per async curve for live scans); repeating it
+        # here cleared each curve's data a second time and re-attempted
+        # the stream registration for every new scan.
+        if self._mode == "sync":
+            self.sync_signal_update.emit()
+            logger.info("Scan status: Sync mode")
+        elif self._mode == "async":
+            self.async_signal_update.emit()
+            logger.info("Scan status: Async mode")
+        else:
+            self.sync_signal_update.emit()
+            self.async_signal_update.emit()
+            logger.info("Scan status: Mixed mode")
+            logger.warning("Mixed mode - integrity of x axis cannot be guaranteed.")
+
     @SafeSlot(dict, dict)
     def on_scan_progress(self, msg: dict, meta: dict):
         """
@@ -1842,7 +1858,6 @@ class Waveform(PlotBase):
         if async_signal_found:
             stream_key = (name, stream)  # AsyncMultiSignal sub-signals share `stream`
             new_endpoint = MessageEndpoints.device_async_signal(self.scan_id, name, stream)
-            old_endpoint = MessageEndpoints.device_async_signal(self.old_scan_id, name, stream)
             legacy_endpoint = False
         else:
             # No BEC async signal (AsyncSignal/AsyncMultiSignal/DynamicSignal) was found
@@ -1850,7 +1865,6 @@ class Waveform(PlotBase):
             # device_async_readback endpoint.
             stream_key = (name, None)  # old endpoint is keyed by device only
             new_endpoint = MessageEndpoints.device_async_readback(self.scan_id, name)
-            old_endpoint = MessageEndpoints.device_async_readback(self.old_scan_id, name)
             legacy_endpoint = True
 
         endpoint_str = getattr(new_endpoint, "endpoint", new_endpoint)
@@ -1858,9 +1872,9 @@ class Waveform(PlotBase):
         shared = self._async_streams_setup.get(stream_key)
         if shared is not None:
             # Another curve already subscribed to this exact stream this scan.
-            shared.append(signal)
+            shared["signals"].append(signal)
             logger.info(
-                f"Async signals {shared} share a single subscription on endpoint "
+                f"Async signals {shared['signals']} share a single subscription on endpoint "
                 f"'{endpoint_str}'; reusing it instead of subscribing again."
             )
             return
@@ -1875,13 +1889,66 @@ class Waveform(PlotBase):
         #         "waveform widget will be removed in a future release."
         #     )
 
-        self._async_streams_setup[stream_key] = [signal]
-        self.bec_dispatcher.disconnect_slot(self.on_async_readback, old_endpoint)
+        # Record the exact endpoint and subscription context next to the subscribers: the
+        # subscription is released with them (_release_async_streams), whatever scan_id holds by
+        # then, and on_async_readback matches readbacks to the curves of this stream.
+        cb_info = {"scan_id": self.scan_id, "device": name, "stream": stream_key[1]}
+        self._async_streams_setup[stream_key] = {
+            "signals": [signal],
+            "endpoint": new_endpoint,
+            "cb_info": cb_info,
+        }
         self.bec_dispatcher.connect_slot(
-            self.on_async_readback, new_endpoint, from_start=True, cb_info={"scan_id": self.scan_id}
+            self.on_async_readback, new_endpoint, from_start=True, cb_info=dict(cb_info)
         )
         logger.info(f"Setup async curve {name} (signal '{signal}', endpoint '{endpoint_str}')")
 
+    def _release_async_streams(self, streams: dict) -> None:
+        """
+        Disconnect async stream subscriptions that the current curve categorisation no longer
+        uses, e.g. those of the previous scan or of a live scan after switching to history.
+
+        Each subscription is released with the exact endpoint and cb_info it was registered
+        with, so it is found regardless of the scan IDs tracked by the widget by now.
+
+        Args:
+            streams(dict): Previous ``_async_streams_setup`` records.
+        """
+        for stream_key, stream in streams.items():
+            current = self._async_streams_setup.get(stream_key)
+            if current is not None and current["cb_info"] == stream["cb_info"]:
+                continue  # the same subscription is still in use
+            self.bec_dispatcher.disconnect_slot(
+                self.on_async_readback, stream["endpoint"], cb_info=stream["cb_info"]
+            )
+
+    def _async_stream_curves(self, cb_info: dict) -> list[Curve]:
+        """
+        Get the async curves fed by the subscription described by cb_info.
+
+        Devices can expose signals with the same entry name, so curves are matched by device
+        and stream, not only by the signal entries of the readback.
+
+        Args:
+            cb_info(dict): Subscription context set up by _setup_async_curve.
+
+        Returns:
+            list[Curve]: The curves of the stream; all async curves if cb_info carries no
+                device, in which case the readback is matched by signal entry only.
+        """
+        device = cb_info.get("device")
+        if device is None:
+            return self._async_curves
+        stream = self._async_streams_setup.get((device, cb_info.get("stream")))
+        if stream is None:
+            return []  # the stream is no longer subscribed
+        return [
+            curve
+            for curve in self._async_curves
+            if curve.config.signal.device == device
+            and curve.config.signal.signal in stream["signals"]
+        ]
+
     @SafeSlot(dict, dict)
     def on_async_readback(self, msg, metadata):
         """
@@ -1899,14 +1966,15 @@ class Waveform(PlotBase):
 
         Args:
             msg(dict): Message with the async data.
-            metadata(dict): Message metadata, including the subscription's scan ID in ``cb_info``.
+            metadata(dict): Message metadata. ``metadata["cb_info"]`` is the subscription
+                context set up by _setup_async_curve (scan ID, device and stream).
         """
         cb_info = metadata.get("cb_info")
         if not isinstance(cb_info, dict):
-            logger.info("Async readback has no subscription context.")
+            logger.info("Async readback has no subscription context, ignoring it.")
             return
-        scan_id = cb_info.get("scan_id")
-        if scan_id != self.scan_id:
+        if cb_info.get("scan_id") != self.scan_id:
+            # A late readback of a previous scan must not overwrite the current scan's data
             logger.info("Scan ID mismatch, ignoring async readback.")
             return
 
@@ -1916,7 +1984,7 @@ class Waveform(PlotBase):
             return
         max_shape = metadata.get("async_update", {}).get("max_shape", [])
         plot_mode = self.x_axis_mode["name"]
-        for curve in self._async_curves:
+        for curve in self._async_stream_curves(cb_info):
             x_data = None  # Reset x_data
             y_data = None  # Reset y_data
             # Get the curve data
@@ -2437,6 +2505,7 @@ class Waveform(PlotBase):
         # Reset sync/async curve lists
         self._async_curves.clear()
         self._sync_curves.clear()
+        previous_async_streams = self._async_streams_setup
         self._async_streams_setup = {}
         found_async = False
         found_sync = False
@@ -2465,6 +2534,9 @@ class Waveform(PlotBase):
                     "device not found in readout priority list."
                 )
                 continue
+        # Subscriptions of the previous source that were not set up again (another scan, a
+        # history scan, removed curves) would keep feeding stale readbacks; release them.
+        self._release_async_streams(previous_async_streams)
         # Determine the mode of the scan
         if found_async and found_sync:
             mode = "mixed"

tests/unit_tests/test_waveform.py

@@ -622,6 +622,176 @@ def test_on_scan_status_ignored_without_device_curves(qtbot, mocked_client_with_
     assert setup_dap_spy.call_count == calls_before
 
 
+def _async_readback_subscriptions(wf) -> list[tuple[str, object]]:
+    """(endpoint, dispatcher wrapper) of every subscription of wf.on_async_readback."""
+    subscriptions = [
+        (topic, slot)
+        for slot in wf.bec_dispatcher._registered_slots.values()
+        if slot.cb == wf.on_async_readback
+        for topic in slot.topics
+    ]
+    return sorted(subscriptions, key=lambda subscription: subscription[0])
+
+
+def _async_readback_endpoints(wf) -> list[str]:
+    return [endpoint for endpoint, _ in _async_readback_subscriptions(wf)]
+
+
+def _deliver_async_readback(wf, endpoint: str, values: dict, update_type: str = "replace"):
+    """Deliver an async message like the dispatcher: through the wrapper subscribed to endpoint."""
+    [wrapper] = [slot for topic, slot in _async_readback_subscriptions(wf) if topic == endpoint]
+    msg = {"signals": {signal: {"value": value} for signal, value in values.items()}}
+    wrapper(msg, {"async_update": {"type": update_type}})
+
+
+def _async_readback_endpoint(scan_id: str, device: str) -> str:
+    from bec_lib.endpoints import MessageEndpoints
+
+    return MessageEndpoints.device_async_readback(scan_id, device).endpoint
+
+
+def test_async_subscription_bound_when_scan_item_arrives_after_scan_status(
+    qtbot, mocked_client, monkeypatch
+):
+    """
+    Regression for #1297: the first scan status arrived before the ScanItem existed, so no
+    async subscription could be set up. A later status of the same scan must bind the live
+    async endpoint, so that async readbacks reach the plot.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    curve = wf.plot(arg1="async_device")
+    scan_storage = {}
+    monkeypatch.setattr(
+        wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id)
+    )
+
+    wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {})
+    assert wf.scan_id == "scan-1"
+    assert _async_readback_endpoints(wf) == []
+
+    scan_storage["scan-1"] = create_dummy_scan_item(scan_id="scan-1")
+    wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {})
+
+    endpoint = _async_readback_endpoint("scan-1", "async_device")
+    assert _async_readback_endpoints(wf) == [endpoint]
+    _deliver_async_readback(wf, endpoint, {"async_device": [4, 5, 6]})
+    np.testing.assert_array_equal(curve.get_data()[1], [4, 5, 6])
+
+
+def test_async_subscription_follows_live_history_live_switch(
+    qtbot, mocked_client, monkeypatch, scan_history_factory
+):
+    """
+    Regression for #1298: switching from a live scan to a history scan must release the live
+    async subscription (registered under the live scan ID), and the next live scan must end up
+    with exactly one subscription, its own.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    wf.plot(arg1="async_device")
+    scan_storage = {
+        scan_id: create_dummy_scan_item(scan_id=scan_id) for scan_id in ("live-1", "live-2")
+    }
+    monkeypatch.setattr(
+        wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id)
+    )
+    inject_scan_history(wf, scan_history_factory, ("hist-1", 1))
+    history_item = wf.client.history.get_by_scan_id("hist-1")
+    # Stored scans carry the readout priority of the scan in their BEC metadata
+    history_item.metadata["bec"]["readout_priority"] = {
+        "monitored": ["bpm4i"],
+        "async": ["async_device"],
+    }
+    monkeypatch.setattr(wf.client.history, "get_by_scan_id", lambda scan_id: history_item)
+
+    wf.on_scan_status({"scan_id": "live-1", "status": "open"}, {})
+    assert _async_readback_endpoints(wf) == [_async_readback_endpoint("live-1", "async_device")]
+
+    wf.update_with_scan_history(scan_id="hist-1")
+    assert wf.scan_id == "hist-1"
+    assert _async_readback_endpoints(wf) == []
+
+    wf.on_scan_status({"scan_id": "live-2", "status": "open"}, {})
+    assert _async_readback_endpoints(wf) == [_async_readback_endpoint("live-2", "async_device")]
+
+
+def test_async_readback_of_previous_scan_is_ignored(qtbot, mocked_client, monkeypatch):
+    """
+    Regression for #1299: an async readback of the previous scan that is delivered after the
+    scan changed (already in flight) must not overwrite the data of the current scan.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    curve = wf.plot(arg1="async_device")
+    scan_storage = {
+        scan_id: create_dummy_scan_item(scan_id=scan_id) for scan_id in ("scan-1", "scan-2")
+    }
+    monkeypatch.setattr(
+        wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_storage.get(scan_id)
+    )
+
+    wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {})
+    [(_, previous_subscription)] = _async_readback_subscriptions(wf)
+    wf.on_scan_status({"scan_id": "scan-2", "status": "open"}, {})
+    _deliver_async_readback(
+        wf, _async_readback_endpoint("scan-2", "async_device"), {"async_device": [1, 2, 3]}
+    )
+
+    previous_subscription(
+        {"signals": {"async_device": {"value": [99]}}}, {"async_update": {"type": "replace"}}
+    )
+
+    np.testing.assert_array_equal(curve.get_data()[1], [1, 2, 3])
+
+
+@pytest.mark.parametrize(
+    "async_signal", [False, True], ids=["device_async_readback", "device_async_signal"]
+)
+def test_async_readback_updates_only_curves_of_its_device(
+    qtbot, mocked_client, monkeypatch, async_signal
+):
+    """
+    Regression for #1300: two async devices expose a signal with the same entry name. A
+    readback of one device must update that device's curve only.
+    """
+    from bec_lib.device import ReadoutPriority
+    from bec_lib.endpoints import MessageEndpoints
+
+    from bec_widgets.tests.utils import FakeDevice
+
+    for name in ("async_a", "async_b"):
+        device = FakeDevice(name, readout_priority=ReadoutPriority.ASYNC)
+        device._info["signals"] = {
+            "value": {
+                "kind_str": "hinted",
+                "component_name": "value",
+                "obj_name": "value",
+                "signal_class": "AsyncSignal",
+            }
+        }
+        monkeypatch.setitem(mocked_client.device_manager.devices, name, device)
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    if async_signal:
+        # Both signals are published as BEC async signal streams named after the signal
+        monkeypatch.setattr(wf, "_check_async_signal_found", lambda name, signal: (True, signal))
+    curve_a = wf.plot(device_y="async_a", signal_y="value", label="curve_a")
+    curve_b = wf.plot(device_y="async_b", signal_y="value", label="curve_b")
+    scan_item = create_dummy_scan_item(scan_id="scan-1")
+    monkeypatch.setattr(wf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: scan_item)
+
+    wf.on_scan_status({"scan_id": "scan-1", "status": "open"}, {})
+
+    def endpoint(device: str) -> str:
+        if async_signal:
+            return MessageEndpoints.device_async_signal("scan-1", device, "value").endpoint
+        return MessageEndpoints.device_async_readback("scan-1", device).endpoint
+
+    assert _async_readback_endpoints(wf) == [endpoint("async_a"), endpoint("async_b")]
+    _deliver_async_readback(wf, endpoint("async_b"), {"value": [7, 8]})
+    _deliver_async_readback(wf, endpoint("async_a"), {"value": [1, 2, 3]})
+
+    np.testing.assert_array_equal(curve_a.get_data()[1], [1, 2, 3])
+    np.testing.assert_array_equal(curve_b.get_data()[1], [7, 8])
+
+
 def test_request_dap_skips_unchanged_static_parent(qtbot, mocked_client_with_dap, monkeypatch):
     """
     DAP curves whose parent is a static custom curve must not be resubmitted by
@@ -1419,7 +1589,7 @@ def test_on_async_readback_through_dispatcher(qtbot, mocked_client):
     wrapper = next(
         slot
         for slot in wf.bec_dispatcher._registered_slots.values()
-        if slot.cb == wf.on_async_readback and slot.cb_info == {"scan_id": wf.scan_id}
+        if slot.cb == wf.on_async_readback and slot.cb_info["scan_id"] == wf.scan_id
     )
     msg = {"signals": {"async_device": {"value": [123, 456]}}}
     metadata = {"async_update": {"type": "replace"}}