From baf7dae1426585a7f5a985d9364917e696d60760 Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Wed, 7 Oct 2026 15:40:17 +0200 Subject: [PATCH] fix(waveform): keep async curve data and subscriptions when curves change during a scan --- .../widgets/plots/waveform/waveform.py | 41 +++++++---- tests/unit_tests/test_waveform.py | 73 +++++++++++++++++++ 2 files changed, 100 insertions(+), 14 deletions(-) diff --git a/bec_widgets/widgets/plots/waveform/waveform.py b/bec_widgets/widgets/plots/waveform/waveform.py index 1e0745bf..d11509dd 100644 --- a/bec_widgets/widgets/plots/waveform/waveform.py +++ b/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) diff --git a/tests/unit_tests/test_waveform.py b/tests/unit_tests/test_waveform.py index aac6e522..dd3b15f6 100644 --- a/tests/unit_tests/test_waveform.py +++ b/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