add Kernel2#createNotebookRendererCommunication

src/vs/vscode.proposed.d.ts

@@ -1429,6 +1429,45 @@ declare module 'vscode' {
 
 	export type NotebookSelector = NotebookFilter | string | ReadonlyArray<NotebookFilter | string>;
 
+	export interface NotebookRendererCommunication {
+
+		/**
+		 *
+		 */
+		dispose(): void;
+
+		/**
+		 * The editor this object communicates with. A single notebook
+		 * document can have multiple attached webviews and editors, when the
+		 * notebook is split for instance. The editor ID lets you differentiate
+		 * between them.
+		 */
+		readonly editor: NotebookEditor;
+
+		/**
+		 *
+		 */
+		readonly rendererId: string;
+
+		/**
+		 * Fired when the output hosting webview posts a message.
+		 */
+		readonly onDidReceiveMessage: Event<any>;
+		/**
+		 * Post a message to the output hosting webview.
+		 *
+		 * Messages are only delivered if the editor is live.
+		 *
+		 * @param message Body of the message. This must be a string or other json serializable object.
+		 */
+		postMessage(message: any): Thenable<boolean>;
+
+		/**
+		 * Convert a uri for the local file system to one that can be used inside outputs webview.
+		 */
+		asWebviewUri(localResource: Uri): Uri;
+	}
+
 	export interface NotebookKernel2 {
 
 		readonly id: string;
@@ -1443,8 +1482,10 @@ declare module 'vscode' {
 		 */
 		readonly onDidChangeNotebookAssociation: Event<{ notebook: NotebookDocument, selected: boolean }>;
 
-		// kernels can establish IPC channels to (visible) notebook editors
-		// createNotebookCommunication(editor: vscode.NotebookEditor): vscode.NotebookCommunication;
+		// kernels can establish IPC channels to notebook editors
+		// todo@API create per editor or allow to postMessage(EDITOR, message) and onDidReceive: Event<{EDITOR, message}>
+		// todo@API have this global on vscode.notebook?
+		createNotebookRendererCommunication(editor: NotebookEditor, rendererId: string): NotebookRendererCommunication;
 
 		// UI properties (get/set)
 		label: string;

src/vs/workbench/api/browser/mainThreadNotebookEditors.ts

@@ -19,11 +19,15 @@ import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editor
 import { editorGroupToViewColumn } from 'vs/workbench/common/editor';
 import { equals } from 'vs/base/common/objects';
 
-class MainThreadEditor {
+class MainThreadNotebook {
+
+	readonly ipcHandles = new Map<number, string>();
+
 	constructor(
 		readonly editor: INotebookEditor,
 		readonly disposables: DisposableStore
 	) { }
+
 	dispose() {
 		this.disposables.dispose();
 	}
@@ -34,7 +38,7 @@ export class MainThreadNotebookEditors implements MainThreadNotebookEditorsShape
 	private readonly _disposables = new DisposableStore();
 
 	private readonly _proxy: ExtHostNotebookShape;
-	private readonly _mainThreadEditors = new Map<string, MainThreadEditor>();
+	private readonly _mainThreadEditors = new Map<string, MainThreadNotebook>();
 
 	private _currentViewColumnInfo?: INotebookEditorViewColumnInfo;
 
@@ -86,7 +90,20 @@ export class MainThreadNotebookEditors implements MainThreadNotebookEditorsShape
 				});
 			}));
 
-			this._mainThreadEditors.set(editor.getId(), new MainThreadEditor(editor, editorDisposables));
+			editorDisposables.add(editor.onDidReceiveMessage(e => {
+				const handles: number[] = [];
+				for (let [handle, rendererId] of wrapper.ipcHandles) {
+					if (rendererId === e.forRenderer) {
+						handles.push(handle);
+					}
+				}
+				if (handles.length > 0) {
+					this._proxy.$acceptEditorIpcMessage(handles, e.message);
+				}
+			}));
+
+			const wrapper = new MainThreadNotebook(editor, editorDisposables);
+			this._mainThreadEditors.set(editor.getId(), wrapper);
 		}
 	}
 
@@ -128,6 +145,35 @@ export class MainThreadNotebookEditors implements MainThreadNotebookEditorsShape
 		return editor.textModel.applyEdits(cellEdits, true, undefined, () => undefined, undefined);
 	}
 
