UI for blocked kernel actions (jupyter-server#222)

* add a error modal in the UI when a kernel action is blocked

* make the console log less aggressive

* add initial test for kernelblock plugin

* test ui element for exact match

data_studio_jupyter_extensions/configurables/provisioner.py

@@ -366,7 +366,9 @@ async def terminate(self, restart: bool = False) -> None:
         """
         pass
 
-    async def _require_status(self, required_status: Union[str, Iterable]) -> str:
+    async def _require_status(
+        self, action, required_status: Union[str, Iterable]
+    ) -> str:
         """Useful for ensuring the kernel is in a desired state before doing another action.
 
         Returns the status.
@@ -381,9 +383,12 @@ async def _require_status(self, required_status: Union[str, Iterable]) -> str:
         # Check that the status is in on of the required states.
         if status_label not in required_status:
             msg = (
-                f"The current kernel status is '{status_label}', but must be "
-                f"in one of the following statuses to continue: {required_status}"
+                f"Cannot {action} the kernel at this time. It is currently in a "
+                f"state, '{status_label}', which blocks this action. Please retry "
+                f"when the kernel is in one of the following states: {required_status}"
             )
+            # Emit this event to the telemetry bus.
+            self._emit_kernel_blocked(action, status_label)
             raise HTTPError(status_code=409, log_message=msg)
         return status_label
 
@@ -401,14 +406,14 @@ async def shutdown_requested(self, restart: bool = False) -> None:
                 "failedwithsecrets",
                 "failed",
             }
-            await self._require_status(statuses)
+            await self._require_status("restart", statuses)
             # Emit a log message to let users known the kernel is restarting.
             self._emit_kernel_status(
                 status="Restarting",
                 description="Terminating the current kernel, then starting a new kernel.",
             )
         else:
-            await self._require_status(["ready", "running"])
+            await self._require_status("shutdown", ["ready", "running"])
             # Emit a log message to let users known the kernel is restarting.
             self._emit_kernel_status(
                 status="Shutdown", description="Terminating the current kernel."
@@ -557,3 +562,16 @@ def _emit_kernel_status(self, status, description=""):
                 "kernel_id": self.kernel_id,
             },
         )
+
+    def _emit_kernel_blocked(self, action, status, message=""):
+        """Emit a kernel status event."""
+        self.telemetry_bus.record_event(
+            schema_name="event.datastudio.jupyter.com/kernel-blocked",
+            version=1,
+            event={
+                "process_id": self.process_id or "Not set yet",
+                "action": action,
+                "status": status,
+                "message": message,
+            },
+        )

data_studio_jupyter_extensions/extensions/telemetry/extension.py

@@ -40,5 +40,6 @@ def initialize_configurables(self):
                 "event.datastudio.jupyter.com/kernel-message",
                 "event.datastudio.jupyter.com/kernel-status",
                 "event.datastudio.jupyter.com/session-message",
+                "event.datastudio.jupyter.com/kernel-blocked",
             ],
         )

data_studio_jupyter_extensions/extensions/telemetry/schemas/kernels/kernel-blocked-message.v1.yaml

@@ -0,0 +1,23 @@
+$id: event.datastudio.jupyter.com/kernel-blocked
+version: 1
+title: Kernel Blocked Message
+description: |
+  Emit a message that the last kernel action was blocked.
+type: object
+properties:
+  process_id:
+    title: Kernel Process ID
+    description: |
+      UUID for this kernel process.
+  action:
+    title: Attempted action
+    description: |
+      The action that was attempted and blocked by the server.
+  message:
+    title: Message type
+    description: |
+      Message returned by the Kernel Provisioner.
+required:
+  - process_id
+  - action
+  - message

data_studio_jupyter_extensions/tests/configurables/test_provisioner.py

@@ -244,10 +244,12 @@ async def test_require_status(notebook_service_client):
         notebook_id="foo",
     )
     # Check a good request
-    r = await provisioner._require_status(required_status=["running"])
+    r = await provisioner._require_status("restart", required_status=["running"])
     assert r == "running"
 
     # Check a failing request
     with pytest.raises(HTTPError) as err:
