Simplify state used for webviews

We previously used nested states to store some additional metadata alongside the real webview state. This is overly complicated. This change switches us to using a single top level state field, while also adding some code to handle migration from the old state structure

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

@@ -22,8 +22,9 @@ import { ACTIVE_GROUP, IEditorService } from 'vs/workbench/services/editor/commo
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 import { extHostNamedCustomer } from '../common/extHostCustomers';
 import { IProductService } from 'vs/platform/product/common/product';
+import { startsWith } from 'vs/base/common/strings';
 
-interface MainThreadWebviewState {
+interface OldMainThreadWebviewState {
 	readonly viewType: string;
 	state: any;
 }
@@ -43,7 +44,7 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 
 
 	private readonly _proxy: ExtHostWebviewsShape;
-	private readonly _webviews = new Map<WebviewPanelHandle, WebviewEditorInput<MainThreadWebviewState>>();
+	private readonly _webviews = new Map<WebviewPanelHandle, WebviewEditorInput>();
 	private readonly _revivers = new Map<string, IDisposable>();
 
 	private _activeWebview: WebviewPanelHandle | undefined = undefined;
@@ -67,8 +68,12 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 		// This reviver's only job is to activate webview extensions
 		// This should trigger the real reviver to be registered from the extension host side.
 		this._register(_webviewEditorService.registerReviver({
-			canRevive: (webview: WebviewEditorInput<any>) => {
-				const viewType = webview.state && webview.state.viewType;
+			canRevive: (webview: WebviewEditorInput) => {
+				if (!webview.state) {
+					return false;
+				}
+
+				const viewType = this.fromInternalWebviewViewType(webview.viewType);
 				if (typeof viewType === 'string') {
 					extensionService.activateByEvent(`onWebviewPanel:${viewType}`);
 				}
@@ -96,11 +101,7 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 		const webview = this._webviewEditorService.createWebview(handle, this.getInternalWebviewViewType(viewType), title, mainThreadShowOptions, reviveWebviewOptions(options), {
 			location: URI.revive(extensionLocation),
 			id: extensionId
-		}, this.createWebviewEventDelegate(handle)) as WebviewEditorInput<MainThreadWebviewState>;
-		webview.state = {
-			viewType: viewType,
-			state: undefined
-		};
+		}, this.createWebviewEventDelegate(handle)) as WebviewEditorInput;
 
 		this._webviews.set(handle, webview);
 
@@ -176,17 +177,31 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 
 		this._revivers.set(viewType, this._webviewEditorService.registerReviver({
 			canRevive: (webview) => {
-				return webview.state && webview.viewType === this.getInternalWebviewViewType(viewType);
+				return !!webview.state && webview.viewType === this.getInternalWebviewViewType(viewType);
 			},
 			reviveWebview: async (webview): Promise<void> => {
-				const viewType = webview.state.viewType;
-				const handle = 'revival-' + MainThreadWebviews.revivalPool++;
+				const viewType = this.fromInternalWebviewViewType(webview.viewType);
+				if (!viewType) {
+					webview.html = MainThreadWebviews.getDeserializationFailedContents(webview.viewType);
+					return;
+				}
+
+				const handle = `revival-${MainThreadWebviews.revivalPool++}`;
 				this._webviews.set(handle, webview);
 				webview._events = this.createWebviewEventDelegate(handle);
 				let state = undefined;
-				if (webview.state.state) {
+				if (webview.state) {
 					try {
-						state = JSON.parse(webview.state.state);
+						// Check for old-style webview state first which stored state inside another state object
+						// TODO: remove this after 1.37 ships.
+						if (
+							typeof (webview.state as unknown as OldMainThreadWebviewState).viewType === 'string' &&
+							'state' in (webview.state as unknown as OldMainThreadWebviewState)
+						) {
+							state = JSON.parse((webview.state as any).state);
+						} else {
+							state = JSON.parse(webview.state);
+						}
 					} catch {
 						// noop
 					}
@@ -216,6 +231,13 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 		return `mainThreadWebview-${viewType}`;
 	}
 
+	private fromInternalWebviewViewType(viewType: string): string | undefined {
+		if (!startsWith(viewType, 'mainThreadWebview-')) {
+			return undefined;
+		}
+		return viewType.replace(/^mainThreadWebview-/, '');
+	}
+
 	private createWebviewEventDelegate(handle: WebviewPanelHandle) {
 		return {
 			onDidClickLink: (uri: URI) => this.onDidClickLink(handle, uri),
@@ -230,7 +252,7 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 				if (!webview || webview.isDisposed()) {
 					return;
 				}
-				(webview as WebviewEditorInput<MainThreadWebviewState>).state!.state = newState;
+				webview.state = newState;
 			}
 		};
 	}

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

@@ -256,7 +256,7 @@ export class WebviewEditor extends BaseEditor {
 				this._webview.initialScrollProgress = input.scrollYPercentage;
 			}
 
-			this._webview.state = input.state ? input.state.state : undefined;
+			this._webview.state = input.state;
 
 			this._content!.setAttribute('aria-flowto', this._webviewContent.id);
 

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

@@ -52,7 +52,7 @@ class WebviewIconsManager {
 	}
 }
 
-export class WebviewEditorInput<State = any> extends EditorInput {
+export class WebviewEditorInput extends EditorInput {
 
 	private readonly iconsManager = new WebviewIconsManager();
 
@@ -70,7 +70,7 @@ export class WebviewEditorInput<State = any> extends EditorInput {
 	private readonly _webviewDisposables = this._register(new DisposableStore());
 	private _group?: GroupIdentifier;
 	private _scrollYPercentage: number = 0;
-	private _state: State | undefined;
+	public state: string | undefined;
 
 	public readonly extension?: {
 		readonly location: URI;
@@ -82,7 +82,7 @@ export class WebviewEditorInput<State = any> extends EditorInput {
 		public readonly viewType: string,
 		name: string,
 		options: WebviewInputOptions,
-		state: State,
+		state: string,
 		events: WebviewEvents,
 		extension: undefined | {
 			readonly location: URI;
@@ -95,7 +95,7 @@ export class WebviewEditorInput<State = any> extends EditorInput {
 		this._name = name;
 		this._options = options;
 		this._events = events;
-		this._state = state;
+		this.state = state;
 		this.extension = extension;
 	}
 
@@ -178,14 +178,6 @@ export class WebviewEditorInput<State = any> extends EditorInput {
 		}
 	}
 
-	public get state(): State | undefined {
-		return this._state;
-	}
-
-	public set state(value: State | undefined) {
-		this._state = value;
-	}
-
 	public get options(): WebviewInputOptions {
 		return this._options;
 	}
@@ -308,7 +300,7 @@ export class RevivedWebviewEditorInput extends WebviewEditorInput {
 		viewType: string,
 		name: string,
 		options: WebviewInputOptions,
-		state: any,
+		state: string,
 		events: WebviewEvents,
 		extension: undefined | {
 			readonly location: URI;

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

@@ -68,12 +68,15 @@ export class WebviewEditorInputFactory implements IEditorInputFactory {
 		const extensionLocation = reviveUri(data.extensionLocation);
 		const extensionId = data.extensionId ? new ExtensionIdentifier(data.extensionId) : undefined;
 		const iconPath = reviveIconPath(data.iconPath);
-		return this._webviewService.reviveWebview(generateUuid(), data.viewType, data.title, iconPath, data.state, data.options, extensionLocation ? {
+		const state = reviveState(data.state);
+
+		return this._webviewService.reviveWebview(generateUuid(), data.viewType, data.title, iconPath, state, data.options, extensionLocation ? {
 			location: extensionLocation,
 			id: extensionId
 		} : undefined, data.group);
 	}
 }
+
 function reviveIconPath(data: SerializedIconPath | undefined) {
 	if (!data) {
 		return undefined;
@@ -98,3 +101,21 @@ function reviveUri(data: string | UriComponents | undefined): URI | undefined {
 		return undefined;
 	}
 }
+
+
+function reviveState(state: unknown | undefined): undefined | string {
+	if (!state) {
+		return undefined;
+	}
+
+	if (typeof state === 'string') {
+		return state;
+	}
+
+	// Likely an old style state. Unwrap to a simple state object
+	// Remove after 1.37
+	if ('state' in (state as any) && typeof (state as any).state === 'string') {
+		return (state as any).state;
+	}
+	return undefined;
+}
\ No newline at end of file