Improve notebook kernel/renderer messaging (#13401)

packages/notebook/src/browser/notebook-editor-widget.tsx

@@ -46,6 +46,11 @@ export function createNotebookEditorWidgetContainer(parent: interfaces.Container
 
 const NotebookEditorProps = Symbol('NotebookEditorProps');
 
+interface RenderMessage {
+    rendererId: string;
+    message: unknown;
+}
+
 export interface NotebookEditorProps {
     uri: URI,
     readonly notebookType: string,
@@ -87,6 +92,18 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
     protected readonly onDidChangeReadOnlyEmitter = new Emitter<boolean | MarkdownString>();
     readonly onDidChangeReadOnly = this.onDidChangeReadOnlyEmitter.event;
 
+    protected readonly onPostKernelMessageEmitter = new Emitter<unknown>();
+    readonly onPostKernelMessage = this.onPostKernelMessageEmitter.event;
+
+    protected readonly onDidPostKernelMessageEmitter = new Emitter<unknown>();
+    readonly onDidPostKernelMessage = this.onDidPostKernelMessageEmitter.event;
+
+    protected readonly onPostRendererMessageEmitter = new Emitter<RenderMessage>();
+    readonly onPostRendererMessage = this.onPostRendererMessageEmitter.event;
+
+    protected readonly onDidReceiveKernelMessageEmitter = new Emitter<unknown>();
+    readonly onDidRecieveKernelMessage = this.onDidReceiveKernelMessageEmitter.event;
+
     protected readonly renderers = new Map<CellKind, CellRenderer>();
     protected _model?: NotebookModel;
     protected _ready: Deferred<NotebookModel> = new Deferred();
@@ -190,4 +207,24 @@ export class NotebookEditorWidget extends ReactWidget implements Navigatable, Sa
         super.onAfterDetach(msg);
         this.notebookEditorService.removeNotebookEditor(this);
     }
+
+    postKernelMessage(message: unknown): void {
+        this.onDidPostKernelMessageEmitter.fire(message);
+    }
+
+    postRendererMessage(rendererId: string, message: unknown): void {
+        this.onPostRendererMessageEmitter.fire({ rendererId, message });
+    }
+
+    recieveKernelMessage(message: unknown): void {
+        this.onDidReceiveKernelMessageEmitter.fire(message);
+    }
+
+    override dispose(): void {
+        this.onDidChangeModelEmitter.dispose();
+        this.onDidPostKernelMessageEmitter.dispose();
+        this.onDidReceiveKernelMessageEmitter.dispose();
+        this.onPostRendererMessageEmitter.dispose();
+        super.dispose();
+    }
 }

packages/notebook/src/browser/notebook-renderer-registry.ts

@@ -30,6 +30,11 @@ export interface NotebookRendererInfo {
     readonly requiresMessaging: boolean;
 }
 
+export interface NotebookPreloadInfo {
+    readonly type: string;
+    readonly entrypoint: string;
+}
+
 @injectable()
 export class NotebookRendererRegistry {
 
@@ -39,6 +44,12 @@ export class NotebookRendererRegistry {
         return this._notebookRenderers;
     }
 
+    private readonly _staticNotebookPreloads: NotebookPreloadInfo[] = [];
+
+    get staticNotebookPreloads(): readonly NotebookPreloadInfo[] {
+        return this._staticNotebookPreloads;
+    }
+
     registerNotebookRenderer(type: NotebookRendererDescriptor, basePath: string): Disposable {
         let entrypoint;
         if (typeof type.entrypoint === 'string') {
@@ -62,5 +73,13 @@ export class NotebookRendererRegistry {
             this._notebookRenderers.splice(this._notebookRenderers.findIndex(renderer => renderer.id === type.id), 1);
         });
     }
+
+    registerStaticNotebookPreload(type: string, entrypoint: string, basePath: string): Disposable {
+        const staticPreload = { type, entrypoint: new Path(basePath).join(entrypoint).toString() };
+        this._staticNotebookPreloads.push(staticPreload);
+        return Disposable.create(() => {
+            this._staticNotebookPreloads.splice(this._staticNotebookPreloads.indexOf(staticPreload), 1);
+        });
+    }
 }
 

packages/notebook/src/browser/renderers/cell-output-webview.ts

@@ -16,10 +16,11 @@
 
 import { Disposable } from '@theia/core';
 import { NotebookCellModel } from '../view-model/notebook-cell-model';
+import { NotebookModel } from '../view-model/notebook-model';
 
 export const CellOutputWebviewFactory = Symbol('outputWebviewFactory');
 
-export type CellOutputWebviewFactory = (cell: NotebookCellModel) => Promise<CellOutputWebview>;
+export type CellOutputWebviewFactory = (cell: NotebookCellModel, notebook: NotebookModel) => Promise<CellOutputWebview>;
 
 export interface CellOutputWebview extends Disposable {
 

packages/notebook/src/browser/service/notebook-kernel-service.ts

@@ -54,6 +54,11 @@ export interface NotebookKernel {
     // ID of the extension providing this kernel
     readonly extensionId: string;
 
+    readonly localResourceRoot: URI;
+    readonly preloadUris: URI[];
+    readonly preloadProvides: string[];
+
+    readonly handle: number;
     label: string;
     description?: string;
     detail?: string;

packages/notebook/src/browser/service/notebook-renderer-messaging-service.ts

@@ -19,8 +19,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Emitter } from '@theia/core';
-import { injectable } from '@theia/core/shared/inversify';
+import { inject, injectable } from '@theia/core/shared/inversify';
 import { Disposable } from '@theia/core/shared/vscode-languageserver-protocol';
+import { NotebookEditorWidgetService } from './notebook-editor-widget-service';
 
 interface RendererMessage {
     editorId: string;
@@ -50,6 +51,9 @@ export class NotebookRendererMessagingService implements Disposable {
     private readonly willActivateRendererEmitter = new Emitter<string>();
     readonly onWillActivateRenderer = this.willActivateRendererEmitter.event;
 
+    @inject(NotebookEditorWidgetService)
+    private readonly editorWidgetService: NotebookEditorWidgetService;
+
     private readonly activations = new Map<string /* rendererId */, undefined | RendererMessage[]>();
     private readonly scopedMessaging = new Map<string /* editorId */, RendererMessaging>();
 
@@ -86,6 +90,10 @@ export class NotebookRendererMessagingService implements Disposable {
 
         const messaging: RendererMessaging = {
             postMessage: (rendererId, message) => this.postMessage(editorId, rendererId, message),
+            receiveMessage: async (rendererId, message) => {
+                this.editorWidgetService.getNotebookEditor(editorId)?.postRendererMessage(rendererId, message);
+                return true;
+            },
             dispose: () => this.scopedMessaging.delete(editorId),
         };
 

packages/notebook/src/browser/view/notebook-code-cell-view.tsx

@@ -61,7 +61,7 @@ export class NotebookCodeCellRenderer implements CellRenderer {
                 </div>
             </div>
             <div className='theia-notebook-cell-with-sidebar'>
-                <NotebookCodeCellOutputs cell={cell} outputWebviewFactory={this.cellOutputWebviewFactory}
+                <NotebookCodeCellOutputs cell={cell} notebook={notebookModel} outputWebviewFactory={this.cellOutputWebviewFactory}
                     renderSidebar={() =>
                         this.notebookCellToolbarFactory.renderSidebar(NotebookCellActionContribution.OUTPUT_SIDEBAR_MENU, notebookModel, cell, cell.outputs[0])} />
             </div>
@@ -166,6 +166,7 @@ export class NotebookCodeCellStatus extends React.Component<NotebookCodeCellStat
 
 interface NotebookCellOutputProps {
     cell: NotebookCellModel;
+    notebook: NotebookModel;
     outputWebviewFactory: CellOutputWebviewFactory;
     renderSidebar: () => React.ReactNode;
 }
@@ -182,14 +183,14 @@ export class NotebookCodeCellOutputs extends React.Component<NotebookCellOutputP
     }
 
     override async componentDidMount(): Promise<void> {
-        const { cell, outputWebviewFactory } = this.props;
+        const { cell, notebook, outputWebviewFactory } = this.props;
         this.toDispose.push(cell.onDidChangeOutputs(async () => {
             if (!this.outputsWebviewPromise && cell.outputs.length > 0) {
-                this.outputsWebviewPromise = outputWebviewFactory(cell).then(webview => {
+                this.outputsWebviewPromise = outputWebviewFactory(cell, notebook).then(webview => {
                     this.outputsWebview = webview;
                     this.forceUpdate();
                     return webview;
-                    });
+                });
                 this.forceUpdate();
             } else if (this.outputsWebviewPromise && cell.outputs.length === 0 && cell.internalMetadata.runEndTime) {
                 (await this.outputsWebviewPromise).dispose();
@@ -199,7 +200,7 @@ export class NotebookCodeCellOutputs extends React.Component<NotebookCellOutputP
             }
         }));
         if (cell.outputs.length > 0) {
-            this.outputsWebviewPromise = outputWebviewFactory(cell).then(webview => {
+            this.outputsWebviewPromise = outputWebviewFactory(cell, notebook).then(webview => {
                 this.outputsWebview = webview;
                 this.forceUpdate();
                 return webview;

packages/plugin-ext/src/common/plugin-api-rpc.ts

@@ -2288,7 +2288,7 @@ export const MAIN_RPC_CONTEXT = {
     NOTEBOOKS_EXT: createProxyIdentifier<NotebooksExt>('NotebooksExt'),
     NOTEBOOK_DOCUMENTS_EXT: createProxyIdentifier<NotebookDocumentsExt>('NotebookDocumentsExt'),
     NOTEBOOK_EDITORS_EXT: createProxyIdentifier<NotebookEditorsExt>('NotebookEditorsExt'),
-    NOTEBOOK_RENDERERS_EXT: createProxyIdentifier<NotebookRenderersExt>('NotebooksExt'),
+    NOTEBOOK_RENDERERS_EXT: createProxyIdentifier<NotebookRenderersExt>('NotebooksRenderersExt'),
     NOTEBOOK_KERNELS_EXT: createProxyIdentifier<NotebookKernelsExt>('NotebookKernelsExt'),
     TERMINAL_EXT: createProxyIdentifier<TerminalServiceExt>('TerminalServiceExt'),
     OUTPUT_CHANNEL_REGISTRY_EXT: createProxyIdentifier<OutputChannelRegistryExt>('OutputChannelRegistryExt'),
@@ -2488,7 +2488,7 @@ export interface NotebookKernelDto {
     id: string;
     notebookType: string;
     extensionId: string;
-    // extensionLocation: UriComponents;
+    extensionLocation: UriComponents;
     label: string;
     detail?: string;
     description?: string;

packages/plugin-ext/src/common/plugin-protocol.ts

@@ -103,6 +103,7 @@ export interface PluginPackageContribution {
     terminal?: PluginPackageTerminal;
     notebooks?: PluginPackageNotebook[];
     notebookRenderer?: PluginNotebookRendererContribution[];
+    notebookPreload?: PluginPackageNotebookPreload[];
 }
 
 export interface PluginPackageNotebook {
@@ -120,6 +121,11 @@ export interface PluginNotebookRendererContribution {
     readonly requiresMessaging?: 'always' | 'optional' | 'never'
 }
 
+export interface PluginPackageNotebookPreload {
+    type: string;
+    entrypoint: string;
+}
+
 export interface PluginPackageAuthenticationProvider {
     id: string;
     label: string;
@@ -610,8 +616,8 @@ export interface PluginContribution {
     terminalProfiles?: TerminalProfile[];
     notebooks?: NotebookContribution[];
     notebookRenderer?: NotebookRendererContribution[];
+    notebookPreload?: notebookPreloadContribution[];
 }
-
 export interface NotebookContribution {
     type: string;
     displayName: string;
@@ -627,6 +633,11 @@ export interface NotebookRendererContribution {
     readonly requiresMessaging?: 'always' | 'optional' | 'never'
 }
 
+export interface notebookPreloadContribution {
+    type: string;
+    entrypoint: string;
+}
+
 export interface AuthenticationProviderInformation {
     id: string;
     label: string;

packages/plugin-ext/src/hosted/node/scanners/scanner-theia.ts

@@ -341,13 +341,19 @@ export class TheiaPluginScanner extends AbstractPluginScanner {
         try {
             contributions.notebooks = rawPlugin.contributes.notebooks;
         } catch (err) {
-            console.error(`Could not read '${rawPlugin.name}' contribution 'notebooks'.`, rawPlugin.contributes.authentication, err);
+            console.error(`Could not read '${rawPlugin.name}' contribution 'notebooks'.`, rawPlugin.contributes.notebooks, err);
         }
 
         try {
             contributions.notebookRenderer = rawPlugin.contributes.notebookRenderer;
         } catch (err) {
-            console.error(`Could not read '${rawPlugin.name}' contribution 'notebooks'.`, rawPlugin.contributes.authentication, err);
+            console.error(`Could not read '${rawPlugin.name}' contribution 'notebook-renderer'.`, rawPlugin.contributes.notebookRenderer, err);
+        }
+
+        try {
+            contributions.notebookPreload = rawPlugin.contributes.notebookPreload;
+        } catch (err) {
+            console.error(`Could not read '${rawPlugin.name}' contribution 'notebooks-preload'.`, rawPlugin.contributes.notebookPreload, err);
         }
 
         try {

packages/plugin-ext/src/main/browser/notebooks/notebook-kernels-main.ts

@@ -23,7 +23,10 @@ import { UriComponents } from '@theia/core/lib/common/uri';
 import { LanguageService } from '@theia/core/lib/browser/language-service';
 import { CellExecuteUpdateDto, CellExecutionCompleteDto, MAIN_RPC_CONTEXT, NotebookKernelDto, NotebookKernelsExt, NotebookKernelsMain } from '../../../common';
 import { RPCProtocol } from '../../../common/rpc-protocol';
-import { CellExecution, NotebookExecutionStateService, NotebookKernelChangeEvent, NotebookKernelService, NotebookService } from '@theia/notebook/lib/browser';
+import {
+    CellExecution, NotebookEditorWidgetService, NotebookExecutionStateService,
+    NotebookKernelChangeEvent, NotebookKernelService, NotebookService
+} from '@theia/notebook/lib/browser';
 import { combinedDisposable } from '@theia/monaco-editor-core/esm/vs/base/common/lifecycle';
 import { interfaces } from '@theia/core/shared/inversify';
 import { NotebookKernelSourceAction } from '@theia/notebook/lib/common';
@@ -54,7 +57,7 @@ abstract class NotebookKernel {
         return this.preloads.map(p => p.provides).flat();
     }
 
-    constructor(data: NotebookKernelDto, private languageService: LanguageService) {
+    constructor(public readonly handle: number, data: NotebookKernelDto, private languageService: LanguageService) {
         this.id = data.id;
         this.viewType = data.notebookType;
         this.extensionId = data.extensionId;
@@ -65,6 +68,7 @@ abstract class NotebookKernel {
         this.detail = data.detail;
         this.supportedLanguages = (data.supportedLanguages && data.supportedLanguages.length > 0) ? data.supportedLanguages : languageService.languages.map(lang => lang.id);
         this.implementsExecutionOrder = data.supportsExecutionOrder ?? false;
+        this.localResourceRoot = URI.fromComponents(data.extensionLocation);
         this.preloads = data.preloads?.map(u => ({ uri: URI.fromComponents(u.uri), provides: u.provides })) ?? [];
     }
 
@@ -125,6 +129,7 @@ export class NotebookKernelsMainImpl implements NotebookKernelsMain {
     private notebookService: NotebookService;
     private languageService: LanguageService;
     private notebookExecutionStateService: NotebookExecutionStateService;
+    private notebookEditorWidgetService: NotebookEditorWidgetService;
 
     private readonly executions = new Map<number, CellExecution>();
 
@@ -138,10 +143,46 @@ export class NotebookKernelsMainImpl implements NotebookKernelsMain {
         this.notebookExecutionStateService = container.get(NotebookExecutionStateService);
         this.notebookService = container.get(NotebookService);
         this.languageService = container.get(LanguageService);
+        this.notebookEditorWidgetService = container.get(NotebookEditorWidgetService);
+
+        this.notebookEditorWidgetService.onDidAddNotebookEditor(editor => {
+            editor.onDidRecieveKernelMessage(async message => {
+                const kernel = this.notebookKernelService.getSelectedOrSuggestedKernel(editor.model!);
+                if (kernel) {
+                    this.proxy.$acceptKernelMessageFromRenderer(kernel.handle, editor.id, message);
+                }
+            });
+        });
     }
 
-    $postMessage(handle: number, editorId: string | undefined, message: unknown): Promise<boolean> {
-        throw new Error('Method not implemented.');
+    async $postMessage(handle: number, editorId: string | undefined, message: unknown): Promise<boolean> {
+        const tuple = this.kernels.get(handle);
+        if (!tuple) {
+            throw new Error('kernel already disposed');
+        }
+        const [kernel] = tuple;
+        let didSend = false;
+        for (const editor of this.notebookEditorWidgetService.getNotebookEditors()) {
+            if (!editor.model) {
+                continue;
+            }
+            if (this.notebookKernelService.getMatchingKernel(editor.model).selected !== kernel) {
+                // different kernel
+                continue;
+            }
+            if (editorId === undefined) {
+                // all editors
+                editor.postKernelMessage(message);
+                didSend = true;
+            } else if (editor.id === editorId) {
+                // selected editors
+                editor.postKernelMessage(message);
+                didSend = true;
+                break;
+            }
+        }
+        return didSend;
+
     }
 
     async $addKernel(handle: number, data: NotebookKernelDto): Promise<void> {
@@ -153,7 +194,7 @@ export class NotebookKernelsMainImpl implements NotebookKernelsMain {
             async cancelNotebookCellExecution(uri: URI, handles: number[]): Promise<void> {
                 await that.proxy.$cancelCells(handle, uri.toComponents(), handles);
             }
-        }(data, this.languageService);
+        }(handle, data, this.languageService);
 
         const listener = this.notebookKernelService.onDidChangeSelectedKernel(e => {
             if (e.oldKernel === kernel.id) {