Re-resolve webview views when extension host restarts

Fixes #108555

Previously webviews were left hanging when the extension host died. With this change, we now try to re-create them once the extension host restarts

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

@@ -128,9 +128,7 @@ export class MainThreadWebviewPanels extends Disposable implements extHostProtoc
 	dispose() {
 		super.dispose();
 
-		for (const disposable of this._editorProviders.values()) {
-			disposable.dispose();
-		}
+		dispose(this._editorProviders.values());
 		this._editorProviders.clear();
 	}
 

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

@@ -5,7 +5,7 @@
 
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, dispose, IDisposable } from 'vs/base/common/lifecycle';
 import { MainThreadWebviews, reviveWebviewExtension } from 'vs/workbench/api/browser/mainThreadWebviews';
 import * as extHostProtocol from 'vs/workbench/api/common/extHost.protocol';
 import { IWebviewViewService, WebviewView } from 'vs/workbench/contrib/webviewView/browser/webviewViewService';
@@ -28,6 +28,15 @@ export class MainThreadWebviewsViews extends Disposable implements extHostProtoc
 		this._proxy = context.getProxy(extHostProtocol.ExtHostContext.ExtHostWebviewViews);
 	}
 
+	dispose() {
+		super.dispose();
+
+		dispose(this._webviewViewProviders.values());
+		this._webviewViewProviders.clear();
+
+		dispose(this._webviewViews.values());
+	}
+
 	public $setWebviewViewTitle(handle: extHostProtocol.WebviewHandle, value: string | undefined): void {
 		const webviewView = this.getWebviewView(handle);
 		webviewView.title = value;
@@ -54,7 +63,7 @@ export class MainThreadWebviewsViews extends Disposable implements extHostProtoc
 
 		const extension = reviveWebviewExtension(extensionData);
 
-		this._webviewViewService.register(viewType, {
+		const registration = this._webviewViewService.register(viewType, {
 			resolve: async (webviewView: WebviewView, cancellation: CancellationToken) => {
 				const handle = webviewView.webview.id;
 
@@ -93,6 +102,8 @@ export class MainThreadWebviewsViews extends Disposable implements extHostProtoc
 				}
 			}
 		});
+
+		this._webviewViewProviders.set(viewType, registration);
 	}
 
 	public $unregisterWebviewViewProvider(viewType: string): void {

src/vs/workbench/contrib/webviewView/browser/webviewViewPane.ts

@@ -5,7 +5,7 @@
 
 import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { Emitter } from 'vs/base/common/event';
-import { toDisposable } from 'vs/base/common/lifecycle';
+import { DisposableStore, MutableDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { setImmediate } from 'vs/base/common/platform';
 import { MenuId } from 'vs/platform/actions/common/actions';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
@@ -34,7 +34,8 @@ const storageKeys = {
 
 export class WebviewViewPane extends ViewPane {
 
-	private _webview?: WebviewOverlay;
+	private readonly _webview = this._register(new MutableDisposable<WebviewOverlay>());
+	private readonly _webviewDisposables = this._register(new DisposableStore());
 	private _activated = false;
 
 	private _container?: HTMLElement;
@@ -71,6 +72,14 @@ export class WebviewViewPane extends ViewPane {
 		this.viewState = this.memento.getMemento(StorageScope.WORKSPACE);
 
 		this._register(this.onDidChangeBodyVisibility(() => this.updateTreeVisibility()));
+
+		this._register(this.webviewViewService.onNewResolverRegistered(e => {
+			if (e.viewType === this.id) {
+				// Potentially re-activate if we have a new resolver
+				this.updateTreeVisibility();
+			}
+		}));
+
 		this.updateTreeVisibility();
 	}
 
@@ -83,14 +92,12 @@ export class WebviewViewPane extends ViewPane {
 	dispose() {
 		this._onDispose.fire();
 
-		this._webview?.dispose();
-
 		super.dispose();
 	}
 
 	focus(): void {
 		super.focus();
-		this._webview?.focus();
+		this._webview.value?.focus();
 	}
 
 	renderBody(container: HTMLElement): void {
@@ -102,7 +109,7 @@ export class WebviewViewPane extends ViewPane {
 			this._resizeObserver = new ResizeObserver(() => {
 				setImmediate(() => {
 					if (this._container) {
-						this._webview?.layoutWebviewOverElement(this._container);
+						this._webview.value?.layoutWebviewOverElement(this._container);
 					}
 				});
 			});
@@ -115,8 +122,8 @@ export class WebviewViewPane extends ViewPane {
 	}
 
 	public saveState() {
-		if (this._webview) {
-			this.viewState[storageKeys.webviewState] = this._webview.state;
+		if (this._webview.value) {
+			this.viewState[storageKeys.webviewState] = this._webview.value.state;
 		}
 
 		this.memento.saveMemento();
@@ -126,21 +133,21 @@ export class WebviewViewPane extends ViewPane {
 	protected layoutBody(height: number, width: number): void {
 		super.layoutBody(height, width);
 
-		if (!this._webview) {
+		if (!this._webview.value) {
 			return;
 		}
 
 		if (this._container) {
-			this._webview.layoutWebviewOverElement(this._container, { width, height });
+			this._webview.value.layoutWebviewOverElement(this._container, { width, height });
 		}
 	}
 
 	private updateTreeVisibility() {
 		if (this.isBodyVisible()) {
 			this.activate();
-			this._webview?.claim(this);
+			this._webview.value?.claim(this);
 		} else {
-			this._webview?.release(this);
+			this._webview.value?.release(this);
 		}
 	}
 
@@ -151,17 +158,20 @@ export class WebviewViewPane extends ViewPane {
 			const webviewId = `webviewView-${this.id.replace(/[^a-z0-9]/gi, '-')}`.toLowerCase();
 			const webview = this.webviewService.createWebviewOverlay(webviewId, {}, {}, undefined);
 			webview.state = this.viewState[storageKeys.webviewState];
-			this._webview = webview;
+			this._webview.value = webview;
 
-			this._register(toDisposable(() => {
-				this._webview?.release(this);
+			if (this._container) {
+				this._webview.value?.layoutWebviewOverElement(this._container);
+			}
+
+			this._webviewDisposables.add(toDisposable(() => {
+				this._webview.value?.release(this);
 			}));
 
-			this._register(webview.onDidUpdateState(() => {
+			this._webviewDisposables.add(webview.onDidUpdateState(() => {
 				this.viewState[storageKeys.webviewState] = webview.state;
 			}));
-
-			const source = this._register(new CancellationTokenSource());
+			const source = this._webviewDisposables.add(new CancellationTokenSource());
 
 			this.withProgress(async () => {
 				await this.extensionService.activateByEvent(`onView:${this.id}`);
@@ -178,6 +188,13 @@ export class WebviewViewPane extends ViewPane {
 					get description(): string | undefined { return self.titleDescription; },
 					set description(value: string | undefined) { self.updateTitleDescription(value); },
 
+					dispose: () => {
+						// Only reset and clear the webview itself. Don't dispose of the view container
+						this._activated = false;
+						this._webview.clear();
+						this._webviewDisposables.clear();
+					},
+
 					show: (preserveFocus) => {
 						this.viewService.openView(this.id, !preserveFocus);
 					}

src/vs/workbench/contrib/webviewView/browser/webviewViewService.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationToken } from 'vs/base/common/cancellation';
-import { Event } from 'vs/base/common/event';
+import { Emitter, Event } from 'vs/base/common/event';
 import { Disposable, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { WebviewOverlay } from 'vs/workbench/contrib/webview/browser/webview';
@@ -20,6 +20,8 @@ export interface WebviewView {
 	readonly onDidChangeVisibility: Event<boolean>;
 	readonly onDispose: Event<void>;
 
+	dispose(): void;
+
 	show(preserveFocus: boolean): void;
 }
 
@@ -31,6 +33,8 @@ export interface IWebviewViewService {
 
 	readonly _serviceBrand: undefined;
 
+	readonly onNewResolverRegistered: Event<{ readonly viewType: string }>;
+
 	register(type: string, resolver: IWebviewViewResolver): IDisposable;
 
 	resolve(viewType: string, webview: WebviewView, cancellation: CancellationToken): Promise<void>;
@@ -40,16 +44,20 @@ export class WebviewViewService extends Disposable implements IWebviewViewServic
 
 	readonly _serviceBrand: undefined;
 
-	private readonly _views = new Map<string, IWebviewViewResolver>();
+	private readonly _resolvers = new Map<string, IWebviewViewResolver>();
 
 	private readonly _awaitingRevival = new Map<string, { webview: WebviewView, resolve: () => void }>();
 
+	private readonly _onNewResolverRegistered = this._register(new Emitter<{ readonly viewType: string }>());
+	public readonly onNewResolverRegistered = this._onNewResolverRegistered.event;
+
 	register(viewType: string, resolver: IWebviewViewResolver): IDisposable {
-		if (this._views.has(viewType)) {
+		if (this._resolvers.has(viewType)) {
 			throw new Error(`View resolver already registered for ${viewType}`);
 		}
 
-		this._views.set(viewType, resolver);
+		this._resolvers.set(viewType, resolver);
+		this._onNewResolverRegistered.fire({ viewType: viewType });
 
 		const pending = this._awaitingRevival.get(viewType);
 		if (pending) {
@@ -60,12 +68,12 @@ export class WebviewViewService extends Disposable implements IWebviewViewServic
 		}
 
 		return toDisposable(() => {
-			this._views.delete(viewType);
+			this._resolvers.delete(viewType);
 		});
 	}
 
 	resolve(viewType: string, webview: WebviewView, cancellation: CancellationToken): Promise<void> {
-		const resolver = this._views.get(viewType);
+		const resolver = this._resolvers.get(viewType);
 		if (!resolver) {
 			if (this._awaitingRevival.has(viewType)) {
 				throw new Error('View already awaiting revival');