From e4e999e3462b4653b9e50c27c88d33a6083eb1d3 Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Tue, 15 Apr 2025 21:43:40 +0200
Subject: [PATCH] wip

---
 bec_widgets/utils/bec_dispatcher.py              | 9 ++++-----
 bec_widgets/widgets/plots/waveform/waveform.py   | 9 +++++++++
 bec_widgets/widgets/utility/logpanel/logpanel.py | 3 ---
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/bec_widgets/utils/bec_dispatcher.py b/bec_widgets/utils/bec_dispatcher.py
index 976b70b0..3b02f21e 100644
--- a/bec_widgets/utils/bec_dispatcher.py
+++ b/bec_widgets/utils/bec_dispatcher.py
@@ -25,9 +25,9 @@ if TYPE_CHECKING:  # pragma: no cover
 class QtThreadSafeCallback(QObject):
     cb_signal = pyqtSignal(dict, dict)
 
-    def __init__(self, cb):
+    def __init__(self, cb: Callable, cb_info: dict | None = None):
         super().__init__()
-        self.topics = None
+        self.cb_info = cb_info
         self.cb = cb
         self.cb_signal.connect(self.cb)
 
@@ -38,7 +38,6 @@ class QtThreadSafeCallback(QObject):
         return id(self.cb)
 
     def __call__(self, msg_content, metadata):
-        logger.info(f"Received message for topic {self.topics}")
         self.cb_signal.emit(msg_content, metadata)
 
 
@@ -138,6 +137,7 @@ class BECDispatcher:
         self,
         slot: Callable,
         topics: Union[EndpointInfo, str, list[Union[EndpointInfo, str]]],
+        cb_info: dict | None = None,
         **kwargs,
     ) -> None:
         """Connect widget's qt slot, so that it is called on new pub/sub topic message.
@@ -147,8 +147,7 @@ class BECDispatcher:
                 the corresponding pub/sub message
             topics (EndpointInfo | str | list): A topic or list of topics that can typically be acquired via bec_lib.MessageEndpoints
         """
-        slot = QtThreadSafeCallback(slot)
-        slot.topics = topics
+        slot = QtThreadSafeCallback(slot, cb_info=cb_info)
         self.client.connector.register(topics, cb=slot, **kwargs)
         topics_str, _ = self.client.connector._convert_endpointinfo(topics)
         self._slots[slot].update(set(topics_str))
diff --git a/bec_widgets/widgets/plots/waveform/waveform.py b/bec_widgets/widgets/plots/waveform/waveform.py
index bfa4cad3..66205b01 100644
--- a/bec_widgets/widgets/plots/waveform/waveform.py
+++ b/bec_widgets/widgets/plots/waveform/waveform.py
@@ -1014,6 +1014,8 @@ class Waveform(PlotBase):
             self.old_scan_id = self.scan_id
             self.scan_id = current_scan_id
             self.scan_item = self.queue.scan_storage.find_scan_by_ID(self.scan_id)  # live scan
+            if self.scan_item is None:
+                raise ValueError(f"Scan item with ID {self.scan_id} not found in the queue.")
             self._slice_index = None  # Reset the slice index
 
             self._mode = self._categorise_device_curves()
@@ -1180,6 +1182,7 @@ class Waveform(PlotBase):
             self.on_async_readback,
             MessageEndpoints.device_async_readback(self.scan_id, name),
             from_start=True,
+            cb_info={"scan_id": self.scan_id},
         )
         logger.info(
             f"remaining subscriptions: {self.bec_dispatcher.client.connector._stream_topics_subscription.values()}"
@@ -1205,6 +1208,12 @@ class Waveform(PlotBase):
             msg(dict): Message with the async data.
             metadata(dict): Metadata of the message.
         """
+        sender = self.sender()
+        if sender and hasattr(sender, "cb_info"):
+            scan_id = sender.cb_info.get("scan_id")
+            if scan_id != self.scan_id:
+                logger.warning("Scan ID mismatch, ignoring async readback.")
+                return
         instruction = metadata.get("async_update", {}).get("type")
         if instruction not in ["add", "add_slice", "replace"]:
             logger.warning(f"Invalid async update instruction: {instruction}")
diff --git a/bec_widgets/widgets/utility/logpanel/logpanel.py b/bec_widgets/widgets/utility/logpanel/logpanel.py
index 97ba3d06..de50c8ee 100644
--- a/bec_widgets/widgets/utility/logpanel/logpanel.py
+++ b/bec_widgets/widgets/utility/logpanel/logpanel.py
@@ -52,9 +52,6 @@ from bec_widgets.widgets.utility.logpanel._util import (
     simple_color_format,
 )
 
-if TYPE_CHECKING:  # pragma: no cover
-    from PySide6.QtCore import SignalInstance
-
 logger = bec_logger.logger
 
 MODULE_PATH = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))