Rework management of custom editor custom models

Fixes #91670

- Move model type logic out of `CustomEditorInput` and into two different implementations of `ICustomEditorModel`
- Add `CustomTextEditorModel` that owns a proper model reference to a text document. This should ensure the text document is disposed of if there are no more references to it
- Move most of the previous `CustomEditorModel` into `mainThreadWebviews` directly. This removes some of the indirection that was previously required (such as using `waitUntil`)

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

@@ -3,31 +3,37 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { createCancelablePromise } from 'vs/base/common/async';
+import { CancelablePromise, createCancelablePromise } from 'vs/base/common/async';
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { Disposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { Emitter, Event } from 'vs/base/common/event';
+import { Disposable, DisposableStore, IDisposable, IReference } from 'vs/base/common/lifecycle';
 import { Schemas } from 'vs/base/common/network';
+import { basename } from 'vs/base/common/path';
 import { isWeb } from 'vs/base/common/platform';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import * as modes from 'vs/editor/common/modes';
 import { localize } from 'vs/nls';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { IFileService } from 'vs/platform/files/common/files';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { ILabelService } from 'vs/platform/label/common/label';
 import { IOpenerService } from 'vs/platform/opener/common/opener';
 import { IProductService } from 'vs/platform/product/common/productService';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import * as extHostProtocol from 'vs/workbench/api/common/extHost.protocol';
 import { editorGroupToViewColumn, EditorViewColumn, viewColumnToEditorGroup } from 'vs/workbench/api/common/shared/editor';
-import { IEditorInput } from 'vs/workbench/common/editor';
+import { IEditorInput, IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
 import { DiffEditorInput } from 'vs/workbench/common/editor/diffEditorInput';
-import { CustomEditorInput, ModelType } from 'vs/workbench/contrib/customEditor/browser/customEditorInput';
+import { CustomEditorInput } from 'vs/workbench/contrib/customEditor/browser/customEditorInput';
 import { ICustomEditorModel, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
+import { CustomTextEditorModel } from 'vs/workbench/contrib/customEditor/common/customTextEditorModel';
 import { WebviewExtensionDescription, WebviewIcons } from 'vs/workbench/contrib/webview/browser/webview';
 import { WebviewInput } from 'vs/workbench/contrib/webview/browser/webviewEditorInput';
 import { ICreateWebViewShowOptions, IWebviewWorkbenchService, WebviewInputOptions } from 'vs/workbench/contrib/webview/browser/webviewWorkbenchService';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
+import { IWorkingCopy, IWorkingCopyBackup, IWorkingCopyService, WorkingCopyCapabilities } from 'vs/workbench/services/workingCopy/common/workingCopyService';
 import { extHostNamedCustomer } from '../common/extHostCustomers';
 
 /**
@@ -79,6 +85,11 @@ class WebviewViewTypeTransformer {
 	}
 }
 
+const enum ModelType {
+	Custom,
+	Text,
+}
+
 const webviewPanelViewType = new WebviewViewTypeTransformer('mainThreadWebview-');
 
 @extHostNamedCustomer(extHostProtocol.MainContext.MainThreadWebviews)
@@ -96,7 +107,6 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 	private readonly _webviewInputs = new WebviewInputStore();
 	private readonly _revivers = new Map<string, IDisposable>();
 	private readonly _editorProviders = new Map<string, IDisposable>();
-	private readonly _customEditorModels = new Map<string, { referenceCount: number }>();
 
 	constructor(
 		context: extHostProtocol.IExtHostContext,
@@ -108,7 +118,7 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		@IProductService private readonly _productService: IProductService,
 		@ITelemetryService private readonly _telemetryService: ITelemetryService,
 		@IWebviewWorkbenchService private readonly _webviewWorkbenchService: IWebviewWorkbenchService,
-		@IFileService private readonly _fileService: IFileService,
+		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 	) {
 		super();
 
@@ -286,16 +296,13 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 
 				webviewInput.webview.options = options;
 				webviewInput.webview.extension = extension;
-				webviewInput.modelType = modelType;
 
 				const resource = webviewInput.resource;
 
-				if (modelType === ModelType.Custom) {
-					const model = await this.retainCustomEditorModel(webviewInput, resource, viewType);
-					webviewInput.onDisposeWebview(() => {
-						this.releaseCustomEditorModel(model);
-					});
-				}
+				const modelRef = await this.getOrCreateCustomEditorModel(modelType, webviewInput, resource, viewType);
+				webviewInput.onDisposeWebview(() => {
+					modelRef.dispose();
+				});
 
 				try {
 					await this._proxy.$resolveWebviewEditor(
@@ -327,71 +334,27 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		this._customEditorService.models.disposeAllModelsForView(viewType);
 	}
 
-	private async retainCustomEditorModel(webviewInput: WebviewInput, resource: URI, viewType: string) {
-		const model = await this._customEditorService.models.resolve(webviewInput.resource, webviewInput.viewType);
-
-		const key = viewType + resource.toString();
-		const existingEntry = this._customEditorModels.get(key);
-		if (existingEntry) {
-			++existingEntry.referenceCount;
-			// no need to hook up listeners again
-			return model;
-		}
-		this._customEditorModels.set(key, { referenceCount: 1 });
-		const { editable } = await this._proxy.$createWebviewCustomEditorDocument(resource, viewType);
-
-		if (editable) {
-			model.onUndo(() => {
-				this._proxy.$undo(resource, viewType);
-			});
-
-			model.onRedo(() => {
-				this._proxy.$redo(resource, viewType);
-			});
-
-			model.onWillSave(e => {
-				e.waitUntil(this._proxy.$onSave(resource.toJSON(), viewType));
-			});
+	private async getOrCreateCustomEditorModel(
+		modelType: ModelType,
+		webviewInput: WebviewInput,
+		resource: URI,
+		viewType: string,
+	): Promise<IReference<ICustomEditorModel>> {
+		const existingModel = this._customEditorService.models.tryRetain(webviewInput.resource, webviewInput.viewType);
+		if (existingModel) {
+			return existingModel;
 		}
 
-		// Save as should always be implemented even if the model is readonly
-		model.onWillSaveAs(e => {
-			if (editable) {
-				e.waitUntil(this._proxy.$onSaveAs(e.resource.toJSON(), viewType, e.targetResource.toJSON()));
-			} else {
-				// Since the editor is readonly, just copy the file over
-				e.waitUntil(this._fileService.copy(e.resource, e.targetResource, false /* overwrite */));
-			}
-		});
-
-		model.onBackup(() => {
-			return createCancelablePromise(token =>
-				this._proxy.$backup(model.resource.toJSON(), viewType, token));
-		});
+		const model = modelType === ModelType.Text
+			? CustomTextEditorModel.create(this._instantiationService, viewType, resource)
+			: MainThreadCustomEditorModel.create(this._instantiationService, this._proxy, viewType, resource);
 
