Reuse LazilyResolvedWebviewEditorInput for custom editors as well

src/vs/workbench/contrib/customEditor/browser/customEditorInput.ts

@@ -6,25 +6,23 @@
 import { memoize } from 'vs/base/common/decorators';
 import { Emitter } from 'vs/base/common/event';
 import { UnownedDisposable } from 'vs/base/common/lifecycle';
+import { Schemas } from 'vs/base/common/network';
 import { basename } from 'vs/base/common/path';
-import { isEqual, DataUri } from 'vs/base/common/resources';
+import { DataUri, isEqual } from 'vs/base/common/resources';
 import { URI } from 'vs/base/common/uri';
 import { WebviewContentState } from 'vs/editor/common/modes';
 import { IDialogService } from 'vs/platform/dialogs/common/dialogs';
-import { IEditorModel } from 'vs/platform/editor/common/editor';
 import { ILabelService } from 'vs/platform/label/common/label';
 import { ConfirmResult, IEditorInput, Verbosity } from 'vs/workbench/common/editor';
 import { WebviewEditorOverlay } from 'vs/workbench/contrib/webview/browser/webview';
-import { WebviewInput } from 'vs/workbench/contrib/webview/browser/webviewEditorInput';
+import { LazilyResolvedWebviewEditorInput } from 'vs/workbench/contrib/webview/browser/webviewEditorInput';
 import { IWebviewEditorService } from 'vs/workbench/contrib/webview/browser/webviewEditorService';
 import { promptSave } from 'vs/workbench/services/textfile/browser/textFileService';
-import { Schemas } from 'vs/base/common/network';
 
-export class CustomFileEditorInput extends WebviewInput {
+export class CustomFileEditorInput extends LazilyResolvedWebviewEditorInput {
 
 	public static typeId = 'workbench.editors.webviewEditor';
 
-	private _hasResolved = false;
 	private readonly _editorResource: URI;
 	private _state = WebviewContentState.Readonly;
 
@@ -33,11 +31,11 @@ export class CustomFileEditorInput extends WebviewInput {
 		viewType: string,
 		id: string,
 		webview: UnownedDisposable<WebviewEditorOverlay>,
-		@ILabelService private readonly labelService: ILabelService,
-		@IWebviewEditorService private readonly _webviewEditorService: IWebviewEditorService,
+		@IWebviewEditorService webviewEditorService: IWebviewEditorService,
 		@IDialogService private readonly dialogService: IDialogService,
+		@ILabelService private readonly labelService: ILabelService,
 	) {
-		super(id, viewType, '', webview);
+		super(id, viewType, '', webview, webviewEditorService);
 		this._editorResource = resource;
 	}
 
@@ -112,14 +110,6 @@ export class CustomFileEditorInput extends WebviewInput {
 		}
 	}
 
-	public async resolve(): Promise<IEditorModel> {
-		if (!this._hasResolved) {
-			this._hasResolved = true;
-			await this._webviewEditorService.resolveWebview(this);
-		}
-		return super.resolve();
-	}
-
 	public setState(newState: WebviewContentState): void {
 		this._state = newState;
 		this._onDidChangeDirty.fire();

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

@@ -4,11 +4,12 @@
  *--------------------------------------------------------------------------------------------*/
 import * as dom from 'vs/base/browser/dom';
 import { memoize } from 'vs/base/common/decorators';
+import { UnownedDisposable as Unowned } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IEditorModel } from 'vs/platform/editor/common/editor';
 import { EditorInput, EditorModel, GroupIdentifier, IEditorInput, Verbosity } from 'vs/workbench/common/editor';
 import { WebviewEditorOverlay } from 'vs/workbench/contrib/webview/browser/webview';
-import { UnownedDisposable as Unowned } from 'vs/base/common/lifecycle';
+import { IWebviewEditorService } from 'vs/workbench/contrib/webview/browser/webviewEditorService';
 
 const WebviewPanelResourceScheme = 'webview-panel';
 
@@ -141,24 +142,20 @@ export class WebviewInput extends EditorInput {
 	}
 }
 
-export class RevivedWebviewEditorInput extends WebviewInput {
-	private _revived: boolean = false;
-
+export class LazilyResolvedWebviewEditorInput extends WebviewInput {
 	constructor(
 		id: string,
 		viewType: string,
 		name: string,
-		private readonly reviver: (input: WebviewInput) => Promise<void>,
-		webview: Unowned<WebviewEditorOverlay>
+		webview: Unowned<WebviewEditorOverlay>,
+		@IWebviewEditorService private readonly _webviewEditorService: IWebviewEditorService,
 	) {
 		super(id, viewType, name, webview);
 	}
 
+	@memoize
 	public async resolve(): Promise<IEditorModel> {
-		if (!this._revived) {
-			this._revived = true;
-			await this.reviver(this);
-		}
+		await this._webviewEditorService.resolveWebview(this);
 		return super.resolve();
 	}
 }

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

@@ -17,7 +17,7 @@ import { GroupIdentifier } from 'vs/workbench/common/editor';
 import { IWebviewService, WebviewContentOptions, WebviewOptions } from 'vs/workbench/contrib/webview/browser/webview';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { ACTIVE_GROUP_TYPE, IEditorService, SIDE_GROUP_TYPE } from 'vs/workbench/services/editor/common/editorService';
-import { RevivedWebviewEditorInput, WebviewInput } from './webviewEditorInput';
+import { LazilyResolvedWebviewEditorInput, WebviewInput } from './webviewEditorInput';
 
 export const IWebviewEditorService = createDecorator<IWebviewEditorService>('webviewEditorService');
 
@@ -190,7 +190,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 		const webview = this.createWebiew(id, extension, options);
 		webview.state = state;
 
-		const webviewInput = new RevivedWebviewEditorInput(id, viewType, title, (webview: WebviewInput) => this.resolveWebview(webview), new UnownedDisposable(webview));
+		const webviewInput = new LazilyResolvedWebviewEditorInput(id, viewType, title, new UnownedDisposable(webview), this);
 		webviewInput.iconPath = iconPath;
 
 		if (typeof group === 'number') {
@@ -219,7 +219,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 
 		// Revived webviews may not have an actively registered reviver but we still want to presist them
 		// since a reviver should exist when it is actually needed.
-		return webview instanceof RevivedWebviewEditorInput;
+		return webview instanceof LazilyResolvedWebviewEditorInput;
 	}
 
 	private async tryRevive(