From 16d88fe8df538bb54e2e961b71098f6d87734df6 Mon Sep 17 00:00:00 2001
From: wyzula-jan <jan.wyzula@psi.ch>
Date: Fri, 17 Oct 2025 13:40:42 +0200
Subject: [PATCH] fix(jupyter_console): cleanup

---
 .../jupyter_console/jupyter_console.py        | 49 ++++++++++++++-----
 1 file changed, 38 insertions(+), 11 deletions(-)

diff --git a/bec_widgets/widgets/editors/jupyter_console/jupyter_console.py b/bec_widgets/widgets/editors/jupyter_console/jupyter_console.py
index 24234db6..e5322689 100644
--- a/bec_widgets/widgets/editors/jupyter_console/jupyter_console.py
+++ b/bec_widgets/widgets/editors/jupyter_console/jupyter_console.py
@@ -2,6 +2,7 @@ from bec_ipython_client.main import BECIPythonClient
 from qtconsole.inprocess import QtInProcessKernelManager
 from qtconsole.manager import QtKernelManager
 from qtconsole.rich_jupyter_widget import RichJupyterWidget
+from qtpy.QtCore import Qt
 from qtpy.QtWidgets import QApplication, QMainWindow
 
 
@@ -9,10 +10,10 @@ class BECJupyterConsole(RichJupyterWidget):  # pragma: no cover:
     def __init__(self, inprocess: bool = False):
         super().__init__()
 
-        self.inprocess = None
-        self.client = None
+        self.inprocess = inprocess
+        self.ipyclient = None
 
-        self.kernel_manager, self.kernel_client = self._init_kernel(inprocess=inprocess)
+        self.kernel_manager, self.kernel_client = self._init_kernel(inprocess=self.inprocess)
         self.set_default_style("linux")
         self._init_bec()
 
@@ -35,14 +36,13 @@ class BECJupyterConsole(RichJupyterWidget):  # pragma: no cover:
             self._init_bec_kernel()
 
     def _init_bec_inprocess(self):
-        self.client = BECIPythonClient()
-        self.client.start()
-
+        self.ipyclient = BECIPythonClient()
+        self.ipyclient.start()
         self.kernel_manager.kernel.shell.push(
             {
-                "bec": self.client,
-                "dev": self.client.device_manager.devices,
-                "scans": self.client.scans,
+                "bec": self.ipyclient,
+                "dev": self.ipyclient.device_manager.devices,
+                "scans": self.ipyclient.scans,
             }
         )
 
@@ -57,20 +57,47 @@ class BECJupyterConsole(RichJupyterWidget):  # pragma: no cover:
             """
         )
 
+    def _cleanup_bec(self):
+        if getattr(self, "ipyclient", None) is not None and self.inprocess is True:
+            self.ipyclient.shutdown()
+            self.ipyclient = None
+
     def shutdown_kernel(self):
+        """
+        Shutdown the Jupyter kernel and clean up resources.
+        """
+        self._cleanup_bec()
         self.kernel_client.stop_channels()
         self.kernel_manager.shutdown_kernel()
+        self.kernel_client = None
+        self.kernel_manager = None
 
     def closeEvent(self, event):
         self.shutdown_kernel()
+        event.accept()
+        super().closeEvent(event)
+
+
+class JupyterConsoleWindow(QMainWindow):  # pragma: no cover:
+    def __init__(self, inprocess: bool = True, parent=None):
+        super().__init__(parent)
+        self.console = BECJupyterConsole(inprocess=inprocess)
+        self.setCentralWidget(self.console)
+        self.setAttribute(Qt.WA_DeleteOnClose, True)
+
+    def closeEvent(self, event):
+        # Explicitly close the console so its own closeEvent runs
+        if getattr(self, "console", None) is not None:
+            self.console.close()
+        event.accept()
+        super().closeEvent(event)
 
 
 if __name__ == "__main__":  # pragma: no cover
     import sys
 
     app = QApplication(sys.argv)
-    win = QMainWindow()
-    win.setCentralWidget(BECJupyterConsole(True))
+    win = JupyterConsoleWindow(inprocess=True)
     win.show()
 
     sys.exit(app.exec_())