notebooks: pipe renderer API postmessages to the renderer itself (#100414)

* notebooks: pipe renderer API postmessages to the renderer itself

Previously the postMessage on acquireNotebookRenderer API was just
a proxy to the global vscode postmessage. Now, it's linked to the
renderer and will cause an optional `onDidReceiveMessage` method on
the renderer to be called.

The message still _also_ goes to the global webview message handling
for advanced use cases, but this change allows the webview<->renderer
communication to be more nicely contained
and separate for most use cases.

* wip

* fixup! pr comments

src/vs/vscode.proposed.d.ts

@@ -1607,6 +1607,20 @@ declare module 'vscode' {
 		 */
 		render(document: NotebookDocument, request: NotebookRenderRequest): string;
 
+		/**
+		 * Call before HTML from the renderer is executed, and will be called for
+		 * every editor associated with notebook documents where the renderer
+		 * is or was used.
+		 *
+		 * The communication object will only send and receive messages to the
+		 * render API, retrieved via `acquireNotebookRendererApi`, acquired with
+		 * this specific renderer's ID.
+		 *
+		 * If you need to keep an association between the communication object
+		 * and the document for use in the `render()` method, you can use a WeakMap.
+		 */
+		resolveNotebook?(document: NotebookDocument, communication: NotebookCommunication): void;
+
 		readonly preloads?: Uri[];
 	}
 
@@ -1735,27 +1749,41 @@ declare module 'vscode' {
 		readonly backupId?: string;
 	}
 
+	/**
+	 * Communication object passed to the {@link NotebookContentProvider} and
+	 * {@link NotebookOutputRenderer} to communicate with the webview.
+	 */
+	export interface NotebookCommunication {
+		/**
+		 * ID of 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 editorId: 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 serilizable 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 NotebookContentProvider {
 		openNotebook(uri: Uri, openContext: NotebookDocumentOpenContext): NotebookData | Promise<NotebookData>;
-		resolveNotebook(document: NotebookDocument, webview: {
-			/**
-			 * 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>;
-
-			/**
-			 * Convert a uri for the local file system to one that can be used inside outputs webview.
-			 */
-			asWebviewUri(localResource: Uri): Uri;
-		}): Promise<void>;
+		resolveNotebook(document: NotebookDocument, webview: NotebookCommunication): Promise<void>;
 		saveNotebook(document: NotebookDocument, cancellation: CancellationToken): Promise<void>;
 		saveNotebookAs(targetResource: Uri, document: NotebookDocument, cancellation: CancellationToken): Promise<void>;
 		readonly onDidChangeNotebook: Event<NotebookDocumentContentChangeEvent>;

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

@@ -253,7 +253,7 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 			return;
 		}
 
-		this._proxy.$acceptDocumentAndEditorsDelta(delta);
+		return this._proxy.$acceptDocumentAndEditorsDelta(delta);
 	}
 
 	registerListeners() {
@@ -476,16 +476,11 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		return this._proxy.$executeNotebook(viewType, uri, undefined, useAttachedKernel, token);
 	}
 
