diff --git a/csaxs_bec/devices/panda_box/panda_box.py b/csaxs_bec/devices/panda_box/panda_box.py
index f881ec4..516d3b0 100644
--- a/csaxs_bec/devices/panda_box/panda_box.py
+++ b/csaxs_bec/devices/panda_box/panda_box.py
@@ -18,6 +18,7 @@ class PandaBoxCSAXS(PandaBox):
         self._timeout_on_completed = 10
 
     def on_stage(self):
+        super().on_stage()
         # TODO, adjust as seen fit.
         # Adjust the acquisition group based on scan parameters if needed
         if self.scan_info.msg.scan_type == "fly":
@@ -35,7 +36,10 @@ class PandaBoxCSAXS(PandaBox):
             captured = 0
             start_time = time.monotonic()
             try:
-                expected_points = int(self.scan_info.msg.num_points * self.scan_info.msg.scan_parameters.get("frames_per_trigger",1))
+                expected_points = int(
+                    self.scan_info.msg.num_points
+                    * self.scan_info.msg.scan_parameters.get("frames_per_trigger", 1)
+                )
                 while captured < expected_points:
                     logger.info(
                         f"Run with captured {captured} and expected points : {expected_points}."
@@ -44,7 +48,9 @@ class PandaBoxCSAXS(PandaBox):
                     captured = int(ret[0].split("=")[-1])
                     time.sleep(0.01)
                     if (time.monotonic() - start_time) > self._timeout_on_completed:
-                        raise TimeoutError(f"Pandabox {self.name} did not complete after {self._timeout_on_completed} with points captured {captured}/{expected_points}")
+                        raise TimeoutError(
+                            f"Pandabox {self.name} did not complete after {self._timeout_on_completed} with points captured {captured}/{expected_points}"
+                        )
             finally:
                 self._disarm()
 
diff --git a/csaxs_bec/devices/panda_box/panda_box_omny.py b/csaxs_bec/devices/panda_box/panda_box_omny.py
index e48f4c2..e41f2d8 100644
--- a/csaxs_bec/devices/panda_box/panda_box_omny.py
+++ b/csaxs_bec/devices/panda_box/panda_box_omny.py
@@ -3,9 +3,8 @@
 import time
 
 from bec_lib.logger import bec_logger
-from ophyd_devices import AsyncMultiSignal, StatusBase
+from ophyd_devices import StatusBase
 from ophyd_devices.devices.panda_box.panda_box import PandaBox, PandaState
-from pandablocks.responses import FrameData
 
 logger = bec_logger.logger
 
@@ -31,7 +30,6 @@ class PandaBoxOMNY(PandaBox):
 
     def on_complete(self):
         """On complete is called after the scan is complete. We need to wait for the capture to complete before we can disarm the PandaBox."""
-        status = super().on_complete()
 
         def _check_capture_complete():
             captured = 0
@@ -55,10 +53,9 @@ class PandaBoxOMNY(PandaBox):
             finally:
                 self._disarm()
 
-        _check_capture_complete()
-
-        if status is not None:
-            status.wait(timeout=self._timeout_on_completed)
+        status_captured = self.task_handler.submit_task(_check_capture_complete, run=True)
+        self.cancel_on_stop(status_captured)
+        return status_captured
 
         # NOTE: This utility class allows to submit a blocking function to a thread and return a status object
         # that can be awaited for. This allows for asynchronous waiting for the capture to complete without blocking