-		return model;
-	}
-
-	private async releaseCustomEditorModel(model: ICustomEditorModel) {
-		const key = model.viewType + model.resource;
-		const entry = this._customEditorModels.get(key);
-		if (!entry) {
-			throw new Error('Model not found');
-		}
-
-		--entry.referenceCount;
-		if (entry.referenceCount <= 0) {
-			this._proxy.$disposeWebviewCustomEditorDocument(model.resource, model.viewType);
-			this._customEditorService.models.disposeModel(model);
-			this._customEditorModels.delete(key);
-		}
+		return this._customEditorService.models.add(resource, viewType, model);
 	}
 
-
-
-	public $onDidChangeCustomDocumentState(resource: UriComponents, viewType: string, state: { dirty: boolean }) {
-		const model = this._customEditorService.models.get(URI.revive(resource), viewType);
-		if (!model) {
+	public async $onDidChangeCustomDocumentState(resource: UriComponents, viewType: string, state: { dirty: boolean }) {
+		const model = await this._customEditorService.models.get(URI.revive(resource), viewType);
+		if (!model || !(model instanceof MainThreadCustomEditorModel)) {
 			throw new Error('Could not find model for webview editor');
 		}
 		model.setDirty(state.dirty);
@@ -515,3 +478,152 @@ function reviveWebviewIcon(
 		? { light: URI.revive(value.light), dark: URI.revive(value.dark) }
 		: undefined;
 }
+
+namespace HotExitState {
+	export const enum Type {
+		Allowed,
+		NotAllowed,
+		Pending,
+	}
+
+	export const Allowed = Object.freeze({ type: Type.Allowed } as const);
+	export const NotAllowed = Object.freeze({ type: Type.NotAllowed } as const);
+
+	export class Pending {
+		readonly type = Type.Pending;
+
+		constructor(
+			public readonly operation: CancelablePromise<void>,
+		) { }
+	}
+
+	export type State = typeof Allowed | typeof NotAllowed | Pending;
+}
+
+class MainThreadCustomEditorModel extends Disposable implements ICustomEditorModel, IWorkingCopy {
+
+	private _hotExitState: HotExitState.State = HotExitState.Allowed;
+	private _dirty = false;
+
+	public static async create(instantiationService: IInstantiationService, proxy: extHostProtocol.ExtHostWebviewsShape, viewType: string, resource: URI) {
+		const { editable } = await proxy.$createWebviewCustomEditorDocument(resource, viewType);
+		return instantiationService.createInstance(MainThreadCustomEditorModel, proxy, viewType, resource, editable);
+	}
+
+	constructor(
+		private readonly _proxy: extHostProtocol.ExtHostWebviewsShape,
+		public readonly viewType: string,
+		private readonly _resource: URI,
+		private readonly _editable: boolean,
+		@IWorkingCopyService workingCopyService: IWorkingCopyService,
+		@ILabelService private readonly _labelService: ILabelService,
+		@IFileService private readonly _fileService: IFileService,
+	) {
+		super();
+		this._register(workingCopyService.registerWorkingCopy(this));
+	}
+
+	dispose() {
+		this._proxy.$disposeWebviewCustomEditorDocument(this.resource, this.viewType);
+		super.dispose();
+	}
+
+	//#region IWorkingCopy
+
+	public get resource() {
+		return this._resource;
+	}
+
+	public get name() {
+		return basename(this._labelService.getUriLabel(this._resource));
+	}
+
+	public get capabilities(): WorkingCopyCapabilities {
+		return 0;
+	}
+
+	public isDirty(): boolean {
+		return this._dirty;
+	}
+
+	private readonly _onDidChangeDirty: Emitter<void> = this._register(new Emitter<void>());
+	readonly onDidChangeDirty: Event<void> = this._onDidChangeDirty.event;
+
+	private readonly _onDidChangeContent: Emitter<void> = this._register(new Emitter<void>());
+	readonly onDidChangeContent: Event<void> = this._onDidChangeContent.event;
+
+	//#endregion
+
+	public setDirty(dirty: boolean): void {
+		this._onDidChangeContent.fire();
+
+		if (this._dirty !== dirty) {
+			this._dirty = dirty;
+			this._onDidChangeDirty.fire();
+		}
+	}
+
+	public async revert(_options?: IRevertOptions) {
+		this._proxy.$revert(this.resource, this.viewType);
+	}
+
+	public undo() {
+		this._proxy.$undo(this.resource, this.viewType);
+	}
+
+	public redo() {
+		this._proxy.$redo(this.resource, this.viewType);
+	}
+
+	public async save(_options?: ISaveOptions): Promise<boolean> {
+		await this._proxy.$onSave(this.resource, this.viewType);
+		this.setDirty(false);
+		return true;
+	}
+
+	public async saveAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
+		if (!this._editable) {
+			// Since the editor is readonly, just copy the file over
+			await this._fileService.copy(resource, targetResource, false /* overwrite */);
+			return true;
+		}
+
+		await this._proxy.$onSaveAs(this.resource, this.viewType, targetResource);
+		this.setDirty(false);
+		return true;
+	}
+
+	public async backup(): Promise<IWorkingCopyBackup> {
+		if (this._hotExitState.type === HotExitState.Type.Pending) {
+			this._hotExitState.operation.cancel();
+		}
+
+		const pendingState = new HotExitState.Pending(
+			createCancelablePromise(token =>
+				this._proxy.$backup(this.resource.toJSON(), this.viewType, token)));
+		this._hotExitState = pendingState;
+
+		try {
+			await pendingState.operation;
+			// Make sure state has not changed in the meantime
+			if (this._hotExitState === pendingState) {
+				this._hotExitState = HotExitState.Allowed;
+			}
+		} catch (e) {
+			// Make sure state has not changed in the meantime
+			if (this._hotExitState === pendingState) {
+				this._hotExitState = HotExitState.NotAllowed;
+			}
+		}
+
+		if (this._hotExitState === HotExitState.Allowed) {
+			return {
+				meta: {
+					viewType: this.viewType,
+				}
+			};
+		}
+
+		throw new Error('Cannot back up in this state');
+	}
+}

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

@@ -657,7 +657,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		}
 
 		const revivedResource = URI.revive(resource);
-		const document = this.getDocument(viewType, revivedResource);
+		const document = this.getCustomDocument(viewType, revivedResource);
 		this._documents.delete(document);
 		document.dispose();
 	}
