messages channel between ext and webview for outputs

src/vs/vscode.proposed.d.ts

@@ -1656,6 +1656,19 @@ declare module 'vscode' {
 	export interface NotebookEditor {
 		readonly document: NotebookDocument;
 		viewColumn?: ViewColumn;
+		/**
+		 * 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 serilizable object.
+		 */
+		postMessage(message: any): Thenable<boolean>;
+
 		/**
 		 * Create a notebook cell. The cell is not inserted into current document when created. Extensions should insert the cell into the document by [TextDocument.cells](#TextDocument.cells)
 		 */

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

@@ -12,6 +12,8 @@ import { INotebookTextModel, INotebookMimeTypeSelector, NOTEBOOK_DISPLAY_ORDER,
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 
 export class MainThreadNotebookDocument extends Disposable {
 	private _textModel: NotebookTextModel;
@@ -54,7 +56,9 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 	constructor(
 		extHostContext: IExtHostContext,
 		@INotebookService private _notebookService: INotebookService,
-		@IConfigurationService private readonly configurationService: IConfigurationService
+		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@IEditorService private readonly editorService: IEditorService,
+
 	) {
 		super();
 		this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostNotebook);
@@ -141,6 +145,21 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 	async executeNotebook(viewType: string, uri: URI): Promise<void> {
 		return this._proxy.$executeNotebook(viewType, uri, undefined);
 	}
+
+	async $postMessage(handle: number, value: any): Promise<boolean> {
+
+		const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
+		if (activeEditorPane?.isNotebookEditor) {
+			const notebookEditor = (activeEditorPane as INotebookEditor);
+
+			if (notebookEditor.viewModel?.handle === handle) {
+				notebookEditor.postMessage(value);
+				return true;
+			}
+		}
+
+		return false;
+	}
 }
 
 export class MainThreadNotebookController implements IMainNotebookController {
@@ -186,6 +205,10 @@ export class MainThreadNotebookController implements IMainNotebookController {
 		this._mainThreadNotebook.executeNotebook(viewType, uri);
 	}
 
+	onDidReceiveMessage(uri: UriComponents, message: any): void {
+		this._proxy.$onDidReceiveMessage(uri, message);
+	}
+
 	// Methods for ExtHost
 	async createNotebookDocument(handle: number, viewType: string, resource: UriComponents): Promise<void> {
 		let document = new MainThreadNotebookDocument(this._proxy, handle, viewType, URI.revive(resource));

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

@@ -688,6 +688,7 @@ export interface MainThreadNotebookShape extends IDisposable {
 	$updateNotebookLanguages(viewType: string, resource: UriComponents, languages: string[]): Promise<void>;
 	$spliceNotebookCells(viewType: string, resource: UriComponents, splices: NotebookCellsSplice[], renderers: number[]): Promise<void>;
 	$spliceNotebookCellOutputs(viewType: string, resource: UriComponents, cellHandle: number, splices: NotebookCellOutputsSplice[], renderers: number[]): Promise<void>;
+	$postMessage(handle: number, value: any): Promise<boolean>;
 }
 
 export interface MainThreadUrlsShape extends IDisposable {
@@ -1527,6 +1528,7 @@ export interface ExtHostNotebookShape {
 	$updateActiveEditor(viewType: string, uri: UriComponents): Promise<void>;
 	$destoryNotebookDocument(viewType: string, uri: UriComponents): Promise<boolean>;
 	$acceptDisplayOrder(displayOrder: INotebookDisplayOrder): void;
+	$onDidReceiveMessage(uri: UriComponents, message: any): void;
 }
 
 export interface ExtHostStorageShape {

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

@@ -331,11 +331,14 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 export class ExtHostNotebookEditor extends Disposable implements vscode.NotebookEditor {
 	private _viewColumn: vscode.ViewColumn | undefined;
 	private static _cellhandlePool: number = 0;
+	onDidReceiveMessage: vscode.Event<any> = this._onDidReceiveMessage.event;
 
 	constructor(
 		viewType: string,
 		readonly id: string,
 		public uri: URI,
+		private _proxy: MainThreadNotebookShape,
+		private _onDidReceiveMessage: Emitter<any>,
 		public document: ExtHostNotebookDocument,
 		private _documentsAndEditors: ExtHostDocumentsAndEditors
 	) {
@@ -377,6 +380,11 @@ export class ExtHostNotebookEditor extends Disposable implements vscode.Notebook
 	set viewColumn(value) {
 		throw readonly('viewColumn');
 	}
+
+	async postMessage(message: any): Promise<boolean> {
+		return this._proxy.$postMessage(this.document.handle, message);
+	}
+
 }
 
 export class ExtHostNotebookOutputRenderer {
@@ -418,7 +426,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 	private readonly _proxy: MainThreadNotebookShape;
 	private readonly _notebookProviders = new Map<string, { readonly provider: vscode.NotebookProvider, readonly extension: IExtensionDescription; }>();
 	private readonly _documents = new Map<string, ExtHostNotebookDocument>();
-	private readonly _editors = new Map<string, ExtHostNotebookEditor>();
+	private readonly _editors = new Map<string, { editor: ExtHostNotebookEditor, onDidReceiveMessage: Emitter<any> }>();
 	private readonly _notebookOutputRenderers = new Map<number, ExtHostNotebookOutputRenderer>();
 	private _outputDisplayOrder: INotebookDisplayOrder | undefined;
 
@@ -442,8 +450,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 					const cellHandle = arg.cell.handle;
 
 					for (let value of this._editors) {
-						if (value[1].document.handle === documentHandle) {
-							const cell = value[1].document.getCell(cellHandle);
+						if (value[1].editor.document.handle === documentHandle) {
+							const cell = value[1].editor.document.getCell(cellHandle);
 							if (cell) {
 								return cell;
 							}
@@ -515,15 +523,19 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 				this._documents.set(URI.revive(uri).toString(), document);
 			}
 
+			const onDidReceiveMessage = new Emitter<any>();
+
 			let editor = new ExtHostNotebookEditor(
 				viewType,
 				`${ExtHostNotebookController._handlePool++}`,
 				URI.revive(uri),
+				this._proxy,
+				onDidReceiveMessage,
 				this._documents.get(URI.revive(uri).toString())!,
 				this._documentsAndEditors
 			);
 
-			this._editors.set(URI.revive(uri).toString(), editor);
+			this._editors.set(URI.revive(uri).toString(), { editor, onDidReceiveMessage });
 			await provider.provider.resolveNotebook(editor);
 			// await editor.document.$updateCells();
 			return editor.document.handle;
@@ -556,7 +568,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 			let editor = this._editors.get(URI.revive(uri).toString());
 			let document = this._documents.get(URI.revive(uri).toString());
 
-			let rawCell = editor?.createCell('', language, type, []) as ExtHostCell;
+			let rawCell = editor?.editor.createCell('', language, type, []) as ExtHostCell;
 			document?.insertCell(index, rawCell!);
 
 			let allDocuments = this._documentsAndEditors.allDocuments();
@@ -629,7 +641,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		let editor = this._editors.get(URI.revive(uri).toString());
 
 		if (editor) {
-			editor.dispose();
+			editor.editor.dispose();
+			editor.onDidReceiveMessage.dispose();
 			this._editors.delete(URI.revive(uri).toString());
 		}
 
@@ -639,4 +652,12 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 	$acceptDisplayOrder(displayOrder: INotebookDisplayOrder): void {
 		this._outputDisplayOrder = displayOrder;
 	}
+
+	$onDidReceiveMessage(uri: UriComponents, message: any): void {
+		let editor = this._editors.get(URI.revive(uri).toString());
+
+		if (editor) {
+			editor.onDidReceiveMessage.fire(message);
+		}
+	}
 }

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

@@ -44,6 +44,8 @@ export interface INotebookEditor {
 	 */
 	viewModel: NotebookViewModel | undefined;
 
+	isNotebookEditor: boolean;
+
 	/**
 	 * Focus the notebook editor cell list
 	 */
@@ -121,6 +123,11 @@ export interface INotebookEditor {
 	 */
 	removeInset(output: IOutput): void;
 
+	/**
+	 * Send message to the webview for outputs.
+	 */
+	postMessage(message: any): void;
+
 	/**
 	 * Trigger the editor to scroll from scroll event programmatically
 	 */

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

@@ -219,6 +219,11 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 
 		this.control = new NotebookCodeEditors(this.list, this.renderedEditors);
 		this.webview = new BackLayerWebView(this.webviewService, this.notebookService, this, this.environmentSerice);
+		this._register(this.webview.onMessage(message => {
+			if (this.viewModel) {
+				this.notebookService.onDidReceiveMessage(this.viewModel.viewType, this.viewModel.uri, message);
+			}
+		}));
 		this.list.rowsContainer.appendChild(this.webview.element);
 		this._register(this.list);
 	}
@@ -697,6 +702,10 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		return this.outputRenderer;
 	}
 
+	postMessage(message: any) {
+		this.webview?.webview.sendMessage(message);
+	}
+
 	//#endregion
 
 	toJSON(): any {

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

@@ -30,6 +30,7 @@ export interface IMainNotebookController {
 	createRawCell(uri: URI, index: number, language: string, type: CellKind): Promise<NotebookCellTextModel | undefined>;
 	deleteCell(uri: URI, index: number): Promise<boolean>
 	executeNotebookActiveCell(uri: URI): void;
+	onDidReceiveMessage(uri: URI, message: any): void;
 	destoryNotebookDocument(notebook: INotebookTextModel): Promise<void>;
 	save(uri: URI): Promise<boolean>;
 }
@@ -54,6 +55,7 @@ export interface INotebookService {
 	destoryNotebookDocument(viewType: string, notebook: INotebookTextModel): void;
 	updateActiveNotebookDocument(viewType: string, resource: URI): void;
 	save(viewType: string, resource: URI): Promise<boolean>;
+	onDidReceiveMessage(viewType: string, uri: URI, message: any): void;
 }
 
 export class NotebookProviderInfoStore {
@@ -325,6 +327,14 @@ export class NotebookService extends Disposable implements INotebookService {
 		return false;
 	}
 
+	onDidReceiveMessage(viewType: string, uri: URI, message: any): void {
+		let provider = this._notebookProviders.get(viewType);
+
+		if (provider) {
+			return provider.controller.onDidReceiveMessage(uri, message);
+		}
+	}
+
 	private _onWillDispose(model: INotebookTextModel): void {
 		let modelId = MODEL_ID(model.uri);
 		let modelData = this._models[modelId];

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -16,15 +16,18 @@ import { IWebviewService, WebviewElement } from 'vs/workbench/contrib/webview/br
 import { WebviewResourceScheme } from 'vs/workbench/contrib/webview/common/resourceLoader';
 import { CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookCellViewModel';
 import { CELL_MARGIN } from 'vs/workbench/contrib/notebook/browser/constants';
+import { Emitter, Event } from 'vs/base/common/event';
 
 export interface IDimentionMessage {
+	__vscode_notebook_message: boolean;
 	type: 'dimension';
 	id: string;
 	data: DOM.Dimension;
 }
 
 
 export interface IScrollAckMessage {
+	__vscode_notebook_message: boolean;
 	type: 'scroll-ack';
 	data: { top: number };
 	version: number;
@@ -78,6 +81,9 @@ export class BackLayerWebView extends Disposable {
 	preloadsCache: Map<string, boolean> = new Map();
 	localResourceRootsCache: URI[] | undefined = undefined;
 	rendererRootsCache: URI[] = [];
+	private readonly _onMessage = this._register(new Emitter<any>());
+	public readonly onMessage: Event<any> = this._onMessage.event;
+
 
 	constructor(public webviewService: IWebviewService, public notebookService: INotebookService, public notebookEditor: INotebookEditor, public environmentSerice: IEnvironmentService) {
 		super();
@@ -154,6 +160,7 @@ export class BackLayerWebView extends Disposable {
 			for (let entry of entries) {
 				if (entry.target.id === id && entry.contentRect) {
 					vscode.postMessage({
+							__vscode_notebook_message: true,
 							type: 'dimension',
 							id: id,
 							data: {
@@ -198,6 +205,7 @@ export class BackLayerWebView extends Disposable {
 					resizeObserve(outputNode, outputId);
 
 					vscode.postMessage({
+						__vscode_notebook_message: true,
 						type: 'dimension',
 						id: outputId,
 						data: {
@@ -255,27 +263,32 @@ export class BackLayerWebView extends Disposable {
 		}));
 
 		this._register(this.webview.onMessage((data: IMessage) => {
-			if (data.type === 'dimension') {
-				let output = this.reversedInsetMapping.get(data.id);
+			if (data.__vscode_notebook_message) {
+				if (data.type === 'dimension') {
+					let output = this.reversedInsetMapping.get(data.id);
 
-				if (!output) {
-					return;
-				}
+					if (!output) {
+						return;
+					}
 
-				let cell = this.insetMapping.get(output)!.cell;
-				let height = data.data.height;
-				let outputHeight = height === 0 ? 0 : height + 16;
+					let cell = this.insetMapping.get(output)!.cell;
+					let height = data.data.height;
+					let outputHeight = height === 0 ? 0 : height + 16;
 
-				if (cell) {
-					let outputIndex = cell.outputs.indexOf(output);
-					cell.updateOutputHeight(outputIndex, outputHeight);
-					this.notebookEditor.layoutNotebookCell(cell, cell.getCellTotalHeight());
+					if (cell) {
+						let outputIndex = cell.outputs.indexOf(output);
+						cell.updateOutputHeight(outputIndex, outputHeight);
+						this.notebookEditor.layoutNotebookCell(cell, cell.getCellTotalHeight());
+					}
+				} else if (data.type === 'scroll-ack') {
+					// const date = new Date();
+					// const top = data.data.top;
+					// console.log('ack top ', top, ' version: ', data.version, ' - ', date.getMinutes() + ':' + date.getSeconds() + ':' + date.getMilliseconds());
 				}
-			} else if (data.type === 'scroll-ack') {
-				// const date = new Date();
-				// const top = data.data.top;
-				// console.log('ack top ', top, ' version: ', data.version, ' - ', date.getMinutes() + ':' + date.getSeconds() + ':' + date.getMilliseconds());
+				return;
 			}
+
+			this._onMessage.fire(data);
 		}));
 	}
 

src/vs/workbench/contrib/notebook/test/testNotebookEditor.ts

@@ -69,6 +69,12 @@ export class TestNotebookEditor implements INotebookEditor {
 	constructor(
 	) { }
 
+	isNotebookEditor = true;
+
+	postMessage(message: any): void {
+		throw new Error('Method not implemented.');
+	}
+
 	setCellSelection(cell: CellViewModel, selection: Range): void {
 		throw new Error('Method not implemented.');
 	}