fix(waveform): keep async curve data and subscriptions when curves change during a scan

bec_widgets/widgets/plots/waveform/waveform.py

@@ -1562,15 +1562,12 @@ class Waveform(PlotBase):
 
     def _curve_clean_up(self, curve: Curve):
         """
-        Clean up the curve by disconnecting the async update signal (even for sync curves).
+        Clean up the curve. Its async stream subscription is not touched here: other curves can
+        share it, so remove_curve releases it when it re-categorises the remaining curves.
 
         Args:
             curve(Curve): The curve to clean up.
         """
-        self.bec_dispatcher.disconnect_slot(
-            self.on_async_readback,
-            MessageEndpoints.device_async_readback(self.scan_id, curve.name()),
-        )
         curve.rpc_register.remove_rpc(curve)
 
         # Stop tracking the removed curve, so history refreshes and DAP requests skip it
@@ -1833,7 +1830,7 @@ class Waveform(PlotBase):
                 return True, entry_data.get("storage_name")
         return False, signal
 
-    def _setup_async_curve(self, curve: Curve):
+    def _setup_async_curve(self, curve: Curve, previous_streams: dict | None = None):
         """
         Setup async curve.
 
@@ -1845,16 +1842,15 @@ class Waveform(PlotBase):
 
         Args:
             curve(Curve): The curve to set up.
+            previous_streams(dict | None): ``_async_streams_setup`` records of the previous
+                categorisation. A stream recorded there with the same subscription context is
+                still subscribed: it is kept as it is, and the curves it already feeds keep
+                their data.
         """
         name = curve.config.signal.device
         signal = curve.config.signal.signal
         async_signal_found, stream = self._check_async_signal_found(name, signal)
 
-        try:
-            curve.clear_data()
-        except KeyError:
-            logger.warning(f"Curve {name} not found in plot item.")
-
         if async_signal_found:
             stream_key = (name, stream)  # AsyncMultiSignal sub-signals share `stream`
             new_endpoint = MessageEndpoints.device_async_signal(self.scan_id, name, stream)
@@ -1868,11 +1864,24 @@ class Waveform(PlotBase):
             legacy_endpoint = True
 
         endpoint_str = getattr(new_endpoint, "endpoint", new_endpoint)
+        cb_info = {"scan_id": self.scan_id, "device": name, "stream": stream_key[1]}
+
+        previous = (previous_streams or {}).get(stream_key)
+        still_subscribed = previous is not None and previous["cb_info"] == cb_info
+        if not (still_subscribed and signal in previous["signals"]):
+            # A new scan or source, or a curve new to this stream, starts empty. A curve the
+            # subscription already feeds keeps its data: the subscription is not set up again,
+            # so nothing would replay it (e.g. when another curve is added or removed mid-scan).
+            try:
+                curve.clear_data()
+            except KeyError:
+                logger.warning(f"Curve {name} not found in plot item.")
 
         shared = self._async_streams_setup.get(stream_key)
         if shared is not None:
             # Another curve already subscribed to this exact stream this scan.
-            shared["signals"].append(signal)
+            if signal not in shared["signals"]:
+                shared["signals"].append(signal)
             logger.info(
                 f"Async signals {shared['signals']} share a single subscription on endpoint "
                 f"'{endpoint_str}'; reusing it instead of subscribing again."
@@ -1892,12 +1901,16 @@ class Waveform(PlotBase):
         # 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,
         }
+        if still_subscribed:
+            logger.info(
+                f"Async stream on endpoint '{endpoint_str}' is still subscribed; keeping it."
+            )
+            return
         self.bec_dispatcher.connect_slot(
             self.on_async_readback, new_endpoint, from_start=True, cb_info=dict(cb_info)
         )
@@ -2523,7 +2536,7 @@ class Waveform(PlotBase):
             if category == SignalCategory.ASYNC or dev_name in readout_priority_async:
                 self._async_curves.append(curve)
                 if hasattr(self.scan_item, "live_data"):
-                    self._setup_async_curve(curve)
+                    self._setup_async_curve(curve, previous_async_streams)
                 found_async = True
             elif category == SignalCategory.SYNC or dev_name in readout_priority_sync:
                 self._sync_curves.append(curve)