@@ -684,12 +684,11 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		switch (entry.type) {
 			case WebviewEditorType.Custom:
 				{
-					const document = this.getDocument(viewType, revivedResource);
+					const document = this.getCustomDocument(viewType, revivedResource);
 					return entry.provider.resolveCustomEditor(document, revivedPanel);
 				}
 			case WebviewEditorType.Text:
 				{
-					await this._extHostDocuments.ensureDocumentData(revivedResource);
 					const document = this._extHostDocuments.getDocument(revivedResource);
 					return entry.provider.resolveCustomTextEditor(document, revivedPanel);
 				}
@@ -701,32 +700,32 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	}
 
 	async $undo(resourceComponents: UriComponents, viewType: string): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		document._undo();
 	}
 
 	async $redo(resourceComponents: UriComponents, viewType: string): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		document._redo();
 	}
 
 	async $revert(resourceComponents: UriComponents, viewType: string): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		document._revert();
 	}
 
 	async $onSave(resourceComponents: UriComponents, viewType: string): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		document._save();
 	}
 
 	async $onSaveAs(resourceComponents: UriComponents, viewType: string, targetResource: UriComponents): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		return document._saveAs(URI.revive(targetResource));
 	}
 
 	async $backup(resourceComponents: UriComponents, viewType: string, cancellation: CancellationToken): Promise<void> {
-		const document = this.getDocument(viewType, resourceComponents);
+		const document = this.getCustomDocument(viewType, resourceComponents);
 		return document._backup(cancellation);
 	}
 
@@ -734,7 +733,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		return this._webviewPanels.get(handle);
 	}
 
