Remove iframe when kernel is killed or restarted

deephaven_ipywidgets/deephaven.py

@@ -7,15 +7,17 @@
 """
 Module for displaying Deephaven widgets within interactive python environments.
 """
+from __future__ import annotations
 
 import __main__
 from ipywidgets import DOMWidget
-from traitlets import Unicode, Integer, Bytes
+from traitlets import Unicode, Integer, Bytes, Bool
 from deephaven_server import Server
 from uuid import uuid4
 from ._frontend import module_name, module_version
 import os
 import base64
+import atexit
 
 
 def _str_object_type(obj):
@@ -34,6 +36,15 @@ def _path_for_object(obj):
     raise TypeError(f"Unknown object type: {name}")
 
 
+def _cleanup(widget: DeephavenWidget):
+    """
+    Remove the widget when the kernel is shutdown
+
+    Args:
+        widget (DeephavenWidget): The widget to remove
+    """
+    widget.set_trait('kernel_active', False)
+
 
 class DeephavenWidget(DOMWidget):
     """A wrapper for viewing DeephavenWidgets in IPython
@@ -52,6 +63,7 @@ class DeephavenWidget(DOMWidget):
     width = Integer().tag(sync=True)
     height = Integer().tag(sync=True)
     token = Unicode().tag(sync=True)
+    kernel_active = Bool().tag(sync=True)
 
     def __init__(self, deephaven_object, height=600, width=0):
         """Create a Deephaven widget for displaying in an interactive Python console.
@@ -122,3 +134,6 @@ def __init__(self, deephaven_object, height=600, width=0):
         self.set_trait('width', width)
         self.set_trait('height', height)
         self.set_trait('token', token)
+        self.set_trait('kernel_active', True)
+
+        atexit.register(_cleanup, self)

package.json

@@ -61,7 +61,7 @@
     "@babel/preset-env": "^7.5.0",
     "@deephaven/eslint-config": "^0.41.0",
     "@deephaven/prettier-config": "^0.41.0",
-    "@jupyterlab/builder": "^3.6.4",
+    "@jupyterlab/builder": "^4.0.5",
     "@phosphor/application": "^1.6.0",
     "@phosphor/widgets": "^1.6.0",
     "@types/jest": "^26.0.0",

src/widget.ts

@@ -104,14 +104,33 @@
     }
   };
 
+  //@ts-ignore
+  onDisconnect = (type): void => {
+    log.info(`Kernel ${type}, removing iframe`);
+    this.iframe.remove();
+  }
+
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   render(): any {
     const iframeUrl = this.model.get('iframe_url');
     const width = this.model.get('width');
     const height = this.model.get('height');
+    let removed = false;
 
     window.addEventListener('message', this.handleAuthentication);
 
+    //@ts-ignore
+    const context = this.model.widget_manager._context;
+
+    this.model.on('change:kernel_active', () => this.onDisconnect("exiting"), this)
+
+    context.sessionContext.statusChanged.connect((sender: any, args: string) => {
+      if (args === "restarting" || args === "terminating" && !removed) {
+        this.onDisconnect(args);
+        removed = true;
+      }
+    })
+
     log.info('init_element for widget', iframeUrl, width, height);
 
     this.iframe = document.createElement('iframe');