Clear kernel usage state when switching away

Author: krassowski

packages/labextension/package.json

@@ -44,6 +44,7 @@
   "dependencies": {
     "@jupyterlab/application": "^3.5.1",
     "@jupyterlab/apputils": "^3.5.1",
+    "@jupyterlab/console": "^3.5.1",
     "@jupyterlab/coreutils": "^5.5.1",
     "@jupyterlab/notebook": "^3.5.1",
     "@jupyterlab/services": "^6.5.1",

packages/labextension/src/index.ts

@@ -1,10 +1,12 @@
 import {
+  ILabShell,
   JupyterFrontEnd,
   JupyterFrontEndPlugin,
 } from '@jupyterlab/application';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { LabIcon } from '@jupyterlab/ui-components';
 import { ICommandPalette } from '@jupyterlab/apputils';
+import { IConsoleTracker } from '@jupyterlab/console';
 import { KernelUsagePanel } from './panel';
 import tachometer from '../style/tachometer.svg';
 
@@ -13,6 +15,7 @@ import { IStatusBar } from '@jupyterlab/statusbar';
 import { ITranslator } from '@jupyterlab/translation';
 
 import { MemoryUsage } from './memoryUsage';
+import { KernelWidgetTracker } from './tracker';
 
 namespace CommandIDs {
   export const getKernelUsage = 'kernel-usage:get';
@@ -46,13 +49,15 @@ const memoryStatusPlugin: JupyterFrontEndPlugin<void> = {
 const kernelUsagePlugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyter-server/resource-usage:kernel-panel-item',
   autoStart: true,
-  optional: [ICommandPalette],
+  optional: [ICommandPalette, ILabShell, IConsoleTracker],
   requires: [ITranslator, INotebookTracker],
   activate: (
     app: JupyterFrontEnd,
     translator: ITranslator,
     notebookTracker: INotebookTracker,
-    palette: ICommandPalette | null
+    palette: ICommandPalette | null,
+    labShell: ILabShell | null,
+    consoleTracker: IConsoleTracker | null
   ) => {
     const trans = translator.load('jupyter-resource-usage');
 
@@ -62,10 +67,15 @@ const kernelUsagePlugin: JupyterFrontEndPlugin<void> = {
     let panel: KernelUsagePanel | null = null;
 
     function createPanel() {
-      if ((!panel || panel.isDisposed) && notebookTracker) {
+      if (!panel || panel.isDisposed) {
+        const tracker = new KernelWidgetTracker({
+          notebookTracker,
+          labShell,
+          consoleTracker,
+        });
+
         panel = new KernelUsagePanel({
-          widgetAdded: notebookTracker.widgetAdded,
-          currentNotebookChanged: notebookTracker.currentChanged,
+          currentChanged: tracker.currentChanged,
           trans: trans,
         });
         shell.add(panel, 'right', { rank: 200 });

packages/labextension/src/panel.ts

@@ -3,16 +3,17 @@ import { ISignal } from '@lumino/signaling';
 import { TranslationBundle } from '@jupyterlab/translation';
 import { StackedPanel } from '@lumino/widgets';
 import { LabIcon } from '@jupyterlab/ui-components';
-import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook';
 import { KernelUsageWidget } from './widget';
+import { IWidgetWithSession } from './types';
+import { KernelWidgetTracker } from './tracker';
+
 import tachometer from '../style/tachometer.svg';
 
 const PANEL_CLASS = 'jp-KernelUsage-view';
 
 export class KernelUsagePanel extends StackedPanel {
   constructor(props: {
-    widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
-    currentNotebookChanged: ISignal<INotebookTracker, NotebookPanel | null>;
+    currentChanged: ISignal<KernelWidgetTracker, IWidgetWithSession | null>;
     trans: TranslationBundle;
   }) {
     super();
@@ -24,9 +25,9 @@ export class KernelUsagePanel extends StackedPanel {
       svgstr: tachometer,
     });
     this.title.closable = true;
+
     const widget = new KernelUsageWidget({
-      widgetAdded: props.widgetAdded,
-      currentNotebookChanged: props.currentNotebookChanged,
+      currentChanged: props.currentChanged,
       panel: this,
       trans: props.trans,
     });

packages/labextension/src/tracker.ts

@@ -0,0 +1,62 @@
+import { ILabShell } from '@jupyterlab/application';
+import { INotebookTracker } from '@jupyterlab/notebook';
+import { IConsoleTracker } from '@jupyterlab/console';
+import { ISignal, Signal } from '@lumino/signaling';
+import { IWidgetWithSession, hasKernelSession } from './types';
+
+/**
+ * Tracks widgets with kernels as well as possible given the available tokens.
+ */
+export class KernelWidgetTracker {
+  constructor(options: KernelWidgetTracker.IOptions) {
+    const { labShell, notebookTracker, consoleTracker } = options;
+    this._currentChanged = new Signal(this);
+    if (labShell) {
+      labShell.currentChanged.connect((_, update) => {
+        const widget = update.newValue;
+        if (widget && hasKernelSession(widget)) {
+          this._currentChanged.emit(widget);
+        } else {
+          this._currentChanged.emit(null);
+        }
+      });
+    } else {
+      notebookTracker.currentChanged.connect((_, widget) => {
+        this._currentChanged.emit(widget);
+      });
+      if (consoleTracker) {
+        consoleTracker.currentChanged.connect((_, widget) => {
+          this._currentChanged.emit(widget);
+        });
+      }
+    }
+  }
+
+  /**
+   * Emits on any change of active widget. The value is a known widget with
+   * kernel session or null if user switched to a widget which does not support
+   * kernel sessions.
+   */
+  get currentChanged(): ISignal<
+    KernelWidgetTracker,
+    IWidgetWithSession | null
+  > {
+    return this._currentChanged;
+  }
+
+  private _currentChanged: Signal<
+    KernelWidgetTracker,
+    IWidgetWithSession | null
+  >;
+}
+
+/**
+ * Namespace for kernel widget tracker.
+ */
+export namespace KernelWidgetTracker {
+  export interface IOptions {
+    notebookTracker: INotebookTracker;
+    labShell: ILabShell | null;
+    consoleTracker: IConsoleTracker | null;
+  }
+}

packages/labextension/src/types.ts

@@ -0,0 +1,26 @@
+import { Widget } from '@lumino/widgets';
+import { ISessionContext } from '@jupyterlab/apputils';
+import { NotebookPanel } from '@jupyterlab/notebook';
+import { ConsolePanel } from '@jupyterlab/console';
+
+/**
+ * Interface used to abstract away widgets with kernel to avoid introducing
+ * spurious dependencies on properties specific to e.g. notebook or console.
+ */
+export interface IWidgetWithSession extends Widget {
+  /**
+   * Session context providing kernel access.
+   */
+  sessionContext: ISessionContext;
+}
+
+/**
+ * Check if given widget is one of the widgets known to have kernel session.
+ *
+ * Note: we could switch to duck-typing in future.
+ */
+export function hasKernelSession(
+  widget: Widget
+): widget is ConsolePanel | NotebookPanel {
+  return widget instanceof ConsolePanel || widget instanceof NotebookPanel;
+}

packages/labextension/src/widget.tsx

@@ -4,11 +4,12 @@ import { ReactWidget, ISessionContext } from '@jupyterlab/apputils';
 import { IChangedArgs } from '@jupyterlab/coreutils';
 import { Kernel } from '@jupyterlab/services';
 import { TranslationBundle } from '@jupyterlab/translation';
-import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook';
 import { requestAPI } from './handler';
 import { KernelUsagePanel } from './panel';
 import useInterval from './useInterval';
 import { formatForDisplay } from './format';
+import { IWidgetWithSession } from './types';
+import { KernelWidgetTracker } from './tracker';
 
 type Usage = {
   timestamp: Date | null;
@@ -43,12 +44,11 @@ type KernelChangeCallback = (
 ) => void;
 let kernelChangeCallback: {
   callback: KernelChangeCallback;
-  panel: NotebookPanel;
+  panel: IWidgetWithSession;
 } | null = null;
 
 const KernelUsage = (props: {
-  widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
-  currentNotebookChanged: ISignal<INotebookTracker, NotebookPanel | null>;
+  currentChanged: ISignal<KernelWidgetTracker, IWidgetWithSession | null>;
   panel: KernelUsagePanel;
   trans: TranslationBundle;
 }) => {
@@ -79,7 +79,7 @@ const KernelUsage = (props: {
   };
 
   useEffect(() => {
-    const createKernelChangeCallback = (panel: NotebookPanel) => {
+    const createKernelChangeCallback = (panel: IWidgetWithSession) => {
       return (
         _sender: ISessionContext,
         args: IChangedArgs<
@@ -102,12 +102,13 @@ const KernelUsage = (props: {
     };
 
     const notebookChangeCallback = (
-      sender: INotebookTracker,
-      panel: NotebookPanel | null
+      _: KernelWidgetTracker,
+      panel: IWidgetWithSession | null
     ) => {
       if (panel === null) {
         // Ideally we would switch to a new "select a notebook to get kernel
         // usage" screen instead of showing outdated info.
+        setKernelId(undefined);
         return;
       }
       if (kernelChangeCallback) {
@@ -131,9 +132,9 @@ const KernelUsage = (props: {
         }
       }
     };
-    props.currentNotebookChanged.connect(notebookChangeCallback);
+    props.currentChanged.connect(notebookChangeCallback);
     return () => {
-      props.currentNotebookChanged.disconnect(notebookChangeCallback);
+      props.currentChanged.disconnect(notebookChangeCallback);
       // In the ideal world we would disconnect kernelChangeCallback from
       // last panel here, but this can lead to a race condition. Instead,
       // we make sure there is ever only one callback active by holding
@@ -239,31 +240,27 @@ const KernelUsage = (props: {
 };
 
 export class KernelUsageWidget extends ReactWidget {
-  private _widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
-  private _currentNotebookChanged: ISignal<
-    INotebookTracker,
-    NotebookPanel | null
+  private _currentChanged: ISignal<
+    KernelWidgetTracker,
+    IWidgetWithSession | null
   >;
   private _panel: KernelUsagePanel;
   private _trans: TranslationBundle;
   constructor(props: {
-    widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
-    currentNotebookChanged: ISignal<INotebookTracker, NotebookPanel | null>;
+    currentChanged: ISignal<KernelWidgetTracker, IWidgetWithSession | null>;
     panel: KernelUsagePanel;
     trans: TranslationBundle;
   }) {
     super();
-    this._widgetAdded = props.widgetAdded;
-    this._currentNotebookChanged = props.currentNotebookChanged;
+    this._currentChanged = props.currentChanged;
     this._panel = props.panel;
     this._trans = props.trans;
   }
 
   protected render(): React.ReactElement<any> {
     return (
       <KernelUsage
-        widgetAdded={this._widgetAdded}
-        currentNotebookChanged={this._currentNotebookChanged}
+        currentChanged={this._currentChanged}
         panel={this._panel}
         trans={this._trans}
       />

yarn.lock

@@ -383,6 +383,29 @@
     react "^17.0.1"
     y-codemirror "^3.0.1"
 
+"@jupyterlab/[email protected]":
+  version "3.5.1"
+  resolved "https://registry.yarnpkg.com/@jupyterlab/console/-/console-3.5.1.tgz#175d6d1b7f103830dde4d7762a174836ee5bce3e"
+  integrity sha512-KQTZjeRg4jMRKie6W7wZxSRRTTRvj4H6jBev16ame7f/jR/pzK5Wt7bZ6nRg8h6J2oxqjsJ5gJ6aqjGPaotuWQ==
+  dependencies:
+    "@jupyterlab/apputils" "^3.5.1"
+    "@jupyterlab/cells" "^3.5.1"
+    "@jupyterlab/codeeditor" "^3.5.1"
+    "@jupyterlab/coreutils" "^5.5.1"
+    "@jupyterlab/nbformat" "^3.5.1"
+    "@jupyterlab/observables" "^4.5.1"
+    "@jupyterlab/rendermime" "^3.5.1"
+    "@jupyterlab/services" "^6.5.1"
+    "@jupyterlab/translation" "^3.5.1"
+    "@jupyterlab/ui-components" "^3.5.1"
+    "@lumino/algorithm" "^1.9.0"
+    "@lumino/coreutils" "^1.11.0"
+    "@lumino/disposable" "^1.10.0"
+    "@lumino/dragdrop" "^1.13.0"
+    "@lumino/messaging" "^1.10.0"
+    "@lumino/signaling" "^1.10.0"
+    "@lumino/widgets" "^1.33.0"
+
 "@jupyterlab/coreutils@^5.5.1":
   version "5.5.1"
   resolved "https://registry.npmjs.org/@jupyterlab/coreutils/-/coreutils-5.5.1.tgz"