notebooks: rerender output if the webview reloads (#97282)

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

@@ -77,6 +77,7 @@ export interface ICreationRequestMessage {
 	outputId: string;
 	top: number;
 	left: number;
+	initiallyHidden?: boolean;
 }
 
 export interface IContentWidgetTopRequest {
@@ -105,6 +106,13 @@ export interface IUpdatePreloadResourceMessage {
 	resources: string[];
 }
 
+interface ICachedInset {
+	outputId: string;
+	cell: CodeCellViewModel;
+	preloads: ReadonlySet<number>;
+	cachedCreation: ICreationRequestMessage;
+}
+
 function html(strings: TemplateStringsArray, ...values: any[]): string {
 	let str = '';
 	strings.forEach((string, i) => {
@@ -119,7 +127,7 @@ let version = 0;
 export class BackLayerWebView extends Disposable {
 	element: HTMLElement;
 	webview!: WebviewElement;
-	insetMapping: Map<IOutput, { outputId: string, cell: CodeCellViewModel, cacheOffset: number | undefined }> = new Map();
+	insetMapping: Map<IOutput, ICachedInset> = new Map();
 	hiddenInsetMapping: Set<IOutput> = new Set();
 	reversedInsetMapping: Map<string, IOutput> = new Map();
 	preloadsCache: Map<string, boolean> = new Map();
@@ -378,6 +386,9 @@ ${loaderJs}
 							height: outputNode.clientHeight
 						}
 					});
+
+					// don't hide until after this step so that the height is right
+					cellOutputContainer.style.display = event.data.initiallyHidden ? 'none' : 'block';
 				}
 				break;
 			case 'view-scroll':
@@ -461,6 +472,14 @@ ${loaderJs}
 			this.openerService.open(link, { fromUserGesture: true });
 		}));
 
+		this._register(this.webview.onDidReload(() => {
+			this.preloadsCache.clear();
+			for (const [output, inset] of this.insetMapping.entries()) {
+				this.updateRendererPreloads(inset.preloads);
+				this.webview.sendMessage({ ...inset.cachedCreation, initiallyHidden: this.hiddenInsetMapping.has(output) });
+			}
+		}));
+
 		this._register(this.webview.onMessage((data: IMessage) => {
 			if (data.__vscode_notebook_message) {
 				if (data.type === 'dimension') {
@@ -550,7 +569,7 @@ ${loaderJs}
 			return true;
 		}
 
-		if (outputOffset === outputCache.cacheOffset) {
+		if (outputOffset === outputCache.cachedCreation.top) {
 			return false;
 		}
 
@@ -564,7 +583,7 @@ ${loaderJs}
 			let outputIndex = item.cell.outputs.indexOf(item.output);
 
 			let outputOffset = item.cellTop + item.cell.getOutputOffset(outputIndex);
-			outputCache.cacheOffset = outputOffset;
+			outputCache.cachedCreation.top = outputOffset;
 			this.hiddenInsetMapping.delete(item.output);
 
 			return {
@@ -614,7 +633,7 @@ ${loaderJs}
 		};
 
 		this.webview.sendMessage(message);
-		this.insetMapping.set(output, { outputId: outputId, cell: cell, cacheOffset: initialTop });
+		this.insetMapping.set(output, { outputId: outputId, cell: cell, preloads, cachedCreation: message });
 		this.hiddenInsetMapping.delete(output);
 		this.reversedInsetMapping.set(outputId, output);
 	}
@@ -669,7 +688,7 @@ ${loaderJs}
 		}, 50);
 	}
 
-	updateRendererPreloads(preloads: Set<number>) {
+	updateRendererPreloads(preloads: ReadonlySet<number>) {
 		let resources: string[] = [];
 		let extensionLocations: URI[] = [];
 		preloads.forEach(preload => {

src/vs/workbench/contrib/webview/browser/baseWebviewElement.ts

@@ -21,6 +21,7 @@ export const enum WebviewMessageChannels {
 	didScroll = 'did-scroll',
 	didFocus = 'did-focus',
 	didBlur = 'did-blur',
+	didLoad = 'did-load',
 	doUpdateState = 'do-update-state',
 	doReload = 'do-reload',
 	loadResource = 'load-resource',
@@ -153,6 +154,9 @@ export abstract class BaseWebview<T extends HTMLElement> extends Disposable {
 	private readonly _onDidClickLink = this._register(new Emitter<string>());
 	public readonly onDidClickLink = this._onDidClickLink.event;
 
+	private readonly _onDidReload = this._register(new Emitter<void>());
+	public readonly onDidReload = this._onDidReload.event;
+
 	private readonly _onMessage = this._register(new Emitter<any>());
 	public readonly onMessage = this._onMessage.event;
 
@@ -216,6 +220,10 @@ export abstract class BaseWebview<T extends HTMLElement> extends Disposable {
 
 	public reload(): void {
 		this.doUpdateContent();
+		const subscription = this._register(this.on(WebviewMessageChannels.didLoad, () => {
+			this._onDidReload.fire();
+			subscription.dispose();
+		}));
 	}
 
 	public set html(value: string) {

src/vs/workbench/contrib/webview/browser/dynamicWebviewEditorOverlay.ts

@@ -130,6 +130,7 @@ export class DynamicWebviewEditorOverlay extends Disposable implements WebviewOv
 			this._webviewEvents.add(webview.onMessage(x => { this._onMessage.fire(x); }));
 			this._webviewEvents.add(webview.onMissingCsp(x => { this._onMissingCsp.fire(x); }));
 			this._webviewEvents.add(webview.onDidWheel(x => { this._onDidWheel.fire(x); }));
+			this._webviewEvents.add(webview.onDidReload(() => { this._onDidReload.fire(); }));
 
 			this._webviewEvents.add(webview.onDidScroll(x => {
 				this._initialScrollProgress = x.scrollYPercentage;
@@ -189,6 +190,9 @@ export class DynamicWebviewEditorOverlay extends Disposable implements WebviewOv
 	private readonly _onDidClickLink = this._register(new Emitter<string>());
 	public readonly onDidClickLink: Event<string> = this._onDidClickLink.event;
 
+	private readonly _onDidReload = this._register(new Emitter<void>());
+	public readonly onDidReload = this._onDidReload.event;
+
 	private readonly _onDidScroll = this._register(new Emitter<{ scrollYPercentage: number; }>());
 	public readonly onDidScroll: Event<{ scrollYPercentage: number; }> = this._onDidScroll.event;
 

src/vs/workbench/contrib/webview/browser/pre/main.js

@@ -517,6 +517,8 @@
 						});
 						pendingMessages = [];
 					}
+
+					host.postMessage('did-load');
 				};
 
 				/**

src/vs/workbench/contrib/webview/browser/webview.ts

@@ -84,6 +84,7 @@ export interface Webview extends IDisposable {
 	readonly onDidScroll: Event<{ scrollYPercentage: number }>;
 	readonly onDidWheel: Event<IMouseWheelEvent>;
 	readonly onDidUpdateState: Event<string | undefined>;
+	readonly onDidReload: Event<void>;
 	readonly onMessage: Event<any>;
 	readonly onMissingCsp: Event<ExtensionIdentifier>;
 

src/vs/workbench/contrib/webview/electron-browser/pre/electron-index.js

@@ -79,4 +79,4 @@
 	});
 
 	require('../../browser/pre/main')(host);
-}());
\ No newline at end of file
+}());