eclipse-theia · msujew · Dec 14, 2023 · Dec 12, 2023
diff --git a/packages/notebook/src/browser/notebook-editor-widget.tsx b/packages/notebook/src/browser/notebook-editor-widget.tsx
@@ -28,6 +28,7 @@ import { inject, injectable, interfaces, postConstruct } from '@theia/core/share
 import { Emitter } from '@theia/core/shared/vscode-languageserver-protocol';
 import { NotebookEditorWidgetService } from './service/notebook-editor-widget-service';
 import { NotebookMainToolbarRenderer } from './view/notebook-main-toolbar';
+import { Deferred } from '@theia/core/lib/common/promise-util';
 
 export const NotebookEditorWidgetContainerFactory = Symbol('NotebookEditorWidgetContainerFactory');
 
@@ -82,11 +83,16 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
 
     protected readonly renderers = new Map<CellKind, CellRenderer>();
     protected _model?: NotebookModel;
+    protected _ready: Deferred<NotebookModel> = new Deferred();
 
     get notebookType(): string {
         return this.props.notebookType;
     }
 
+    get ready(): Promise<NotebookModel> {
+        return this._ready.promise;
+    }
+
     get model(): NotebookModel | undefined {
         return this._model;
     }
@@ -103,17 +109,17 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
 
         this.renderers.set(CellKind.Markup, this.markdownCellRenderer);
         this.renderers.set(CellKind.Code, this.codeCellRenderer);
-        this.waitForData();
-
+        this._ready.resolve(this.waitForData());
     }
 
