notebooks: add ID for rich outputs, expose to renderer and webview

Some API changes:

 - I changed the second param of NotebookOutputRenderer.render to a
   `NotebookRenderRequest` since we were up to four parameters already.
 - The request contains a unique `outputId`, which is generated when
   output is first created or restored (the ID is not stable if a
	 notebook is closed and reopened)
 - I changed the renderer events to `onDidCreateOutput` and
   `onWillDestroyOutput` since they're for outputs, not cells.
 - Additionally, they're called with objects that contain the output
   ID to render.

src/vs/vscode.proposed.d.ts

@@ -1526,14 +1526,21 @@ declare module 'vscode' {
 		subTypes?: string[];
 	}
 
+	export interface NotebookRenderRequest {
+		output: CellDisplayOutput;
+		mimeType: string;
+		outputId: string;
+	}
+
 	export interface NotebookOutputRenderer {
 		/**
 		 *
 		 * @returns HTML fragment. We can probably return `CellOutput` instead of string ?
 		 *
 		 */
-		render(document: NotebookDocument, output: CellDisplayOutput, mimeType: string): string;
-		preloads?: Uri[];
+		render(document: NotebookDocument, request: NotebookRenderRequest): string;
+
+		readonly preloads?: Uri[];
 	}
 
 	export interface NotebookCellsChangeData {

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

@@ -51,7 +51,7 @@ import { TunnelDto } from 'vs/workbench/api/common/extHostTunnelService';
 import { TunnelOptions } from 'vs/platform/remote/common/tunnel';
 import { Timeline, TimelineChangeEvent, TimelineOptions, TimelineProviderDescriptor, InternalTimelineOptions } from 'vs/workbench/contrib/timeline/common/timeline';
 import { revive } from 'vs/base/common/marshalling';
-import { INotebookMimeTypeSelector, IProcessedOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto, INotebookKernelInfoDto, IMainCellDto, IOutputRenderRequest, IOutputRenderResponse, IRawOutput } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookMimeTypeSelector, IProcessedOutput, INotebookDisplayOrder, NotebookCellMetadata, NotebookDocumentMetadata, ICellEditOperation, NotebookCellsChangedEvent, NotebookDataDto, INotebookKernelInfoDto, IMainCellDto, IOutputRenderRequest, IOutputRenderResponse } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { CallHierarchyItem } from 'vs/workbench/contrib/callHierarchy/common/callHierarchy';
 import { Dto } from 'vs/base/common/types';
 import { ISerializableEnvironmentVariableCollection } from 'vs/workbench/contrib/terminal/common/environmentVariable';