-	async $postMessage(handle: number, value: any): Promise<boolean> {
-
-		const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
-		if (activeEditorPane?.isNotebookEditor) {
-			const notebookEditor = (activeEditorPane.getControl() as INotebookEditor);
-
-			if (notebookEditor.viewModel?.handle === handle) {
-				notebookEditor.postMessage(value);
-				return true;
-			}
+	async $postMessage(editorId: string, forRendererId: string | undefined, value: any): Promise<boolean> {
+		const editor = this._notebookService.getNotebookEditor(editorId) as INotebookEditor | undefined;
+		if (editor?.isNotebookEditor) {
+			editor.postMessage(forRendererId, value);
+			return true;
 		}
 
 		return false;
@@ -645,8 +640,8 @@ export class MainThreadNotebookController implements IMainNotebookController {
 		return this._mainThreadNotebook.executeNotebook(viewType, uri, useAttachedKernel, token);
 	}
 
-	onDidReceiveMessage(editorId: string, message: any): void {
-		this._proxy.$onDidReceiveMessage(editorId, message);
+	onDidReceiveMessage(editorId: string, rendererType: string | undefined, message: unknown): void {
+		this._proxy.$onDidReceiveMessage(editorId, rendererType, message);
 	}
 
 	async removeNotebookDocument(notebook: INotebookTextModel): Promise<void> {

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

@@ -711,7 +711,7 @@ export interface MainThreadNotebookShape extends IDisposable {
 	$updateNotebookMetadata(viewType: string, resource: UriComponents, metadata: NotebookDocumentMetadata): Promise<void>;
 	$updateNotebookCellMetadata(viewType: string, resource: UriComponents, handle: number, metadata: NotebookCellMetadata | undefined): Promise<void>;
 	$spliceNotebookCellOutputs(viewType: string, resource: UriComponents, cellHandle: number, splices: NotebookCellOutputsSplice[], renderers: number[]): Promise<void>;
-	$postMessage(handle: number, value: any): Promise<boolean>;
+	$postMessage(editorId: string, forRendererId: string | undefined, value: any): Promise<boolean>;
 
 	$onDidEdit(resource: UriComponents, viewType: string, editId: number, label: string | undefined): void;
 	$onContentChange(resource: UriComponents, viewType: string): void;
@@ -1610,7 +1610,7 @@ export interface ExtHostNotebookShape {
 	$acceptDisplayOrder(displayOrder: INotebookDisplayOrder): void;
 	$renderOutputs(uriComponents: UriComponents, id: string, request: IOutputRenderRequest<UriComponents>): Promise<IOutputRenderResponse<UriComponents> | undefined>;
 	$renderOutputs2<T>(uriComponents: UriComponents, id: string, request: IOutputRenderRequest<T>): Promise<IOutputRenderResponse<T> | undefined>;
-	$onDidReceiveMessage(editorId: string, message: any): void;
+	$onDidReceiveMessage(editorId: string, rendererId: string | undefined, message: unknown): void;
 	$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void;
 	$acceptEditorPropertiesChanged(uriComponents: UriComponents, data: INotebookEditorPropertiesChangeData): void;
 	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): Promise<void>;

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

@@ -494,6 +494,10 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 			cell.detachTextDocument();
 		}
 	}
+
+	allEditors() {
+		return this._documentsAndEditors.allEditors();
+	}
 }
 
 export class NotebookEditorCellEditBuilder implements vscode.NotebookEditorCellEdit {
@@ -553,27 +557,48 @@ export class NotebookEditorCellEditBuilder implements vscode.NotebookEditorCellE
 	}
 }
 
-class ExtHostWebviewComm extends Disposable {
-
-	onDidReceiveMessage: vscode.Event<any> = this._onDidReceiveMessage.event;
+class ExtHostWebviewCommWrapper extends Disposable {
+	private readonly _onDidReceiveDocumentMessage = new Emitter<any>();
+	private readonly _rendererIdToEmitters = new Map<string, Emitter<any>>();
 
 	constructor(
-		readonly id: string,
+		private _editorId: string,
 		public uri: URI,
 		private _proxy: MainThreadNotebookShape,
-		private _onDidReceiveMessage: Emitter<any>,
 		private _webviewInitData: WebviewInitData,
 		public document: ExtHostNotebookDocument,
 	) {
 		super();
 	}
 
-	async postMessage(message: any): Promise<boolean> {
-		return this._proxy.$postMessage(this.document.handle, message);
+	public onDidReceiveMessage(forRendererId: string | undefined, message: any) {
+		this._onDidReceiveDocumentMessage.fire(message);
+		if (forRendererId !== undefined) {
+			this._rendererIdToEmitters.get(forRendererId)?.fire(message);
+		}
 	}
 
-	asWebviewUri(localResource: vscode.Uri): vscode.Uri {
-		return asWebviewUri(this._webviewInitData, this.id, localResource);
+	public readonly contentProviderComm: vscode.NotebookCommunication = {
+		editorId: this._editorId,
+		onDidReceiveMessage: this._onDidReceiveDocumentMessage.event,
+		postMessage: (message: any) => this._proxy.$postMessage(this._editorId, undefined, message),
+		asWebviewUri: (uri: vscode.Uri) => this._asWebviewUri(uri),
+	};
+
+	public getRendererComm(rendererId: string): vscode.NotebookCommunication {
+		const emitter = new Emitter<any>();
+		this._rendererIdToEmitters.set(rendererId, emitter);
+		return {
+			editorId: this._editorId,
+			onDidReceiveMessage: emitter.event,
+			postMessage: (message: any) => this._proxy.$postMessage(this._editorId, rendererId, message),
+			asWebviewUri: (uri: vscode.Uri) => this._asWebviewUri(uri),
+		};
+	}
+
+
+	private _asWebviewUri(localResource: vscode.Uri): vscode.Uri {
+		return asWebviewUri(this._webviewInitData, this._editorId, localResource);
 	}
 }
 
@@ -618,7 +643,7 @@ export class ExtHostNotebookEditor extends Disposable implements vscode.Notebook
 		readonly id: string,
 		public uri: URI,
 		private _proxy: MainThreadNotebookShape,
-		private _webComm: ExtHostWebviewComm,
+		private _webComm: vscode.NotebookCommunication,
 		public document: ExtHostNotebookDocument,
 		private _documentsAndEditors: ExtHostDocumentsAndEditors
 	) {
@@ -721,6 +746,7 @@ export class ExtHostNotebookEditor extends Disposable implements vscode.Notebook
 
 export class ExtHostNotebookOutputRenderer {
 	private static _handlePool: number = 0;
+	private resolvedComms = new WeakSet<ExtHostWebviewCommWrapper>();
 	readonly handle = ExtHostNotebookOutputRenderer._handlePool++;
 
 	constructor(
@@ -740,13 +766,19 @@ export class ExtHostNotebookOutputRenderer {
 		return false;
 	}
 
+	resolveNotebook(document: ExtHostNotebookDocument, comm: ExtHostWebviewCommWrapper) {
+		if (!this.resolvedComms.has(comm) && this.renderer.resolveNotebook) {
+			this.renderer.resolveNotebook(document, comm.getRendererComm(this.type));
+			this.resolvedComms.add(comm);
+		}
+	}
+
 	render(document: ExtHostNotebookDocument, output: vscode.CellDisplayOutput, outputId: string, mimeType: string): string {
 		let html = this.renderer.render(document, { output, outputId, mimeType });
 
 		return html;
 	}
 }
-
 export interface ExtHostNotebookOutputRenderingHandler {
 	outputDisplayOrder: INotebookDisplayOrder | undefined;
 	findBestMatchedRenderer(mimeType: string): ExtHostNotebookOutputRenderer[];
@@ -759,8 +791,9 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 	private readonly _documents = new Map<string, ExtHostNotebookDocument>();
 	private readonly _unInitializedDocuments = new Map<string, ExtHostNotebookDocument>();
 	private readonly _editors = new Map<string, { editor: ExtHostNotebookEditor }>();
-	private readonly _webviewComm = new Map<string, { comm: ExtHostWebviewComm, onDidReceiveMessage: Emitter<any> }>();
+	private readonly _webviewComm = new Map<string, ExtHostWebviewCommWrapper>();
 	private readonly _notebookOutputRenderers = new Map<string, ExtHostNotebookOutputRenderer>();
+	private readonly _renderersUsedInNotebooks = new WeakMap<ExtHostNotebookDocument, Set<ExtHostNotebookOutputRenderer>>();
 	private readonly _onDidChangeNotebookCells = new Emitter<vscode.NotebookCellsChangeEvent>();
 	readonly onDidChangeNotebookCells = this._onDidChangeNotebookCells.event;
 	private readonly _onDidChangeCellOutputs = new Emitter<vscode.NotebookCellOutputsChangeEvent>();
@@ -854,6 +887,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		}
 
 		const renderer = this._notebookOutputRenderers.get(id)!;
+		this.provideCommToNotebookRenderers(document, renderer);
+
 		const cellsResponse: IOutputRenderResponseCellInfo<UriComponents>[] = request.items.map(cellInfo => {
 			const cell = document.getCell2(cellInfo.key)!;
 			const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
@@ -890,6 +925,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		}
 
 		const renderer = this._notebookOutputRenderers.get(id)!;
+		this.provideCommToNotebookRenderers(document, renderer);
+
 		const cellsResponse: IOutputRenderResponseCellInfo<T>[] = request.items.map(cellInfo => {
 			const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
 				return {
@@ -1035,19 +1072,36 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 			return;
 		}
 
+		let webComm = this._webviewComm.get(editorId);
+		if (!webComm) {
+			webComm = new ExtHostWebviewCommWrapper(editorId, revivedUri, this._proxy, this._webviewInitData, document);
+			this._webviewComm.set(editorId, webComm);
+		}
+
 		if (!provider.provider.resolveNotebook) {
 			return;
 		}
 
-		let webComm = this._webviewComm.get(editorId)?.comm;
+		await provider.provider.resolveNotebook(document, webComm.contentProviderComm);
+	}
 
-		if (webComm) {
-			await provider.provider.resolveNotebook(document, webComm);
-		} else {
-			const onDidReceiveMessage = new Emitter<any>();
-			webComm = new ExtHostWebviewComm(editorId, revivedUri, this._proxy, onDidReceiveMessage, this._webviewInitData, document);
-			this._webviewComm.set(editorId, { comm: webComm, onDidReceiveMessage });
-			await provider.provider.resolveNotebook(document, webComm);
+	private provideCommToNotebookRenderers(document: ExtHostNotebookDocument, renderer: ExtHostNotebookOutputRenderer) {
+		let alreadyRegistered = this._renderersUsedInNotebooks.get(document);
+		if (!alreadyRegistered) {
+			alreadyRegistered = new Set();
+			this._renderersUsedInNotebooks.set(document, alreadyRegistered);
+		}
+
+		if (alreadyRegistered.has(renderer)) {
+			return;
+		}
+
+		alreadyRegistered.add(renderer);
+		for (const editor of document.allEditors()) {
+			const comm = this._webviewComm.get(editor.id);
+			if (comm) {
+				renderer.resolveNotebook(document, comm);
+			}
 		}
 	}
 
@@ -1182,12 +1236,8 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 		return editor;
 	}
 
-	$onDidReceiveMessage(editorId: string, message: any): void {
-		let messageEmitter = this._webviewComm.get(editorId)?.onDidReceiveMessage;
-
-		if (messageEmitter) {
-			messageEmitter.fire(message);
-		}
+	$onDidReceiveMessage(editorId: string, forRendererType: string | undefined, message: any): void {
+		this._webviewComm.get(editorId)?.onDidReceiveMessage(forRendererType, message);
 	}
 
 	$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void {
@@ -1226,12 +1276,11 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 
 	private _createExtHostEditor(document: ExtHostNotebookDocument, editorId: string, selections: number[]) {
 		const revivedUri = document.uri;
-		let webComm = this._webviewComm.get(editorId)?.comm;
+		let webComm = this._webviewComm.get(editorId);
 
 		if (!webComm) {
-			const onDidReceiveMessage = new Emitter<any>();
-			webComm = new ExtHostWebviewComm(editorId, revivedUri, this._proxy, onDidReceiveMessage, this._webviewInitData, document);
-			this._webviewComm.set(editorId, { comm: webComm!, onDidReceiveMessage });
+			webComm = new ExtHostWebviewCommWrapper(editorId, revivedUri, this._proxy, this._webviewInitData, document);
+			this._webviewComm.set(editorId, webComm);
 		}
 
 		let editor = new ExtHostNotebookEditor(
@@ -1239,7 +1288,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 			editorId,
 			revivedUri,
 			this._proxy,
-			webComm,
+			webComm.contentProviderComm,
 			document,
 			this._documentsAndEditors
 		);
@@ -1255,6 +1304,10 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 
 		this._editors.get(editorId)?.editor.dispose();
 
+		for (const renderer of this._renderersUsedInNotebooks.get(document) ?? []) {
+			renderer.resolveNotebook(document, webComm);
+		}
+
 		this._editors.set(editorId, { editor });
 	}
 

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

@@ -270,7 +270,7 @@ export interface INotebookEditor extends IEditor {
 	/**
 	 * Send message to the webview for outputs.
 	 */
-	postMessage(message: any): void;
+	postMessage(forRendererId: string | undefined, message: any): void;
 
 	/**
 	 * Trigger the editor to scroll from scroll event programmatically

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

@@ -528,9 +528,9 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 			this._onDidFocusEmitter.fire();
 		});
 
-		this._localStore.add(this._webview.onMessage(message => {
+		this._localStore.add(this._webview.onMessage(({ message, forRenderer }) => {
 			if (this.viewModel) {
-				this.notebookService.onDidReceiveMessage(this.viewModel.viewType, this.getId(), message);
+				this.notebookService.onDidReceiveMessage(this.viewModel.viewType, this.getId(), forRenderer, message);
 			}
 		}));
 	}
@@ -1213,8 +1213,12 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		return this._outputRenderer;
 	}
 
-	postMessage(message: any) {
-		this._webview?.webview.postMessage(message);
+	postMessage(forRendererId: string | undefined, message: any) {
+		if (forRendererId === undefined) {
+			this._webview?.webview.postMessage(message);
+		} else {
+			this._webview?.postRendererMessage(forRendererId, message);
+		}
 	}
 
 	//#endregion

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

@@ -683,6 +683,10 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		this._onNotebookEditorAdd.fire(editor);
 	}
 
+	getNotebookEditor(editorId: string) {
+		return this._notebookEditors.get(editorId);
+	}
+
 	listNotebookEditors(): INotebookEditor[] {
 		return [...this._notebookEditors].map(e => e[1]);
 	}
@@ -750,11 +754,11 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		return;
 	}
 
-	onDidReceiveMessage(viewType: string, editorId: string, message: any): void {
+	onDidReceiveMessage(viewType: string, editorId: string, rendererType: string | undefined, message: any): void {
 		let provider = this._notebookProviders.get(viewType);
 
 		if (provider) {
-			return provider.controller.onDidReceiveMessage(editorId, message);
+			return provider.controller.onDidReceiveMessage(editorId, rendererType, message);
 		}
 	}
 

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

@@ -76,7 +76,7 @@ export interface IBlurOutputMessage {
 }
 
 export interface IClickedDataUrlMessage {
-	__vscode_notebook_message: string;
+	__vscode_notebook_message: boolean;
 	type: 'clicked-data-url';
 	data: string;
 	downloadName?: string;
@@ -156,6 +156,13 @@ export interface IUpdatePreloadResourceMessage {
 	source: 'renderer' | 'kernel';
 }
 
+export interface ICustomRendererMessage {
+	__vscode_notebook_message: boolean;
+	type: 'customRendererMessage';
+	rendererId: string;
+	message: unknown;
+}
+
 export type FromWebviewMessage =
 	| WebviewIntialized
 	| IDimensionMessage
@@ -163,7 +170,9 @@ export type FromWebviewMessage =
 	| IMouseLeaveMessage
 	| IWheelMessage
 	| IScrollAckMessage
-	| IBlurOutputMessage;
+	| IBlurOutputMessage
+	| ICustomRendererMessage
+	| IClickedDataUrlMessage;
 
 export type ToWebviewMessage =
 	| IClearMessage
@@ -175,7 +184,8 @@ export type ToWebviewMessage =
 	| IHideOutputMessage
 	| IShowOutputMessage
 	| IUpdatePreloadResourceMessage
-	| IFocusOutputMessage;
+	| IFocusOutputMessage
+	| ICustomRendererMessage;
 
 export type AnyMessage = FromWebviewMessage | ToWebviewMessage;
 
@@ -194,7 +204,10 @@ function html(strings: TemplateStringsArray, ...values: any[]): string {
 	return str;
 }
 
-type IMessage = IDimensionMessage | IScrollAckMessage | IWheelMessage | IMouseEnterMessage | IMouseLeaveMessage | IBlurOutputMessage | WebviewIntialized | IClickedDataUrlMessage;
+export interface INotebookWebviewMessage {
+	message: unknown;
+	forRenderer?: string;
+}
 
 let version = 0;
 export class BackLayerWebView extends Disposable {
@@ -207,8 +220,8 @@ export class BackLayerWebView extends Disposable {
 	localResourceRootsCache: URI[] | undefined = undefined;
 	rendererRootsCache: URI[] = [];
 	kernelRootsCache: URI[] = [];
-	private readonly _onMessage = this._register(new Emitter<any>());
-	public readonly onMessage: Event<any> = this._onMessage.event;
+	private readonly _onMessage = this._register(new Emitter<INotebookWebviewMessage>());
+	public readonly onMessage: Event<INotebookWebviewMessage> = this._onMessage.event;
 	private _loaded!: Promise<void>;
 	private _initalized?: Promise<void>;
 	private _disposed = false;
@@ -267,6 +280,15 @@ export class BackLayerWebView extends Disposable {
 		</html>`;
 	}
 
+	postRendererMessage(rendererId: string, message: any) {
+		this._sendMessageToWebview({
+			__vscode_notebook_message: true,
+			type: 'customRendererMessage',
+			message,
+			rendererId
+		});
+	}
+
 	private resolveOutputId(id: string): { cell: CodeCellViewModel, output: IProcessedOutput } | undefined {
 		const output = this.reversedInsetMapping.get(id);
 		if (!output) {
@@ -341,7 +363,7 @@ ${loaderJs}
 			}
 		}));
 
-		this._register(this.webview.onMessage((data: IMessage) => {
+		this._register(this.webview.onMessage((data: FromWebviewMessage) => {
 			if (data.__vscode_notebook_message) {
 				if (data.type === 'dimension') {
 					let height = data.data.height;
@@ -397,11 +419,13 @@ ${loaderJs}
 					}
 				} else if (data.type === 'clicked-data-url') {
 					this._onDidClickDataLink(data);
+				} else if (data.type === 'customRendererMessage') {
+					this._onMessage.fire({ message: data.message, forRenderer: data.rendererId });
 				}
 				return;
 			}
 
-			this._onMessage.fire(data);
+			this._onMessage.fire({ message: data });
 		}));
 	}
 
@@ -459,7 +483,7 @@ ${loaderJs}
 			resolveFunc = resolve;
 		});
 
-		let dispose = webview.onMessage((data: IMessage) => {
+		let dispose = webview.onMessage((data: FromWebviewMessage) => {
 			if (data.__vscode_notebook_message && data.type === 'initialized') {
 				resolveFunc();
 				dispose.dispose();

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

@@ -262,6 +262,7 @@ function webviewPreloads() {
 
 	const onWillDestroyOutput = createEmitter<[string | undefined /* namespace */, IDestroyCellInfo | undefined /* cell uri */]>();
 	const onDidCreateOutput = createEmitter<[string | undefined /* namespace */, ICreateCellInfo]>();
+	const onDidReceiveMessage = createEmitter<[string, unknown]>();
 
 	const matchesNs = (namespace: string, query: string | undefined) => namespace === '*' || query === namespace || query === 'undefined';
 
@@ -271,7 +272,14 @@ function webviewPreloads() {
 		}
 
 		return {
-			postMessage: vscode.postMessage,
+			postMessage(message: unknown) {
+				vscode.postMessage({
+					__vscode_notebook_message: true,
+					type: 'customRendererMessage',
+					rendererId: namespace,
+					message,
+				});
+			},
 			setState(newState: T) {
 				vscode.setState({ ...vscode.getState(), [namespace]: newState });
 			},
@@ -279,6 +287,7 @@ function webviewPreloads() {
 				const state = vscode.getState();
 				return typeof state === 'object' && state ? state[namespace] as T : undefined;
 			},
+			onDidReceiveMessage: mapEmitter(onDidReceiveMessage, ([ns, data]) => ns === namespace ? data : dontEmit),
 			onWillDestroyOutput: mapEmitter(onWillDestroyOutput, ([ns, data]) => matchesNs(namespace, ns) ? data : dontEmit),
 			onDidCreateOutput: mapEmitter(onDidCreateOutput, ([ns, data]) => matchesNs(namespace, ns) ? data : dontEmit),
 		};
@@ -406,6 +415,9 @@ function webviewPreloads() {
 					focusFirstFocusableInCell(event.data.id);
 					break;
 				}
+			case 'customRendererMessage':
+				onDidReceiveMessage.fire([event.data.rendererId, event.data.message]);
+				break;
 		}
 	});
 

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

@@ -21,7 +21,7 @@ export interface IMainNotebookController {
 	createNotebook(viewType: string, uri: URI, backup: INotebookTextModelBackup | undefined, forceReload: boolean, editorId?: string, backupId?: string): Promise<NotebookTextModel | undefined>;
 	resolveNotebookEditor(viewType: string, uri: URI, editorId: string): Promise<void>;
 	executeNotebook(viewType: string, uri: URI, useAttachedKernel: boolean, token: CancellationToken): Promise<void>;
-	onDidReceiveMessage(editorId: string, message: any): void;
+	onDidReceiveMessage(editorId: string, rendererType: string | undefined, message: any): void;
 	executeNotebookCell(uri: URI, handle: number, useAttachedKernel: boolean, token: CancellationToken): Promise<void>;
 	removeNotebookDocument(notebook: INotebookTextModel): Promise<void>;
 	save(uri: URI, token: CancellationToken): Promise<boolean>;
@@ -65,13 +65,14 @@ export interface INotebookService {
 	save(viewType: string, resource: URI, token: CancellationToken): Promise<boolean>;
 	saveAs(viewType: string, resource: URI, target: URI, token: CancellationToken): Promise<boolean>;
 	backup(viewType: string, uri: URI, token: CancellationToken): Promise<string | undefined>;
-	onDidReceiveMessage(viewType: string, editorId: string, message: any): void;
+	onDidReceiveMessage(viewType: string, editorId: string, rendererType: string | undefined, message: unknown): void;
 	setToCopy(items: NotebookCellTextModel[]): void;
 	getToCopy(): NotebookCellTextModel[] | undefined;
 
 	// editor events
 	addNotebookEditor(editor: IEditor): void;
 	removeNotebookEditor(editor: IEditor): void;
+	getNotebookEditor(editorId: string): IEditor | undefined;
 	listNotebookEditors(): readonly IEditor[];
 	listVisibleNotebookEditors(): readonly IEditor[];
 	listNotebookDocuments(): readonly NotebookTextModel[];

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

@@ -110,7 +110,7 @@ export class TestNotebookEditor implements INotebookEditor {
 
 	isNotebookEditor = true;
 
-	postMessage(message: any): void {
+	postMessage(): void {
 		throw new Error('Method not implemented.');
 	}