+	//#region --- IPC
+
+	async $createNotebookIPC(editorId: string, handle: number, rendererId: string): Promise<boolean> {
+		const wrapper = this._mainThreadEditors.get(editorId);
+		if (!wrapper) {
+			return false;
+		}
+		wrapper.ipcHandles.set(handle, rendererId);
+		return true;
+	}
+
+	async $removeNotebookIpc(editorId: string, handle: number) {
+		const wrapper = this._mainThreadEditors.get(editorId);
+		if (wrapper) {
+			wrapper.ipcHandles.delete(handle);
+		}
+	}
+
+	async $postMessage(editorId: string, handle: number, forRendererId: string | undefined, message: unknown): Promise<boolean> {
+		const wrapper = this._mainThreadEditors.get(editorId);
+		if (!wrapper || !wrapper.ipcHandles.has(handle)) {
+			return false;
+		}
+		wrapper.editor.postMessage(forRendererId, message);
+		return true;
+	}
+
+	//#endregion
+
 	async $tryShowNotebookDocument(resource: UriComponents, viewType: string, options: INotebookDocumentShowOptions): Promise<string> {
 		const editorOptions = new NotebookEditorOptions({
 			cellSelections: options.selections,

src/vs/workbench/api/common/extHost.protocol.ts

@@ -892,6 +892,10 @@ export interface MainThreadNotebookEditorsShape extends IDisposable {
 	$removeNotebookEditorDecorationType(key: string): void;
 	$trySetDecorations(id: string, range: ICellRange, decorationKey: string): void;
 	$tryApplyEdits(editorId: string, modelVersionId: number, cellEdits: ICellEditOperation[]): Promise<boolean>
+
+	$createNotebookIPC(editorId: string, handle: number, rendererId: string): Promise<boolean>
+	$removeNotebookIpc(editorId: string, handle: number): void;
+	$postMessage(editorId: string, handle: number, forRendererId: string | undefined, message: unknown): Promise<boolean>;
 }
 
 export interface MainThreadNotebookDocumentsShape extends IDisposable {
@@ -1949,6 +1953,8 @@ export type INotebookEditorViewColumnInfo = Record<string, number>;
 export interface ExtHostNotebookEditorsShape {
 	$acceptEditorPropertiesChanged(id: string, data: INotebookEditorPropertiesChangeData): void;
 	$acceptEditorViewColumns(data: INotebookEditorViewColumnInfo): void;
+
+	$acceptEditorIpcMessage(handles: number[], message: unknown): void;
 }
 
 export interface ExtHostNotebookKernelsShape {

src/vs/workbench/api/common/extHostNotebook.ts

@@ -660,6 +660,65 @@ export class ExtHostNotebookController implements ExtHostNotebookShape {
 		}
 	}
 
+	private readonly _editorIpcObjects = new Map<number, { emitter: Emitter<any> }>();
+
+	createNotebookCommunication(editor: vscode.NotebookEditor, rendererId: string): vscode.NotebookRendererCommunication {
+		// todo@jrieken should this be created for a specific renderer?
+		// how else would this be working when sending/receiving messages
+
+		let actualEditor: ExtHostNotebookEditor | undefined;
+		for (let candidate of this._editors.values()) {
+			if (candidate.editor.editor === editor) {
+				actualEditor = candidate.editor;
+				break;
+			}
+		}
+
+		if (!actualEditor) {
+			throw new Error(`the provided editor is NOT KNOWN`);
+		}
+
+		const editorId = actualEditor.id;
+		const that = this;
+		const handle = this._handlePool++;
+
+		const emitter = new Emitter<unknown>();
+
+		const registration = this._notebookEditorsProxy.$createNotebookIPC(editorId, handle, rendererId);
+
+		const result: vscode.NotebookRendererCommunication = {
+			editor,
+			rendererId,
+			onDidReceiveMessage: emitter.event,
+			dispose(): void {
+				emitter.dispose();
+				that._editorIpcObjects.delete(handle);
+				that._notebookEditorsProxy.$removeNotebookIpc(editorId, handle);
+			},
+			async postMessage(message: unknown) {
+				if (!that._editors.has(editorId)) {
+					return false;
+				}
+				if (!await registration) {
+					return false;
+				}
+				return that._notebookEditorsProxy.$postMessage(editorId, handle, rendererId, message);
+			},
+			asWebviewUri(localResource) {
+				return asWebviewUri(that._webviewInitData, editorId, localResource);
+			}
+		};
+
+		this._editorIpcObjects.set(handle, { emitter });
+		return result;
+	}
+
+	$acceptEditorIpcMessage(handles: number[], message: unknown): void {
+		for (const handle of handles) {
+			this._editorIpcObjects.get(handle)?.emitter.fire(message);
+		}
+	}
+
 	$acceptEditorViewColumns(data: INotebookEditorViewColumnInfo): void {
 		for (const id in data) {
 			const editor = this._editors.get(id);

src/vs/workbench/api/common/extHostNotebookKernels.ts

@@ -114,15 +114,26 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 				_update();
 			},
 			createNotebookCellExecutionTask(cell) {
+				if (isDisposed) {
+					throw new Error('object disposed');
+				}
 				//todo@jrieken
 				return that._extHostNotebook.createNotebookCellExecution(cell.document.uri, cell.index, data.id)!;
 			},
+			createNotebookRendererCommunication(editor, rendererId) {
+				if (isDisposed) {
+					throw new Error('object disposed');
+				}
+				return that._extHostNotebook.createNotebookCommunication(editor, rendererId);
+			},
 			dispose: () => {
-				isDisposed = true;
-				this._kernelData.delete(handle);
-				commandDisposables.dispose();
-				emitter.dispose();
-				this._proxy.$removeKernel(handle);
+				if (!isDisposed) {
+					isDisposed = true;
+					this._kernelData.delete(handle);
+					commandDisposables.dispose();
+					emitter.dispose();
+					this._proxy.$removeKernel(handle);
+				}
 			}
 		};
 	}

src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts

@@ -30,6 +30,7 @@ import { EditorOptions, IEditorPane } from 'vs/workbench/common/editor';
 import { IResourceEditorInput } from 'vs/platform/editor/common/editor';
 import { IConstructorSignature1 } from 'vs/platform/instantiation/common/instantiation';
 import { CellEditorStatusBar } from 'vs/workbench/contrib/notebook/browser/view/renderers/cellWidgets';
+import { INotebookWebviewMessage } from 'vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView';
 
 export const NOTEBOOK_EDITOR_ID = 'workbench.editor.notebook';
 export const NOTEBOOK_DIFF_EDITOR_ID = 'workbench.editor.notebookTextDiffEditor';
@@ -495,6 +496,9 @@ export interface INotebookEditor extends ICommonNotebookEditor {
 	 */
 	hideInset(output: IDisplayOutputViewModel): void;
 
+
+	onDidReceiveMessage: Event<INotebookWebviewMessage>;
+
 	/**
 	 * Send message to the webview for outputs.
 	 */

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -52,7 +52,7 @@ import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/no
 import { errorStateIcon, successStateIcon } from 'vs/workbench/contrib/notebook/browser/notebookIcons';
 import { NotebookCellList } from 'vs/workbench/contrib/notebook/browser/view/notebookCellList';
 import { OutputRenderer } from 'vs/workbench/contrib/notebook/browser/view/output/outputRenderer';
-import { BackLayerWebView } from 'vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView';
+import { BackLayerWebView, INotebookWebviewMessage } from 'vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView';
 import { CellContextKeyManager } from 'vs/workbench/contrib/notebook/browser/view/renderers/cellContextKeys';
 import { CellDragAndDropController } from 'vs/workbench/contrib/notebook/browser/view/renderers/cellDnd';
 import { CodeCellRenderer, ListTopCellToolbar, MarkdownCellRenderer, NotebookCellListDelegate } from 'vs/workbench/contrib/notebook/browser/view/renderers/cellRenderer';
@@ -1063,9 +1063,10 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 				}
 			});
 
-			this._localStore.add(this._webview.onMessage(({ message, forRenderer }) => {
+			this._localStore.add(this._webview.onMessage(e => {
 				if (this.viewModel) {
-					this.notebookService.onDidReceiveMessage(this.viewModel.viewType, this.getId(), forRenderer, message);
+					this.notebookService.onDidReceiveMessage(this.viewModel.viewType, this.getId(), e.forRenderer, e.message);
+					this._onDidReceiveMessage.fire(e);
 				}
 			}));
 
@@ -2297,6 +2298,12 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		return this._outputRenderer;
 	}
 
+	//#region --- webview IPC ----
+
+	private readonly _onDidReceiveMessage = new Emitter<INotebookWebviewMessage>();
+
+	readonly onDidReceiveMessage: Event<INotebookWebviewMessage> = this._onDidReceiveMessage.event;
+
 	postMessage(forRendererId: string | undefined, message: any) {
 		if (this._webview?.isResolved()) {
 			if (forRendererId === undefined) {
@@ -2307,6 +2314,8 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		}
 	}
 
+	//#endregion
+
 	addClassName(className: string) {
 		this._overlayContainer.classList.add(className);
 	}