@@ -693,7 +693,7 @@ export type NotebookCellsSplice = [
 export type NotebookCellOutputsSplice = [
 	number /* start */,
 	number /* delete count */,
-	IRawOutput[]
+	IProcessedOutput[]
 ];
 
 export interface MainThreadNotebookShape extends IDisposable {

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

@@ -13,13 +13,14 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'
 import { CellKind, ExtHostNotebookShape, IMainContext, MainContext, MainThreadNotebookShape, NotebookCellOutputsSplice, MainThreadDocumentsShape, INotebookEditorPropertiesChangeData, INotebookDocumentsAndEditorsDelta } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
 import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
-import { CellEditType, CellUri, diff, ICellEditOperation, ICellInsertEdit, INotebookDisplayOrder, INotebookEditData, NotebookCellsChangedEvent, NotebookCellsSplice2, ICellDeleteEdit, notebookDocumentMetadataDefaults, NotebookCellsChangeType, NotebookDataDto, IOutputRenderRequest, IOutputRenderResponse, IOutputRenderResponseOutputInfo, IOutputRenderResponseCellInfo } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellEditType, CellUri, diff, ICellEditOperation, ICellInsertEdit, INotebookDisplayOrder, INotebookEditData, NotebookCellsChangedEvent, NotebookCellsSplice2, ICellDeleteEdit, notebookDocumentMetadataDefaults, NotebookCellsChangeType, NotebookDataDto, IOutputRenderRequest, IOutputRenderResponse, IOutputRenderResponseOutputInfo, IOutputRenderResponseCellInfo, IRawOutput, CellOutputKind, IProcessedOutput } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { ExtHostDocumentData } from 'vs/workbench/api/common/extHostDocumentData';
 import { NotImplementedProxy } from 'vs/base/common/types';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import { asWebviewUri, WebviewInitData } from 'vs/workbench/api/common/shared/webview';
+import { generateUuid } from 'vs/base/common/uuid';
 
 interface IObservable<T> {
 	proxy: T;
@@ -48,15 +49,18 @@ interface INotebookEventEmitter {
 	emitCellLanguageChange(event: vscode.NotebookCellLanguageChangeEvent): void;
 }
 
+const addIdToOutput = (output: IRawOutput, id = generateUuid()): IProcessedOutput => output.outputKind === CellOutputKind.Rich
+	? ({ ...output, outputId: id }) : output;
+
 export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 
 	// private originalSource: string[];
 	private _outputs: any[];
-	private _onDidChangeOutputs = new Emitter<ISplice<vscode.CellOutput>[]>();
-	onDidChangeOutputs: Event<ISplice<vscode.CellOutput>[]> = this._onDidChangeOutputs.event;
+	private _onDidChangeOutputs = new Emitter<ISplice<IProcessedOutput>[]>();
+	onDidChangeOutputs: Event<ISplice<IProcessedOutput>[]> = this._onDidChangeOutputs.event;
 	// private _textDocument: vscode.TextDocument | undefined;
 	// private _initalVersion: number = -1;
-	private _outputMapping = new Set<vscode.CellOutput>();
+	private _outputMapping = new WeakMap<vscode.CellOutput, string | undefined /* output ID */>();
 	private _metadata: vscode.NotebookCellMetadata;
 
 	private _metadataChangeListener: IDisposable;
@@ -91,6 +95,10 @@ export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 		);
 
 		this._outputs = outputs;
+		for (const output of this._outputs) {
+			this._outputMapping.set(output, output.outputId);
+			delete output.outputId;
+		}
 
 		const observableMetadata = getObservable(_metadata || {});
 		this._metadata = observableMetadata.proxy;
@@ -104,22 +112,33 @@ export class ExtHostCell extends Disposable implements vscode.NotebookCell {
 	}
 
 	set outputs(newOutputs: vscode.CellOutput[]) {
-		let diffs = diff<vscode.CellOutput>(this._outputs || [], newOutputs || [], (a) => {
+		let rawDiffs = diff<vscode.CellOutput>(this._outputs || [], newOutputs || [], (a) => {
 			return this._outputMapping.has(a);
 		});
 
-		diffs.forEach(diff => {
+		const transformedDiffs: ISplice<IProcessedOutput>[] = rawDiffs.map(diff => {
 			for (let i = diff.start; i < diff.start + diff.deleteCount; i++) {
 				this._outputMapping.delete(this._outputs[i]);
 			}
 
-			diff.toInsert.forEach(output => {
-				this._outputMapping.add(output);
-			});
+			return {
+				deleteCount: diff.deleteCount,
+				start: diff.start,
+				toInsert: diff.toInsert.map((output): IProcessedOutput => {
+					if (output.outputKind === CellOutputKind.Rich) {
+						const uuid = generateUuid();
+						this._outputMapping.set(output, uuid);
+						return { ...output, outputId: uuid };
+					}
+
+					this._outputMapping.set(output, undefined);
+					return output;
+				})
+			};
 		});
 
 		this._outputs = newOutputs;
-		this._onDidChangeOutputs.fire(diffs);
+		this._onDidChangeOutputs.fire(transformedDiffs);
 	}
 
 	get metadata() {
@@ -375,7 +394,7 @@ export class ExtHostNotebookDocument extends Disposable implements vscode.Notebo
 		this._emitter.emitCellLanguageChange(event);
 	}
 
-	async eventuallyUpdateCellOutputs(cell: ExtHostCell, diffs: ISplice<vscode.CellOutput>[]) {
+	async eventuallyUpdateCellOutputs(cell: ExtHostCell, diffs: ISplice<IProcessedOutput>[]) {
 		let renderers = new Set<number>();
 		let outputDtos: NotebookCellOutputsSplice[] = diffs.map(diff => {
 			let outputs = diff.toInsert;
@@ -447,8 +466,8 @@ export class NotebookEditorCellEditBuilder implements vscode.NotebookEditorCellE
 			source: sourceArr,
 			language,
 			cellKind: type,
-			outputs: outputs,
-			metadata
+			outputs: outputs.map(o => addIdToOutput(o)),
+			metadata,
 		};
 
 		this._collectedEdits.push({
@@ -628,8 +647,8 @@ export class ExtHostNotebookOutputRenderer {
 		return false;
 	}
 
-	render(document: ExtHostNotebookDocument, output: vscode.CellDisplayOutput, mimeType: string): string {
-		let html = this.renderer.render(document, output, mimeType);
+	render(document: ExtHostNotebookDocument, output: vscode.CellDisplayOutput, outputId: string, mimeType: string): string {
+		let html = this.renderer.render(document, { output, outputId, mimeType });
 
 		return html;
 	}
@@ -736,13 +755,14 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 
 		const renderer = this._notebookOutputRenderers.get(id)!;
 		const cellsResponse: IOutputRenderResponseCellInfo<UriComponents>[] = request.items.map(cellInfo => {
-			const cell = document.getCell2(cellInfo.key);
+			const cell = document.getCell2(cellInfo.key)!;
 			const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
 				return {
 					index: output.index,
+					outputId: output.outputId,
 					mimeType: output.mimeType,
 					handlerId: id,
-					transformedOutput: renderer.render(document, cell!.outputs[output.index] as vscode.CellDisplayOutput, output.mimeType)
+					transformedOutput: renderer.render(document, cell.outputs[output.index] as vscode.CellDisplayOutput, output.outputId, output.mimeType)
 				};
 			});
 
@@ -774,9 +794,10 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 			const outputResponse: IOutputRenderResponseOutputInfo[] = cellInfo.outputs.map(output => {
 				return {
 					index: output.index,
+					outputId: output.outputId,
 					mimeType: output.mimeType,
 					handlerId: id,
-					transformedOutput: renderer.render(document, output.output! as vscode.CellDisplayOutput, output.mimeType)
+					transformedOutput: renderer.render(document, output.output! as vscode.CellDisplayOutput, output.outputId, output.mimeType)
 				};
 			});
 
@@ -873,7 +894,10 @@ export class ExtHostNotebookController implements ExtHostNotebookShape, ExtHostN
 					...rawCells.metadata
 				},
 				languages: rawCells.languages,
-				cells: rawCells.cells,
+				cells: rawCells.cells.map(cell => ({
+					...cell,
+					outputs: cell.outputs.map(o => addIdToOutput(o))
+				})),
 			};
 
 			return dto;

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

@@ -29,6 +29,7 @@ import { IAccessibilityService } from 'vs/platform/accessibility/common/accessib
 import { Memento } from 'vs/workbench/common/memento';
 import { StorageScope, IStorageService } from 'vs/platform/storage/common/storage';
 import { IExtensionPointUser } from 'vs/workbench/services/extensions/common/extensionsRegistry';
+import { generateUuid } from 'vs/base/common/uuid';
 
 function MODEL_ID(resource: URI): string {
 	return resource.toString();
@@ -435,14 +436,14 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 			outputs.forEach((output, index) => {
 				if (output.outputKind === CellOutputKind.Rich) {
 					// TODO no string[] casting
-					const ret = this._transformMimeTypes(output, textModel.metadata.displayOrder as string[] || []);
+					const ret = this._transformMimeTypes(output, output.outputId, textModel.metadata.displayOrder as string[] || []);
 					const orderedMimeTypes = ret.orderedMimeTypes!;
 					const pickedMimeTypeIndex = ret.pickedMimeTypeIndex!;
 					output.pickedMimeTypeIndex = pickedMimeTypeIndex;
 					output.orderedMimeTypes = orderedMimeTypes;
 
 					if (orderedMimeTypes[pickedMimeTypeIndex!].rendererId && orderedMimeTypes[pickedMimeTypeIndex].rendererId !== BUILTIN_RENDERER_ID) {
-						outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType });
+						outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType, outputId: output.outputId });
 						renderers.add(orderedMimeTypes[pickedMimeTypeIndex].rendererId!);
 					}
 				}
@@ -469,14 +470,14 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 					const outputRequest: IOutputRenderRequestOutputInfo[] = [];
 					outputs.map((output, index) => {
 						if (output.outputKind === CellOutputKind.Rich) {
-							const ret = this._transformMimeTypes(output, textModel.metadata.displayOrder as string[] || []);
+							const ret = this._transformMimeTypes(output, output.outputId, textModel.metadata.displayOrder as string[] || []);
 							const orderedMimeTypes = ret.orderedMimeTypes!;
 							const pickedMimeTypeIndex = ret.pickedMimeTypeIndex!;
 							output.pickedMimeTypeIndex = pickedMimeTypeIndex;
 							output.orderedMimeTypes = orderedMimeTypes;
 
 							if (orderedMimeTypes[pickedMimeTypeIndex!].rendererId && orderedMimeTypes[pickedMimeTypeIndex].rendererId !== BUILTIN_RENDERER_ID) {
-								outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType, output: output });
+								outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType, output: output, outputId: output.outputId });
 								renderers.add(orderedMimeTypes[pickedMimeTypeIndex].rendererId!);
 							}
 						}
@@ -505,14 +506,14 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 			const outputRequest: IOutputRenderRequestOutputInfo[] = [];
 			outputs.map((output, index) => {
 				if (output.outputKind === CellOutputKind.Rich) {
-					const ret = this._transformMimeTypes(output, textModel.metadata.displayOrder as string[] || []);
+					const ret = this._transformMimeTypes(output, output.outputId, textModel.metadata.displayOrder as string[] || []);
 					const orderedMimeTypes = ret.orderedMimeTypes!;
 					const pickedMimeTypeIndex = ret.pickedMimeTypeIndex!;
 					output.pickedMimeTypeIndex = pickedMimeTypeIndex;
 					output.orderedMimeTypes = orderedMimeTypes;
 
 					if (orderedMimeTypes[pickedMimeTypeIndex!].rendererId && orderedMimeTypes[pickedMimeTypeIndex].rendererId !== BUILTIN_RENDERER_ID) {
-						outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType, output: output });
+						outputRequest.push({ index, handlerId: orderedMimeTypes[pickedMimeTypeIndex].rendererId!, mimeType: orderedMimeTypes[pickedMimeTypeIndex].mimeType, output: output, outputId: output.outputId });
 						renderers.add(orderedMimeTypes[pickedMimeTypeIndex].rendererId!);
 					}
 				}
@@ -536,6 +537,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 				outputs: [
 					{
 						index: 0,
+						outputId: generateUuid(),
 						handlerId: rendererId,
 						mimeType: mimeType,
 						output: output
@@ -560,7 +562,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		return;
 	}
 
-	private _transformMimeTypes(output: IDisplayOutput, documentDisplayOrder: string[]): ITransformedDisplayOutputDto {
+	private _transformMimeTypes(output: IDisplayOutput, outputId: string, documentDisplayOrder: string[]): ITransformedDisplayOutputDto {
 		let mimeTypes = Object.keys(output.data);
 		let coreDisplayOrder = this._displayOrder;
 		const sorted = sortMimeTypes(mimeTypes, coreDisplayOrder?.userOrder || [], documentDisplayOrder, coreDisplayOrder?.defaultOrder || []);
@@ -605,6 +607,7 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 
 		return {
 			outputKind: output.outputKind,
+			outputId,
 			data: output.data,
 			orderedMimeTypes: orderMimeTypes,
 			pickedMimeTypeIndex: 0

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

@@ -541,7 +541,7 @@ ${loaderJs}
 			}
 		}
 
-		let outputId = UUID.generateUuid();
+		let outputId = output.outputKind === CellOutputKind.Rich ? output.outputId : UUID.generateUuid();
 		let apiNamespace: string | undefined;
 		if (output.outputKind === CellOutputKind.Rich && output.pickedMimeTypeIndex !== undefined) {
 			const pickedMimeTypeRenderer = output.orderedMimeTypes?.[output.pickedMimeTypeIndex];

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

@@ -251,9 +251,17 @@ function webviewPreloads() {
 		return mapped.event;
 	}
 
+	interface ICreateCellInfo {
+		outputId: string;
+		element: HTMLElement;
+	}
+
+	interface IDestroyCellInfo {
+		outputId: string;
+	}
 
-	const onWillDestroyCell = createEmitter<[string | undefined /* namespace */, string | undefined /* cell uri */]>();
-	const onDidCreateCell = createEmitter<[string | undefined /* namespace */, HTMLElement]>();
+	const onWillDestroyOutput = createEmitter<[string | undefined /* namespace */, IDestroyCellInfo | undefined /* cell uri */]>();
+	const onDidCreateOutput = createEmitter<[string | undefined /* namespace */, ICreateCellInfo]>();
 
 	const matchesNs = (namespace: string, query: string | undefined) => namespace === '*' || query === namespace || query === 'undefined';
 
@@ -271,8 +279,8 @@ function webviewPreloads() {
 				const state = vscode.getState();
 				return typeof state === 'object' && state ? state[namespace] as T : undefined;
 			},
-			onWillDestroyCell: mapEmitter(onWillDestroyCell, ([ns, cellUri]) => matchesNs(namespace, ns) ? cellUri : dontEmit),
-			onDidCreateCell: mapEmitter(onDidCreateCell, ([ns, element]) => matchesNs(namespace, ns) ? element : dontEmit),
+			onWillDestroyOutput: mapEmitter(onWillDestroyOutput, ([ns, data]) => matchesNs(namespace, ns) ? data : dontEmit),
+			onDidCreateOutput: mapEmitter(onDidCreateOutput, ([ns, data]) => matchesNs(namespace, ns) ? data : dontEmit),
 		};
 	};
 
@@ -319,7 +327,7 @@ function webviewPreloads() {
 					// eval
 					domEval(outputNode);
 					resizeObserve(outputNode, outputId);
-					onDidCreateCell.fire([event.data.apiNamespace, outputNode]);
+					onDidCreateOutput.fire([event.data.apiNamespace, { element: outputNode, outputId }]);
 
 					vscode.postMessage({
 						__vscode_notebook_message: true,
@@ -347,7 +355,7 @@ function webviewPreloads() {
 					break;
 				}
 			case 'clear':
-				onWillDestroyCell.fire([undefined, undefined]);
+				onWillDestroyOutput.fire([undefined, undefined]);
 				document.getElementById('container')!.innerHTML = '';
 				for (let i = 0; i < observers.length; i++) {
 					observers[i].disconnect();
@@ -358,7 +366,7 @@ function webviewPreloads() {
 			case 'clearOutput':
 				{
 					const id = event.data.id;
-					onWillDestroyCell.fire([event.data.apiNamespace, event.data.cellUri]);
+					onWillDestroyOutput.fire([event.data.apiNamespace, { outputId: id }]);
 					let output = document.getElementById(id);
 					if (output && output.parentNode) {
 						document.getElementById(id)!.parentNode!.removeChild(output);

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

@@ -186,6 +186,7 @@ export interface IOrderedMimeType {
 
 export interface ITransformedDisplayOutputDto {
 	outputKind: CellOutputKind.Rich;
+	outputId: string;
 	data: { [key: string]: any; }
 	metadata?: NotebookCellOutputMetadata;
 
@@ -206,6 +207,7 @@ export type IRawOutput = IDisplayOutput | IStreamOutput | IErrorOutput;
 
 export interface IOutputRenderRequestOutputInfo {
 	index: number;
+	outputId: string;
 	handlerId: string;
 	mimeType: string;
 	output?: IRawOutput;
@@ -222,6 +224,7 @@ export interface IOutputRenderRequest<T> {
 
 export interface IOutputRenderResponseOutputInfo {
 	index: number;
+	outputId: string;
 	mimeType: string;
 	handlerId: string;
 	transformedOutput: string;