-    protected async waitForData(): Promise<void> {
+    protected async waitForData(): Promise<NotebookModel> {
         this._model = await this.props.notebookData;
         this.saveable.delegate = this._model;
         this.toDispose.push(this._model);
         // Ensure that the model is loaded before adding the editor
         this.notebookEditorService.addNotebookEditor(this);
         this.update();
+        return this._model;
     }
 
     protected override onActivateRequest(msg: Message): void {
@@ -130,11 +136,11 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
     }
 
     undo(): void {
-        this.model?.undo();
+        this._model?.undo();
     }
 
     redo(): void {
-        this.model?.redo();
+        this._model?.redo();
     }
 
     protected render(): ReactNode {

diff --git a/packages/notebook/src/browser/service/notebook-editor-widget-service.ts b/packages/notebook/src/browser/service/notebook-editor-widget-service.ts
@@ -19,13 +19,13 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Disposable, DisposableCollection, Emitter } from '@theia/core';
+import { Emitter } from '@theia/core';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
 import { ApplicationShell } from '@theia/core/lib/browser';
 import { NotebookEditorWidget } from '../notebook-editor-widget';
 
 @injectable()
-export class NotebookEditorWidgetService implements Disposable {
+export class NotebookEditorWidgetService {
 
     @inject(ApplicationShell)
     protected applicationShell: ApplicationShell;
@@ -40,27 +40,22 @@ export class NotebookEditorWidgetService implements Disposable {
     private readonly onDidChangeFocusedEditorEmitter = new Emitter<NotebookEditorWidget | undefined>();
     readonly onDidChangeFocusedEditor = this.onDidChangeFocusedEditorEmitter.event;
 
-    private readonly toDispose = new DisposableCollection();
-
     focusedEditor?: NotebookEditorWidget = undefined;
 
     @postConstruct()
     protected init(): void {
-        this.toDispose.push(this.applicationShell.onDidChangeActiveWidget(event => {
-            if (event.newValue instanceof NotebookEditorWidget && event.newValue !== this.focusedEditor) {
-                this.focusedEditor = event.newValue;
-                this.onDidChangeFocusedEditorEmitter.fire(this.focusedEditor);
-            } else {
+        this.applicationShell.onDidChangeActiveWidget(event => {
+            if (event.newValue instanceof NotebookEditorWidget) {
+                if (event.newValue !== this.focusedEditor) {
+                    this.focusedEditor = event.newValue;
+                    this.onDidChangeFocusedEditorEmitter.fire(this.focusedEditor);
+                }
+            } else if (event.newValue) {
+                // Only unfocus editor if a new widget has been focused
+                this.focusedEditor = undefined;
                 this.onDidChangeFocusedEditorEmitter.fire(undefined);
             }
-        }));
-    }
-
-    dispose(): void {
-        this.onNotebookEditorAddEmitter.dispose();
-        this.onNotebookEditorRemoveEmitter.dispose();
-        this.onDidChangeFocusedEditorEmitter.dispose();
-        this.toDispose.dispose();
+        });
     }
 
     // --- editor management

diff --git a/packages/notebook/src/browser/service/notebook-kernel-quick-pick-service.ts b/packages/notebook/src/browser/service/notebook-kernel-quick-pick-service.ts
@@ -18,13 +18,14 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-import { ArrayUtils, Command, CommandService, DisposableCollection, Event, nls, QuickInputButton, QuickInputService, QuickPickInput, QuickPickItem, URI, } from '@theia/core';
+import { ArrayUtils, CommandService, DisposableCollection, Event, nls, QuickInputButton, QuickInputService, QuickPickInput, QuickPickItem, URI, } from '@theia/core';
 import { inject, injectable } from '@theia/core/shared/inversify';
 import { NotebookKernelService, NotebookKernel, NotebookKernelMatchResult, SourceCommand } from './notebook-kernel-service';
 import { NotebookModel } from '../view-model/notebook-model';
 import { NotebookEditorWidget } from '../notebook-editor-widget';
 import { codicon, OpenerService } from '@theia/core/lib/browser';
 import { NotebookKernelHistoryService } from './notebook-kernel-history-service';
+import { NotebookCommand, NotebookModelResource } from '../../common';
 import debounce = require('@theia/core/shared/lodash.debounce');
 
 export const JUPYTER_EXTENSION_ID = 'ms-toolsai.jupyter';
@@ -43,7 +44,7 @@ function isSourcePick(item: QuickPickInput<QuickPickItem>): item is SourcePick {
 }
 type InstallExtensionPick = QuickPickItem & { extensionIds: string[] };
 
-type KernelSourceQuickPickItem = QuickPickItem & { command: Command; documentation?: string };
+type KernelSourceQuickPickItem = QuickPickItem & { command: NotebookCommand; documentation?: string };
 function isKernelSourceQuickPickItem(item: QuickPickItem): item is KernelSourceQuickPickItem {
     return 'command' in item;
 }
@@ -469,10 +470,8 @@ export class NotebookKernelQuickPickService {
         quickPick.show();
     }
 
-    private async executeCommand<T>(notebook: NotebookModel, command: string | Command): Promise<T | undefined | void> {
-        const id = typeof command === 'string' ? command : command.id;
-
-        return this.commandService.executeCommand(id, { uri: notebook.uri });
-
+    private async executeCommand<T>(notebook: NotebookModel, command: NotebookCommand): Promise<T | undefined | void> {
+        const args = (command.arguments || []).concat([NotebookModelResource.create(notebook.uri)]);
+        return this.commandService.executeCommand(command.id, ...args);
     }
 }
diff --git a/packages/notebook/src/common/notebook-common.ts b/packages/notebook/src/common/notebook-common.ts
@@ -14,11 +14,19 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { Command, URI } from '@theia/core';
+import { Command, URI, isObject } from '@theia/core';
 import { MarkdownString } from '@theia/core/lib/common/markdown-rendering/markdown-string';
 import { BinaryBuffer } from '@theia/core/lib/common/buffer';
 import { UriComponents } from '@theia/core/lib/common/uri';
 
+export interface NotebookCommand {
+    id: string;
+    title?: string;
+    tooltip?: string;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    arguments?: any[];
+}
+
 export enum CellKind {
     Markup = 1,
     Code = 2
@@ -159,6 +167,19 @@ export interface NotebookCellContentChangeEvent {
     readonly index: number;
 }
 
+export interface NotebookModelResource {
+    notebookModelUri: URI;
+}
+
+export namespace NotebookModelResource {
+    export function is(item: unknown): item is NotebookModelResource {
+        return isObject<NotebookModelResource>(item) && item.notebookModelUri instanceof URI;
+    }
+    export function create(uri: URI): NotebookModelResource {
+        return { notebookModelUri: uri };
+    }
+}
+
 export enum NotebookCellExecutionState {
     Unconfirmed = 1,
     Pending = 2,

diff --git a/packages/plugin-ext/src/main/browser/notebooks/notebook-documents-and-editors-main.ts b/packages/plugin-ext/src/main/browser/notebooks/notebook-documents-and-editors-main.ts
@@ -23,14 +23,15 @@ import { interfaces } from '@theia/core/shared/inversify';
 import { UriComponents } from '@theia/core/lib/common/uri';
 import { NotebookEditorWidget, NotebookService, NotebookEditorWidgetService } from '@theia/notebook/lib/browser';
 import { NotebookModel } from '@theia/notebook/lib/browser/view-model/notebook-model';
-import { MAIN_RPC_CONTEXT, NotebookDocumentsAndEditorsDelta, NotebookDocumentsAndEditorsMain, NotebookModelAddedData, NotebooksExt } from '../../../common';
+import { MAIN_RPC_CONTEXT, NotebookDocumentsAndEditorsDelta, NotebookDocumentsAndEditorsMain, NotebookEditorAddData, NotebookModelAddedData, NotebooksExt } from '../../../common';
 import { RPCProtocol } from '../../../common/rpc-protocol';
 import { NotebookDto } from './notebook-dto';
 import { WidgetManager } from '@theia/core/lib/browser';
 import { NotebookEditorsMainImpl } from './notebook-editors-main';
 import { NotebookDocumentsMainImpl } from './notebook-documents-main';
 import { diffMaps, diffSets } from '../../../common/collections';
 import { Mutex } from 'async-mutex';
+import throttle = require('@theia/core/shared/lodash.throttle');
 
 interface NotebookAndEditorDelta {
     removedDocuments: UriComponents[];
@@ -106,12 +107,12 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
         this.notebookEditorService = container.get(NotebookEditorWidgetService);
         this.WidgetManager = container.get(WidgetManager);
 
-        this.notebookService.onDidAddNotebookDocument(async () => this.updateState(), this, this.disposables);
-        this.notebookService.onDidRemoveNotebookDocument(async () => this.updateState(), this, this.disposables);
+        this.notebookService.onDidAddNotebookDocument(async () => this.throttleStateUpdate(), this, this.disposables);
+        this.notebookService.onDidRemoveNotebookDocument(async () => this.throttleStateUpdate(), this, this.disposables);
         // this.WidgetManager.onActiveEditorChanged(() => this.updateState(), this, this.disposables);
         this.notebookEditorService.onDidAddNotebookEditor(async editor => this.handleEditorAdd(editor), this, this.disposables);
         this.notebookEditorService.onDidRemoveNotebookEditor(async editor => this.handleEditorRemove(editor), this, this.disposables);
-        this.notebookEditorService.onDidChangeFocusedEditor(async editor => this.updateState(editor), this, this.disposables);
+        this.notebookEditorService.onDidChangeFocusedEditor(async editor => this.throttleStateUpdate(editor), this, this.disposables);
     }
 
     dispose(): void {
@@ -129,16 +130,18 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
         } else {
             this.editorListeners.set(editor.id, [disposable]);
         }
-        await this.updateState();
+        await this.throttleStateUpdate();
     }
 
     private handleEditorRemove(editor: NotebookEditorWidget): void {
         const listeners = this.editorListeners.get(editor.id);
         listeners?.forEach(listener => listener.dispose());
         this.editorListeners.delete(editor.id);
-        this.updateState();
+        this.throttleStateUpdate();
     }
 
+    private throttleStateUpdate = throttle((focusedEditor?: NotebookEditorWidget) => this.updateState(focusedEditor), 100);
+
     private async updateState(focusedEditor?: NotebookEditorWidget): Promise<void> {
         await this.updateMutex.runExclusive(async () => this.doUpdateState(focusedEditor));
     }
@@ -149,9 +152,7 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
         const visibleEditorsMap = new Map<string, NotebookEditorWidget>();
 
         for (const editor of this.notebookEditorService.getNotebookEditors()) {
-            if (editor.model) {
-                editors.set(editor.id, editor);
-            }
+            editors.set(editor.id, editor);
         }
 
         const activeNotebookEditor = this.notebookEditorService.focusedEditor;
@@ -167,7 +168,7 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
 
         const notebookEditors = this.WidgetManager.getWidgets(NotebookEditorWidget.ID) as NotebookEditorWidget[];
         for (const notebookEditor of notebookEditors) {
-            if (notebookEditor?.model && editors.has(notebookEditor.id) && notebookEditor.isVisible) {
+            if (editors.has(notebookEditor.id) && notebookEditor.isVisible) {
                 visibleEditorsMap.set(notebookEditor.id, notebookEditor);
             }
         }
@@ -191,7 +192,7 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
             newActiveEditor: delta.newActiveEditor,
             visibleEditors: delta.visibleEditors,
             addedDocuments: delta.addedDocuments.map(NotebooksAndEditorsMain.asModelAddData),
-            // addedEditors: delta.addedEditors.map(this.asEditorAddData, this),
+            addedEditors: delta.addedEditors.map(NotebooksAndEditorsMain.asEditorAddData),
         };
 
         // send to extension FIRST
@@ -235,4 +236,17 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
             cells: e.cells.map(NotebookDto.toNotebookCellDto)
         };
     }
+
+    private static asEditorAddData(notebookEditor: NotebookEditorWidget): NotebookEditorAddData {
+        const uri = notebookEditor.getResourceUri();
+        if (!uri) {
+            throw new Error('Notebook editor without resource URI');
+        }
+        return {
+            id: notebookEditor.id,
+            documentUri: uri.toComponents(),
+            selections: [],
+            visibleRanges: []
+        };
+    }
 }
diff --git a/packages/plugin-ext/src/plugin/command-registry.ts b/packages/plugin-ext/src/plugin/command-registry.ts
@@ -204,11 +204,16 @@ export class CommandsConverter {
 
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     private executeSafeCommand<R>(...args: any[]): PromiseLike<R | undefined> {
-        const command = this.commandsMap.get(args[0]);
+        const handle = args[0];
+        if (typeof handle !== 'number') {
+            return Promise.reject(`Invalid handle ${handle}`);
+        }
+        const command = this.commandsMap.get(handle);
         if (!command || !command.command) {
-            return Promise.reject(`command ${args[0]} not found`);
+            return Promise.reject(`Safe command with handle ${handle} not found`);
         }
-        return this.commands.executeCommand(command.command, ...(command.arguments || []));
+        const allArgs = (command.arguments ?? []).concat(args.slice(1));
+        return this.commands.executeCommand(command.command, ...allArgs);
     }
 
 }
diff --git a/packages/plugin-ext/src/plugin/notebook/notebooks.ts b/packages/plugin-ext/src/plugin/notebook/notebooks.ts
@@ -36,6 +36,7 @@ import { NotebookDocument } from './notebook-document';
 import { NotebookEditor } from './notebook-editor';
 import { EditorsAndDocumentsExtImpl } from '../editors-and-documents';
 import { DocumentsExtImpl } from '../documents';
+import { NotebookModelResource } from '@theia/notebook/lib/common';
 
 export class NotebooksExtImpl implements NotebooksExt {
 
@@ -82,11 +83,12 @@ export class NotebooksExtImpl implements NotebooksExt {
         this.notebookEditors = rpc.getProxy(PLUGIN_RPC_CONTEXT.NOTEBOOK_EDITORS_MAIN);
 
         commands.registerArgumentProcessor({
-            processArgument: (arg: { uri: URI }) => {
-                if (arg && arg.uri && this.documents.has(arg.uri.toString())) {
-                    return this.documents.get(arg.uri.toString())?.apiNotebook;
+            processArgument: arg => {
+                if (NotebookModelResource.is(arg)) {
+                    return this.documents.get(arg.notebookModelUri.toString())?.apiNotebook;
+                } else {
+                    return arg;
                 }
-                return arg;
             }
         });
     }