-        r = await provisioner._require_status(required_status=["terminated"])
+        r = await provisioner._require_status(
+            "shutdown", required_status=["terminated"]
+        )
     assert err.value.status_code == 409

src/index.ts

@@ -6,14 +6,16 @@ import { FileRadarPlugin } from './fileradar';
 import { SessionExpPlugin } from './sessionexp';
 import { SyntaxHighlightPlugin } from './syntax-highlight';
 import { HTMLReplacementPlugin } from './overridehtml';
+import { KernelBlockedPlugin } from './kernelblocked';
 
 const plugins: JupyterFrontEndPlugin<any>[] = [
   kernelStatusPlugin,
   ExternalLinksPlugin,
   FileRadarPlugin,
   SessionExpPlugin,
   SyntaxHighlightPlugin,
-  HTMLReplacementPlugin
+  HTMLReplacementPlugin,
+  KernelBlockedPlugin
 ];
 
 export default plugins;

src/kernelblocked.ts

@@ -0,0 +1,43 @@
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+
+import { PageConfig, URLExt } from '@jupyterlab/coreutils';
+import { showErrorMessage } from '@jupyterlab/apputils';
+
+/**
+ * A plugin to capture Kernel action blocking events and surface to the user.
+ */
+export const KernelBlockedPlugin: JupyterFrontEndPlugin<void> = {
+  id: 'data_studio:kernel_blocked_plugin',
+  autoStart: true,
+  requires: [],
+  activate: async (app: JupyterFrontEnd) => {
+    console.log('JupyterLab extension "Kernel Blocked Dialog" is activated!');
+
+    let url =
+      PageConfig.getOption('studioSubscribeURL') ||
+      URLExt.join(PageConfig.getWsUrl(), 'subscribe');
+
+    const ws = new WebSocket(url);
+
+    ws.addEventListener('message', function (event) {
+      const data = JSON.parse(event.data);
+      console.log('Kernel action blocked:', data);
+
+      // Use session messages to update the path map and also log them.
+      if (data.__schema__ == 'event.datastudio.jupyter.com/kernel-blocked') {
+        const msg =
+          'Cannot ' +
+          data.action +
+          ' the kernel at this time. ' +
+          "It is currently in a state, '" +
+          data.status +
+          "', which blocks this action.";
+
+        showErrorMessage('409: Kernel Action Blocked', msg);
+      }
+    });
+  }
+};

test/kernelblocked.spec.ts

@@ -0,0 +1,52 @@
+import { JupyterLab } from '@jupyterlab/application';
+
+import { waitForDialog } from '@jupyterlab/testutils';
+
+import WS from 'jest-websocket-mock';
+
+import { KernelBlockedPlugin } from '../src/kernelblocked';
+
+import { PageConfig } from '@jupyterlab/coreutils';
+
+describe('kernel blocked', () => {
+  let app: JupyterLab;
+
+  beforeEach(async () => {
+    // Wait for the server to start before creating the application
+    // so we pick up the page config (base url, etc.)
+    app = new JupyterLab();
+  });
+
+  it('should capture blocked kernel messages from telemetry bus', async () => {
+    // Prepare a move websocket.
+    const url = 'ws://localhost:5555';
+    PageConfig.setOption('studioSubscribeURL', url);
+    const server = new WS(url, { jsonProtocol: true });
+
+    // Activate the extension.
+    await KernelBlockedPlugin.activate(app);
+
+    await server.connected;
+
+    // Mock a message from the server
+    const message = {
+      __timestamp__: '2021-08-26T11:36:40.946609Z',
+      __schema__: 'event.datastudio.jupyter.com/kernel-blocked',
+      __schema_version__: 1,
+      __metadata_version__: 1,
+      action: 'restart',
+      process_id: 'test-id'
+    };
+    server.send(message);
+
+    // Wait for the error dialog.
+    await waitForDialog();
+
+    const dialog = document.body.getElementsByClassName('jp-Dialog')[0];
+    const header = dialog.getElementsByClassName('jp-Dialog-header')[0];
+
+    expect(header.innerHTML).toBe('409: Kernel Action Blocked');
+    server.close();
+    WS.clean();
+  });
+});