-	private getDocument(viewType: string, resource: UriComponents): CustomDocument {
+	private getCustomDocument(viewType: string, resource: UriComponents): CustomDocument {
 		const document = this._documents.get(viewType, URI.revive(resource));
 		if (!document) {
 			throw new Error('No webview editor custom document found');

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

@@ -20,22 +20,16 @@ import { IWebviewService, WebviewEditorOverlay } from 'vs/workbench/contrib/webv
 import { IWebviewWorkbenchService, LazilyResolvedWebviewEditorInput } from 'vs/workbench/contrib/webview/browser/webviewWorkbenchService';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { AutoSaveMode, IFilesConfigurationService } from 'vs/workbench/services/filesConfiguration/common/filesConfigurationService';
-import { ITextFileService } from 'vs/workbench/services/textfile/common/textfiles';
-
-export const enum ModelType {
-	Custom = 'custom',
-	Text = 'text',
-}
+import { IReference } from 'vs/base/common/lifecycle';
 
 export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 
-
 	public static typeId = 'workbench.editors.webviewEditor';
 
 	private readonly _editorResource: URI;
 	get resource() { return this._editorResource; }
 
-	private _model?: { readonly type: ModelType.Custom, readonly model: ICustomEditorModel } | { readonly type: ModelType.Text };
+	private _modelRef?: IReference<ICustomEditorModel>;
 
 	constructor(
 		resource: URI,
@@ -50,15 +44,11 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 		@IFileDialogService private readonly fileDialogService: IFileDialogService,
 		@IFilesConfigurationService private readonly filesConfigurationService: IFilesConfigurationService,
 		@IEditorService private readonly editorService: IEditorService,
-		@ITextFileService private readonly textFileService: ITextFileService,
-
 	) {
 		super(id, viewType, '', webview, webviewService, webviewWorkbenchService);
 		this._editorResource = resource;
 	}
 
-	public modelType?: ModelType;
-
 	public getTypeId(): string {
 		return CustomEditorInput.typeId;
 	}
@@ -110,20 +100,10 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 	}
 
 	public isDirty(): boolean {
-		if (!this._model) {
+		if (!this._modelRef) {
 			return false;
 		}
-
-		switch (this._model.type) {
-			case ModelType.Text:
-				return this.textFileService.isDirty(this.resource);
-
-			case ModelType.Custom:
-				return this._model.model.isDirty();
-
-			default:
-				throw new Error('Unknown model type');
-		}
+		return this._modelRef.object.isDirty();
 	}
 
 	public isSaving(): boolean {
@@ -139,28 +119,16 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 	}
 
 	public async save(groupId: GroupIdentifier, options?: ISaveOptions): Promise<IEditorInput | undefined> {
-		if (!this._model) {
+		if (!this._modelRef) {
 			return undefined;
 		}
 
-		switch (this._model.type) {
-			case ModelType.Text:
-				{
-					const result = await this.textFileService.save(this.resource, options);
-					return result ? this : undefined;
-				}
-			case ModelType.Custom:
-				{
-					const result = await this._model.model.save(options);
-					return result ? this : undefined;
-				}
-			default:
-				throw new Error('Unknown model type');
-		}
+		const result = await this._modelRef.object.save(options);
+		return result ? this : undefined;
 	}
 
 	public async saveAs(groupId: GroupIdentifier, options?: ISaveOptions): Promise<IEditorInput | undefined> {
-		if (!this._model) {
+		if (!this._modelRef) {
 			return undefined;
 		}
 
@@ -170,66 +138,25 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 			return undefined; // save cancelled
 		}
 
-		switch (this._model.type) {
-			case ModelType.Text:
-				if (!await this.textFileService.saveAs(this.resource, target, options)) {
-					return undefined;
-				}
-				break;
-
-			case ModelType.Custom:
-				if (!await this._model.model.saveAs(this._editorResource, target, options)) {
-					return undefined;
-				}
-				break;
-
-			default:
-				throw new Error('Unknown model type');
+		if (!await this._modelRef.object.saveAs(this._editorResource, target, options)) {
+			return undefined;
 		}
 
 		return this.handleMove(groupId, target) || this.editorService.createInput({ resource: target, forceFile: true });
 	}
 
 	public async revert(group: GroupIdentifier, options?: IRevertOptions): Promise<void> {
-		if (!this._model) {
-			return;
-		}
-
-		switch (this._model.type) {
-			case ModelType.Text:
-				return this.textFileService.revert(this.resource, options);
-
-			case ModelType.Custom:
-				return this._model.model.revert(options);
-
-			default:
-				throw new Error('Unknown model type');
-		}
+		return this._modelRef?.object.revert(options);
 	}
 
 	public async resolve(): Promise<IEditorModel> {
 		const editorModel = await super.resolve();
-		if (!this._model) {
-			switch (this.modelType) {
-				case ModelType.Custom:
-					const model = await this.customEditorService.models.resolve(this.resource, this.viewType);
-					this._model = { type: ModelType.Custom, model };
-					this._register(model.onDidChangeDirty(() => this._onDidChangeDirty.fire()));
-
-					break;
-
-				case ModelType.Text:
-					this._model = { type: ModelType.Text, };
-					this.textFileService.files.onDidChangeDirty(e => {
-						if (isEqual(this.resource, e.resource)) {
-							this._onDidChangeDirty.fire();
-						}
-					});
-
-					break;
-
-				default:
-					throw new Error('Unknown model type');
+		if (!this._modelRef) {
+			const modelRef = await this.customEditorService.models.tryRetain(this.resource, this.viewType);
+			if (modelRef) {
+				this._modelRef = modelRef;
+				this._register(this._modelRef);
+				this._register(this._modelRef.object.onDidChangeDirty(() => this._onDidChangeDirty.fire()));
 			}
 		}
 
@@ -256,40 +183,10 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 	}
 
 	public undo(): void {
-		if (!this._model) {
-			return;
-		}
-
-		switch (this._model.type) {
-			case ModelType.Custom:
-				this._model.model.undo();
-				return;
-
-			case ModelType.Text:
-				this.textFileService.files.get(this.resource)?.textEditorModel?.undo();
-				return;
-
-			default:
-				throw new Error('Unknown model type');
-		}
+		this._modelRef?.object.undo();
 	}
 
 	public redo(): void {
-		if (!this._model) {
-			return;
-		}
-
-		switch (this._model.type) {
-			case ModelType.Custom:
-				this._model.model.redo();
-				return;
-
-			case ModelType.Text:
-				this.textFileService.files.get(this.resource)?.textEditorModel?.redo();
-				return;
-
-			default:
-				throw new Error('Unknown model type');
-		}
+		this._modelRef?.object.redo();
 	}
 }

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

@@ -26,7 +26,6 @@ export class CustomEditorInputFactory extends WebviewEditorInputFactory {
 		const data = {
 			...this.toJson(input),
 			editorResource: input.resource.toJSON(),
-			modelType: input.modelType
 		};
 
 		try {
@@ -55,9 +54,6 @@ export class CustomEditorInputFactory extends WebviewEditorInputFactory {
 		if (typeof data.group === 'number') {
 			customInput.updateGroup(data.group);
 		}
-		if ((data as any).modelType) {
-			customInput.modelType = (data as any).modelType;
-		}
 		return customInput;
 	}
 }

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

@@ -16,7 +16,6 @@ import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/c
 import { IEditorOptions, ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { FileOperation, IFileService } from 'vs/platform/files/common/files';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
-import { ILabelService } from 'vs/platform/label/common/label';
 import { IQuickInputService, IQuickPickItem } from 'vs/platform/quickinput/common/quickInput';
 import * as colorRegistry from 'vs/platform/theme/common/colorRegistry';
 import { registerThemingParticipant } from 'vs/platform/theme/common/themeService';
@@ -25,12 +24,11 @@ import { IWorkbenchContribution } from 'vs/workbench/common/contributions';
 import { EditorInput, EditorOptions, IEditor, IEditorInput } from 'vs/workbench/common/editor';
 import { DiffEditorInput } from 'vs/workbench/common/editor/diffEditorInput';
 import { webviewEditorsExtensionPoint } from 'vs/workbench/contrib/customEditor/browser/extensionPoint';
-import { CONTEXT_FOCUSED_CUSTOM_EDITOR_IS_EDITABLE, CONTEXT_CUSTOM_EDITORS, CustomEditorInfo, CustomEditorInfoCollection, CustomEditorPriority, CustomEditorSelector, ICustomEditor, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
+import { CONTEXT_CUSTOM_EDITORS, CONTEXT_FOCUSED_CUSTOM_EDITOR_IS_EDITABLE, CustomEditorInfo, CustomEditorInfoCollection, CustomEditorPriority, CustomEditorSelector, ICustomEditor, ICustomEditorService } from 'vs/workbench/contrib/customEditor/common/customEditor';
 import { CustomEditorModelManager } from 'vs/workbench/contrib/customEditor/common/customEditorModelManager';
 import { IWebviewService, webviewHasOwnEditFunctionsContext } from 'vs/workbench/contrib/webview/browser/webview';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IEditorService, IOpenEditorOverride } from 'vs/workbench/services/editor/common/editorService';
-import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
 import { CustomEditorInput } from './customEditorInput';
 
 export const defaultEditorId = 'default';
@@ -104,7 +102,6 @@ export class CustomEditorService extends Disposable implements ICustomEditorServ
 
 	constructor(
 		@IContextKeyService contextKeyService: IContextKeyService,
-		@IWorkingCopyService workingCopyService: IWorkingCopyService,
 		@IFileService fileService: IFileService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IEditorService private readonly editorService: IEditorService,
@@ -112,11 +109,10 @@ export class CustomEditorService extends Disposable implements ICustomEditorServ
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IQuickInputService private readonly quickInputService: IQuickInputService,
 		@IWebviewService private readonly webviewService: IWebviewService,
-		@ILabelService labelService: ILabelService
 	) {
 		super();
 
-		this._models = new CustomEditorModelManager(workingCopyService, labelService);
+		this._models = new CustomEditorModelManager();
 
 		this._customEditorContextKey = CONTEXT_CUSTOM_EDITORS.bindTo(contextKeyService);
 		this._focusedCustomEditorIsEditable = CONTEXT_FOCUSED_CUSTOM_EDITOR_IS_EDITABLE.bindTo(contextKeyService);

src/vs/workbench/contrib/customEditor/common/customEditor.ts

@@ -4,7 +4,6 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { distinct, mergeSort } from 'vs/base/common/arrays';
-import { CancelablePromise } from 'vs/base/common/async';
 import { Event } from 'vs/base/common/event';
 import * as glob from 'vs/base/common/glob';
 import { basename } from 'vs/base/common/resources';
@@ -12,9 +11,9 @@ import { URI } from 'vs/base/common/uri';
 import { RawContextKey } from 'vs/platform/contextkey/common/contextkey';
 import { ITextEditorOptions } from 'vs/platform/editor/common/editor';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { IEditor, IRevertOptions, ISaveOptions, IEditorInput } from 'vs/workbench/common/editor';
+import { IEditor, IEditorInput, IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
 import { IEditorGroup } from 'vs/workbench/services/editor/common/editorGroupsService';
-import { IWorkingCopy } from 'vs/workbench/services/workingCopy/common/workingCopyService';
+import { IDisposable, IReference } from 'vs/base/common/lifecycle';
 
 export const ICustomEditorService = createDecorator<ICustomEditorService>('customEditorService');
 
@@ -44,39 +43,22 @@ export interface ICustomEditorService {
 }
 
 export interface ICustomEditorModelManager {
-	get(resource: URI, viewType: string): ICustomEditorModel | undefined;
+	get(resource: URI, viewType: string): Promise<ICustomEditorModel | undefined>;
 
-	resolve(resource: URI, viewType: string): Promise<ICustomEditorModel>;
+	tryRetain(resource: URI, viewType: string): Promise<IReference<ICustomEditorModel>> | undefined;
 
-	disposeModel(model: ICustomEditorModel): void;
+	add(resource: URI, viewType: string, model: Promise<ICustomEditorModel>): Promise<IReference<ICustomEditorModel>>;
 
 	disposeAllModelsForView(viewType: string): void;
 }
 
-export interface CustomEditorSaveEvent {
-	readonly resource: URI;
-	readonly waitUntil: (until: Promise<any>) => void;
-}
-
-export interface CustomEditorSaveAsEvent {
-	readonly resource: URI;
-	readonly targetResource: URI;
-	readonly waitUntil: (until: Promise<any>) => void;
-}
-
-export interface ICustomEditorModel extends IWorkingCopy {
+export interface ICustomEditorModel extends IDisposable {
 	readonly viewType: string;
+	readonly resource: URI;
 
-	readonly onUndo: Event<void>;
-	readonly onRedo: Event<void>;
-	readonly onRevert: Event<void>;
-
-	readonly onWillSave: Event<CustomEditorSaveEvent>;
-	readonly onWillSaveAs: Event<CustomEditorSaveAsEvent>;
-
-	onBackup(f: () => CancelablePromise<void>): void;
+	isDirty(): boolean;
+	readonly onDidChangeDirty: Event<void>;
 
-	setDirty(dirty: boolean): void;
 	undo(): void;
 	redo(): void;
 	revert(options?: IRevertOptions): Promise<void>;

src/vs/workbench/contrib/customEditor/common/customEditorModel.ts

-/*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the MIT License. See License.txt in the project root for license information.
- *--------------------------------------------------------------------------------------------*/
-
-import { CancelablePromise } from 'vs/base/common/async';
-import { Emitter, Event } from 'vs/base/common/event';
-import { Disposable } from 'vs/base/common/lifecycle';
-import { URI } from 'vs/base/common/uri';
-import { IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
-import { CustomEditorSaveAsEvent, CustomEditorSaveEvent, ICustomEditorModel } from 'vs/workbench/contrib/customEditor/common/customEditor';
-import { IWorkingCopyBackup, WorkingCopyCapabilities } from 'vs/workbench/services/workingCopy/common/workingCopyService';
-import { ILabelService } from 'vs/platform/label/common/label';
-import { basename } from 'vs/base/common/path';
-
-namespace HotExitState {
-	export const enum Type {
-		NotSupported,
-		Allowed,
-		NotAllowed,
-		Pending,
-	}
-
-	export const NotSupported = Object.freeze({ type: Type.NotSupported } as const);
-	export const Allowed = Object.freeze({ type: Type.Allowed } as const);
-	export const NotAllowed = Object.freeze({ type: Type.NotAllowed } as const);
-
-	export class Pending {
-		readonly type = Type.Pending;
-
-		constructor(
-			public readonly operation: CancelablePromise<void>,
-		) { }
-	}
-
-	export type State = typeof NotSupported | typeof Allowed | typeof NotAllowed | Pending;
-}
-
-export class CustomEditorModel extends Disposable implements ICustomEditorModel {
-
-	private _hotExitState: HotExitState.State = HotExitState.NotSupported;
-	private _dirty = false;
-
-	constructor(
-		public readonly viewType: string,
-		private readonly _resource: URI,
-		private readonly labelService: ILabelService,
-	) {
-		super();
-	}
-
-	//#region IWorkingCopy
-
-	public get resource() {
-		return this._resource;
-	}
-
-	public get name() {
-		return basename(this.labelService.getUriLabel(this._resource));
-	}
-
-	public get capabilities(): WorkingCopyCapabilities {
-		return 0;
-	}
-
-	public isDirty(): boolean {
-		return this._dirty;
-	}
-
-	private readonly _onDidChangeDirty: Emitter<void> = this._register(new Emitter<void>());
-	readonly onDidChangeDirty: Event<void> = this._onDidChangeDirty.event;
-
-	private readonly _onDidChangeContent: Emitter<void> = this._register(new Emitter<void>());
-	readonly onDidChangeContent: Event<void> = this._onDidChangeContent.event;
-
-	//#endregion
-
-	private readonly _onUndo = this._register(new Emitter<void>());
-	public readonly onUndo = this._onUndo.event;
-
-	private readonly _onRedo = this._register(new Emitter<void>());
-	public readonly onRedo = this._onRedo.event;
-
-	private readonly _onRevert = this._register(new Emitter<void>());
-	public readonly onRevert = this._onRevert.event;
-
-	private readonly _onWillSave = this._register(new Emitter<CustomEditorSaveEvent>());
-	public readonly onWillSave = this._onWillSave.event;
-
-	private readonly _onWillSaveAs = this._register(new Emitter<CustomEditorSaveAsEvent>());
-	public readonly onWillSaveAs = this._onWillSaveAs.event;
-
-	private _onBackup: undefined | (() => CancelablePromise<void>);
-
-	public onBackup(f: () => CancelablePromise<void>) {
-		if (this._onBackup) {
-			throw new Error('Backup already implemented');
-		}
-		this._onBackup = f;
-
-		if (this._hotExitState === HotExitState.NotSupported) {
-			this._hotExitState = this.isDirty() ? HotExitState.NotAllowed : HotExitState.Allowed;
-		}
-	}
-
-	public setDirty(dirty: boolean): void {
-		this._onDidChangeContent.fire();
-
-		if (this._dirty !== dirty) {
-			this._dirty = dirty;
-			this._onDidChangeDirty.fire();
-		}
-	}
-
-	public async revert(_options?: IRevertOptions) {
-		if (this._dirty) {
-			this._onRevert.fire();
-		}
-	}
-
-	public undo() {
-		this._onUndo.fire();
-	}
-
-	public redo() {
-		this._onRedo.fire();
-	}
-
-	public async save(_options?: ISaveOptions): Promise<boolean> {
-		const untils: Promise<any>[] = [];
-		const handler: CustomEditorSaveEvent = {
-			resource: this._resource,
-			waitUntil: (until: Promise<any>) => untils.push(until)
-		};
-
-		try {
-			this._onWillSave.fire(handler);
-			await Promise.all(untils);
-		} catch {
-			return false;
-		}
-
-		this.setDirty(false);
-
-		return true;
-	}
-
-	public async saveAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
-		const untils: Promise<any>[] = [];
-		const handler: CustomEditorSaveAsEvent = {
-			resource,
-			targetResource,
-			waitUntil: (until: Promise<any>) => untils.push(until)
-		};
-
-		try {
-			this._onWillSaveAs.fire(handler);
-			await Promise.all(untils);
-		} catch {
-			return false;
-		}
-
-		this.setDirty(false);
-
-		return true;
-	}
-
-	public async backup(): Promise<IWorkingCopyBackup> {
-		if (this._hotExitState === HotExitState.NotSupported) {
-			throw new Error('Not supported');
-		}
-
-		if (this._hotExitState.type === HotExitState.Type.Pending) {
-			this._hotExitState.operation.cancel();
-		}
-		this._hotExitState = HotExitState.NotAllowed;
-
-		const pendingState = new HotExitState.Pending(this._onBackup!());
-		this._hotExitState = pendingState;
-
-		try {
-			await pendingState.operation;
-			// Make sure state has not changed in the meantime
-			if (this._hotExitState === pendingState) {
-				this._hotExitState = HotExitState.Allowed;
-			}
-		} catch (e) {
-			// Make sure state has not changed in the meantime
-			if (this._hotExitState === pendingState) {
-				this._hotExitState = HotExitState.NotAllowed;
-			}
-		}
-
-		if (this._hotExitState === HotExitState.Allowed) {
-			return {
-				meta: {
-					viewType: this.viewType,
-				}
-			};
-		}
-		throw new Error('Cannot back up in this state');
-	}
-}

src/vs/workbench/contrib/customEditor/common/customEditorModelManager.ts

@@ -3,58 +3,64 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { DisposableStore } from 'vs/base/common/lifecycle';
+import { IReference } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { ICustomEditorModel, ICustomEditorModelManager } from 'vs/workbench/contrib/customEditor/common/customEditor';
-import { CustomEditorModel } from 'vs/workbench/contrib/customEditor/common/customEditorModel';
-import { IWorkingCopyService } from 'vs/workbench/services/workingCopy/common/workingCopyService';
-import { ILabelService } from 'vs/platform/label/common/label';
+import { once } from 'vs/base/common/functional';
 
 export class CustomEditorModelManager implements ICustomEditorModelManager {
-	private readonly _models = new Map<string, { readonly model: CustomEditorModel, readonly disposables: DisposableStore }>();
 
-	constructor(
-		@IWorkingCopyService private readonly _workingCopyService: IWorkingCopyService,
-		@ILabelService private readonly _labelService: ILabelService
-	) { }
+	private readonly _references = new Map<string, {
+		readonly viewType: string,
+		readonly model: Promise<ICustomEditorModel>,
+		counter: number
+	}>();
 
-
-	public get(resource: URI, viewType: string): ICustomEditorModel | undefined {
-		return this._models.get(this.key(resource, viewType))?.model;
+	public async get(resource: URI, viewType: string): Promise<ICustomEditorModel | undefined> {
+		const key = this.key(resource, viewType);
+		const entry = this._references.get(key);
+		return entry?.model;
 	}
 
-	public async resolve(resource: URI, viewType: string): Promise<ICustomEditorModel> {
-		const existing = this.get(resource, viewType);
-		if (existing) {
-			return existing;
+	public tryRetain(resource: URI, viewType: string): Promise<IReference<ICustomEditorModel>> | undefined {
+		const key = this.key(resource, viewType);
+
+		const entry = this._references.get(key);
+		if (!entry) {
+			return undefined;
 		}
 
-		const model = new CustomEditorModel(viewType, resource, this._labelService);
-		const disposables = new DisposableStore();
-		disposables.add(this._workingCopyService.registerWorkingCopy(model));
-		this._models.set(this.key(resource, viewType), { model, disposables });
-		return model;
+		entry.counter++;
+
+		return entry.model.then(model => {
+			return {
+				object: model,
+				dispose: once(() => {
+					if (--entry!.counter <= 0) {
+						entry.model.then(x => x.dispose());
+						this._references.delete(key);
+					}
+				}),
+			};
+		});
 	}
 
-	public disposeModel(model: ICustomEditorModel): void {
-		let foundKey: string | undefined;
-		for (const [key, value] of this._models) {
-			if (model === value.model) {
-				value.disposables.dispose();
-				value.model.dispose();
-				foundKey = key;
-			}
-		}
-		if (typeof foundKey === 'string') {
-			this._models.delete(foundKey);
+	public add(resource: URI, viewType: string, model: Promise<ICustomEditorModel>): Promise<IReference<ICustomEditorModel>> {
+		const key = this.key(resource, viewType);
+		const existing = this._references.get(key);
+		if (existing) {
+			throw new Error('Model already exists');
 		}
-		return;
+
+		this._references.set(key, { viewType, model, counter: 0 });
+		return this.tryRetain(resource, viewType)!;
 	}
 
 	public disposeAllModelsForView(viewType: string): void {
-		for (const [, value] of this._models) {
-			if (value.model.viewType === viewType) {
-				this.disposeModel(value.model);
+		for (const [key, value] of this._references) {
+			if (value.viewType === viewType) {
+				value.model.then(x => x.dispose());
+				this._references.delete(key);
 			}
 		}
 	}

src/vs/workbench/contrib/customEditor/common/customTextEditorModel.ts

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Emitter, Event } from 'vs/base/common/event';
+import { Disposable, IReference } from 'vs/base/common/lifecycle';
+import { isEqual } from 'vs/base/common/resources';
+import { URI } from 'vs/base/common/uri';
+import { IResolvedTextEditorModel, ITextModelService } from 'vs/editor/common/services/resolverService';
+import { IRevertOptions, ISaveOptions } from 'vs/workbench/common/editor';
+import { ICustomEditorModel } from 'vs/workbench/contrib/customEditor/common/customEditor';
+import { ITextFileService } from 'vs/workbench/services/textfile/common/textfiles';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+
+export class CustomTextEditorModel extends Disposable implements ICustomEditorModel {
+
+	public static async create(
+		instantiationService: IInstantiationService,
+		viewType: string,
+		resource: URI
+	): Promise<CustomTextEditorModel> {
+		return instantiationService.invokeFunction(async accessor => {
+			const textModelResolverService = accessor.get(ITextModelService);
+			const textFileService = accessor.get(ITextFileService);
+			const model = await textModelResolverService.createModelReference(resource);
+			return new CustomTextEditorModel(viewType, resource, model, textFileService);
+		});
+	}
+
+	private constructor(
+		public readonly viewType: string,
+		private readonly _resource: URI,
+		model: IReference<IResolvedTextEditorModel>,
+		@ITextFileService private readonly textFileService: ITextFileService,
+	) {
+		super();
+
+		this._register(model);
+
+		this._register(this.textFileService.files.onDidChangeDirty(e => {
+			if (isEqual(this.resource, e.resource)) {
+				this._onDidChangeDirty.fire();
+				this._onDidChangeContent.fire();
+			}
+		}));
+	}
+
+	public get resource() {
+		return this._resource;
+	}
+
+	public isDirty(): boolean {
+		return this.textFileService.isDirty(this.resource);
+	}
+
+	private readonly _onDidChangeDirty: Emitter<void> = this._register(new Emitter<void>());
+	readonly onDidChangeDirty: Event<void> = this._onDidChangeDirty.event;
+
+	private readonly _onDidChangeContent: Emitter<void> = this._register(new Emitter<void>());
+	readonly onDidChangeContent: Event<void> = this._onDidChangeContent.event;
+
+	public async revert(options?: IRevertOptions) {
+		return this.textFileService.revert(this.resource, options);
+	}
+
+	public undo() {
+		this.textFileService.files.get(this.resource)?.textEditorModel?.undo();
+	}
+
+	public redo() {
+		this.textFileService.files.get(this.resource)?.textEditorModel?.redo();
+	}
+
+	public async save(options?: ISaveOptions): Promise<boolean> {
+		return !!await this.textFileService.save(this.resource, options);
+	}
+
+	public async saveAs(resource: URI, targetResource: URI, options?: ISaveOptions): Promise<boolean> {
+		return !!await this.textFileService.saveAs(resource, targetResource, options);
+	}
+}