tests/unit_tests/test_waveform.py

@@ -792,6 +792,79 @@ def test_async_readback_updates_only_curves_of_its_device(
     np.testing.assert_array_equal(curve_b.get_data()[1], [7, 8])
 
 
+@pytest.mark.parametrize("change", ["remove_dap_curve", "remove_history_curve", "add_curve"])
+def test_curve_changes_mid_scan_keep_live_async_stream_and_data(
+    qtbot, mocked_client_with_dap, monkeypatch, scan_history_factory, change
+):
+    """
+    DAP and history curves share the DeviceSignal of their live counterpart. Adding or removing
+    such a curve (or any other curve) during a live scan re-categorises the curves; the live
+    async stream must stay subscribed without being registered again, and the live curve must
+    keep the data it already received.
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
+    live = wf.plot(arg1="bpm4i")
+    inject_scan_history(wf, scan_history_factory, ("hist-1", 1))
+    scan_item = create_dummy_scan_item(scan_id="scan-1")
+    # bpm4i is read out asynchronously in this scan
+    scan_item.status_message.info["readout_priority"] = {"monitored": [], "async": ["bpm4i"]}
+    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"}, {})
+    endpoint = _async_readback_endpoint("scan-1", "bpm4i")
+    [(_, subscription)] = _async_readback_subscriptions(wf)
+    _deliver_async_readback(wf, endpoint, {"bpm4i": [1, 2, 3]}, update_type="add")
+
+    connector = wf.client.connector
+    register_spy = mock.patch.object(connector, "register", wraps=connector.register)
+    unregister_spy = mock.patch.object(connector, "unregister", wraps=connector.unregister)
+    with register_spy as register, unregister_spy as unregister:
+        if change == "remove_dap_curve":
+            # The fit request is not under test (it cannot serialise the mocked scan item)
+            monkeypatch.setattr(wf, "request_dap", lambda *args, **kwargs: None)
+            dap_curve = wf.add_dap_curve(device_label=live.name(), dap_name="GaussianModel")
+            wf.remove_curve(dap_curve.name())
+        elif change == "remove_history_curve":
+            history_curve = wf.plot(device_y="bpm4i", signal_y="bpm4i", scan_id="hist-1")
+            wf.remove_curve(history_curve.name())
+        else:
+            wf.plot(arg1="samx")
+        live_stream_calls = [
+            call
+            for call in register.call_args_list + unregister.call_args_list
+            if getattr(call.args[0], "endpoint", call.args[0]) == endpoint
+        ]
+
+    assert live_stream_calls == []
+    assert _async_readback_subscriptions(wf) == [(endpoint, subscription)]
+    np.testing.assert_array_equal(live.get_data()[1], [1, 2, 3])
+    _deliver_async_readback(wf, endpoint, {"bpm4i": [4]}, update_type="add")
+    np.testing.assert_array_equal(live.get_data()[1], [1, 2, 3, 4])
+
+
+def test_removing_curve_named_like_its_device_keeps_shared_async_stream(
+    qtbot, mocked_client, monkeypatch
+):
+    """
+    Removing one of two curves fed by the same async stream must keep the stream subscribed for
+    the other one, also when the removed curve's label equals the device name (which is what
+    the old per-curve disconnect derived the endpoint from).
+    """
+    wf = create_widget(qtbot, Waveform, client=mocked_client)
+    wf.plot(arg1="async_device", label="async_device")
+    remaining = wf.plot(arg1="async_device", label="async_device_copy")
+    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"}, {})
+    endpoint = _async_readback_endpoint("scan-1", "async_device")
+    _deliver_async_readback(wf, endpoint, {"async_device": [1, 2]}, update_type="add")
+
+    wf.remove_curve("async_device")
+
+    assert _async_readback_endpoints(wf) == [endpoint]
+    _deliver_async_readback(wf, endpoint, {"async_device": [3]}, update_type="add")
+    np.testing.assert_array_equal(remaining.get_data()[1], [1, 2, 3])
+
+
 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