Fixing save of untitled text based custom editors

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

@@ -752,18 +752,23 @@ class MainThreadCustomEditorModel extends Disposable implements ICustomEditorMod
 		});
 	}
 
-	public async save(_options?: ISaveOptions): Promise<boolean> {
+	public async save(options?: ISaveOptions): Promise<boolean> {
+		return !!await this.saveCustomEditor(options);
+	}
+
+	public async saveCustomEditor(_options?: ISaveOptions): Promise<URI | undefined> {
 		if (!this._editable) {
-			return false;
+			return undefined;
 		}
+		// TODO: handle save untitled case
 		await createCancelablePromise(token => this._proxy.$onSave(this._editorResource, this.viewType, token));
 		this.change(() => {
 			this._savePoint = this._currentEditIndex;
 		});
-		return true;
+		return this._editorResource;
 	}
 
-	public async saveAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
+	public async saveCustomEditorAs(resource: URI, targetResource: URI, _options?: ISaveOptions): Promise<boolean> {
 		if (this._editable) {
 			await this._proxy.$onSaveAs(this._editorResource, this.viewType, targetResource);
 			this.change(() => {

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

@@ -6,6 +6,7 @@
 import { memoize } from 'vs/base/common/decorators';
 import { Lazy } from 'vs/base/common/lazy';
 import { IReference } from 'vs/base/common/lifecycle';
+import { Schemas } from 'vs/base/common/network';
 import { basename } from 'vs/base/common/path';
 import { isEqual } from 'vs/base/common/resources';
 import { assertIsDefined } from 'vs/base/common/types';
@@ -99,6 +100,10 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 		return this._modelRef ? this._modelRef.object.isReadonly() : false;
 	}
 
+	public isUntitled(): boolean {
+		return this.resource.scheme === Schemas.untitled;
+	}
+
 	public isDirty(): boolean {
 		if (!this._modelRef) {
 			return false;
@@ -120,8 +125,16 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 
 	public async save(groupId: GroupIdentifier, options?: ISaveOptions): Promise<IEditorInput | undefined> {
 		const modelRef = assertIsDefined(this._modelRef);
-		const result = await modelRef.object.save(options);
-		return result ? this : undefined;
+		const target = await modelRef.object.saveCustomEditor(options);
+		if (!target) {
+			return undefined; // save cancelled
+		}
+
+		if (!isEqual(target, this.resource)) {
+			return this.customEditorService.createInput(target, this.viewType, groupId);
+		}
+
+		return this;
 	}
 
 	public async saveAs(groupId: GroupIdentifier, options?: ISaveOptions): Promise<IEditorInput | undefined> {
@@ -133,7 +146,7 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 			return undefined; // save cancelled
 		}
 
-		if (!await modelRef.object.saveAs(this._editorResource, target, options)) {
+		if (!await modelRef.object.saveCustomEditorAs(this._editorResource, target, options)) {
 			return undefined;
 		}
 

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

@@ -57,8 +57,8 @@ export interface ICustomEditorModel extends IDisposable {
 
 	revert(options?: IRevertOptions): Promise<void>;
 
-	save(options?: ISaveOptions): Promise<boolean>;
-	saveAs(resource: URI, targetResource: URI, currentOptions?: ISaveOptions): Promise<boolean>;
+	saveCustomEditor(options?: ISaveOptions): Promise<URI | undefined>;
+	saveCustomEditorAs(resource: URI, targetResource: URI, currentOptions?: ISaveOptions): Promise<boolean>;
 }
 
 export const enum CustomEditorPriority {

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

@@ -68,11 +68,11 @@ export class CustomTextEditorModel extends Disposable implements ICustomEditorMo
 		return this.textFileService.revert(this.resource, options);
 	}
 
-	public async save(options?: ISaveOptions): Promise<boolean> {
-		return !!await this.textFileService.save(this.resource, options);
+	public saveCustomEditor(options?: ISaveOptions): Promise<URI | undefined> {
+		return this.textFileService.save(this.resource, options);
 	}
 
-	public async saveAs(resource: URI, targetResource: URI, options?: ISaveOptions): Promise<boolean> {
+	public async saveCustomEditorAs(resource: URI, targetResource: URI, options?: ISaveOptions): Promise<boolean> {
 		return !!await this.textFileService.saveAs(resource, targetResource, options);
 	}
 }