fix(jupyter): Rework comm creation to make it compatible w/ other ext…

…ensions

package.json

@@ -56,6 +56,7 @@
     "dependencies": {
         "@jupyterlab/application": "^3.5.0 || ^4.0.5",
         "@jupyterlab/coreutils": "^6.0.0",
+        "@jupyterlab/notebook": "^4.0.7",
         "@jupyterlab/services": "^7.0.0"
     },
     "devDependencies": {

src/comm.ts

@@ -2,10 +2,7 @@ import {
   IKernelConnection,
   IComm
 } from '@jupyterlab/services/lib/kernel/kernel';
-import {
-  ICommMsgMsg,
-  ICommCloseMsg
-} from '@jupyterlab/services/lib/kernel/messages';
+import { ICommMsgMsg } from '@jupyterlab/services/lib/kernel/messages';
 
 import { ConcreteEmitter } from './emitter';
 
@@ -16,6 +13,8 @@ export type CommMessage = {
 
 export type CommEvents = {
   message: CommMessage;
+  open: void;
+  close: void;
 };
 
 export class TrameJupyterComm extends ConcreteEmitter<CommEvents> {
@@ -27,10 +26,28 @@ export class TrameJupyterComm extends ConcreteEmitter<CommEvents> {
 
     this.kernel = kernel;
     this.comm = null;
+
+    this.kernel.disposed.connect(this.onClose.bind(this));
+    this.kernel.statusChanged.connect((_kernel, status) => {
+      if (
+        status === 'restarting' ||
+        status === 'autorestarting' ||
+        status === 'terminating' ||
+        status === 'dead'
+      ) {
+        this.onClose();
+      }
+    });
   }
 
   open() {
-    if (!this.comm || this.comm.isDisposed) {
+    if (this.kernel.isDisposed) {
+      throw new Error(
+        `Can't open a comm for disposed kernel ${this.kernel.id}`
+      );
+    }
+
+    if ((!this.comm || this.comm.isDisposed) && !this.kernel.isDisposed) {
       this.comm = this.kernel.createComm('wslink_comm');
       this.comm.open();
       this.comm.onMsg = this.onMessage.bind(this);
@@ -39,8 +56,8 @@ export class TrameJupyterComm extends ConcreteEmitter<CommEvents> {
   }
 
   send(message: CommMessage) {
-    if (this.comm) {
-      this.comm.send(message.data, undefined, message.buffers);
+    if (this.isUseable()) {
+      this.comm!.send(message.data, undefined, message.buffers);
     } else {
       console.error('trame::jupyter-comm::send -- NO COMM');
     }
@@ -50,7 +67,13 @@ export class TrameJupyterComm extends ConcreteEmitter<CommEvents> {
     this.emit('message', { data: msg.content.data, buffers: msg.buffers });
   }
 
-  onClose(msg: ICommCloseMsg<'iopub' | 'shell'>) {
+  onClose(...args: any) {
     console.error('trame::jupyter-comm::close -- NO COMM');
+    this.comm = null;
+    this.emit('close', undefined);
+  }
+
+  isUseable() {
+    return !this.kernel.isDisposed && this.comm && !this.comm.isDisposed;
   }
 }

src/index.ts

@@ -1,11 +1,43 @@
+import { IDisposable } from '@lumino/disposable';
+
 import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
 
-import { ActiveManager } from './active';
+import { DocumentRegistry } from '@jupyterlab/docregistry';
+import { NotebookPanel, INotebookModel } from '@jupyterlab/notebook';
+import { Kernel } from '@jupyterlab/services';
+
 import { TrameJupyterComm } from './comm';
 import { TrameJupyterWebSocket } from './websocket';
+import { ContextManager } from './manager';
+import { Registry } from './registry';
+import { getExtensionLocation } from './location';
+
+/**
+ * A notebook widget extension that creates a kernel manager each time a notebook is opened.
+ */
+export class WidgetExtension
+  implements DocumentRegistry.IWidgetExtension<NotebookPanel, INotebookModel>
+{
+  private _endpoint: string;
+  private _www: string;
+
+  constructor(endpoint: string, www: string) {
+    this._endpoint = endpoint;
+    this._www = www;
+  }
+
+  createNew(
+    _panel: NotebookPanel,
+    context: DocumentRegistry.IContext<INotebookModel>
+  ): IDisposable {
+    const manager = new ContextManager(context, this._endpoint, this._www);
+
+    return manager;
+  }
+}
 
 /**
  * Initialization data for the trame-jupyter-extension extension.
@@ -14,41 +46,52 @@ const plugin: JupyterFrontEndPlugin<void> = {
   id: 'trame-jupyter-extension:plugin',
   description: 'A JupyterLab extension for trame communication layer',
   autoStart: true,
-  activate: (app: JupyterFrontEnd) => {
-    const activeManager = new ActiveManager(app);
-    const comms: Record<string, TrameJupyterComm> = {};
+  activate: async (app: JupyterFrontEnd) => {
+    const kernelsRegistry = new Registry<Kernel.IKernelConnection>();
+    const commsRegistry = new Registry<TrameJupyterComm>();
 
     function init(childWindow: any) {
       const kernelId = childWindow.frameElement.dataset.kernelId;
-      if (!comms[kernelId]) {
-        const kc = activeManager.getKernelConnection(kernelId);
-        if (!kc) {
-          throw new Error(
-            `trame: Could not get kernel connection to ${kernelId}`
-          );
-        }
-        comms[kernelId] = new TrameJupyterComm(kc);
+      const kc = kernelsRegistry.getItem(kernelId);
+      let comm = commsRegistry.getItem(kernelId);
 
-        // Open kernel connection at creation
-        comms[kernelId].open();
+      if (!kc) {
+        throw new Error(
+          `trame: Could not get kernel connection to ${kernelId}`
+        );
       }
-      if (comms[kernelId]) {
-        return {
-          createWebSocket: () => {
-            return new TrameJupyterWebSocket(childWindow, comms[kernelId]);
-          }
-        };
+
+      if (!comm || !comm.isUseable()) {
+        comm = new TrameJupyterComm(kc);
+        comm.open();
+        commsRegistry.setItem(kernelId, comm);
+        comm.addEventListener('close', () => {
+          commsRegistry.setItem(kernelId, null);
+        });
       }
+
+      return {
+        createWebSocket: () => {
+          return new TrameJupyterWebSocket(childWindow, comm!);
+        }
+      };
     }
 
     const namespace = {
       app,
-      activeManager,
-      comms,
+      kernelsRegistry,
+      commsRegistry,
       init
     };
 
     (window as any).trameJupyter = namespace;
+
+    const { endpoint, www } = await getExtensionLocation();
+
+    app.docRegistry.addWidgetExtension(
+      'Notebook',
+      new WidgetExtension(endpoint, www)
+    );
   }
 };
 

src/location.ts

@@ -0,0 +1,42 @@
+import { URLExt } from '@jupyterlab/coreutils';
+import { ServerConnection } from '@jupyterlab/services';
+
+async function getExtensionLocation(): Promise<{
+  endpoint: string;
+  www: string;
+}> {
+  let www = '';
+  let endpoint = '';
+
+  const settings = ServerConnection.makeSettings();
+  endpoint = URLExt.join(settings.baseUrl, 'trame-jupyter-server');
+  const requestUrl = URLExt.join(
+    settings.baseUrl,
+    'trame-jupyter-server',
+    'location'
+  );
+  let response: Response;
+  try {
+    response = await ServerConnection.makeRequest(requestUrl, {}, settings);
+  } catch (error) {
+    throw new ServerConnection.NetworkError(error as any);
+  }
+
+  const data: any = await response.text();
+
+  if (data.length > 0) {
+    try {
+      www = JSON.parse(data).www;
+    } catch (error) {
+      console.log('Not a JSON response body.', response);
+    }
+  }
+
+  if (!response.ok) {
+    throw new ServerConnection.ResponseError(response, data.message || data);
+  }
+
+  return { endpoint, www };
+}
+
+export { getExtensionLocation };