Merge pull request #286325 from microsoft/dev/dmitriv/custom-editors-bugs

Fix a set of issues around transfer and save of custom editor documents

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

@@ -176,6 +176,7 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 
 				const disposeSub = webviewInput.webview.onDidDispose(() => {
 					disposeSub.dispose();
+					inputDisposeSub.dispose();
 
 					// If the model is still dirty, make sure we have time to save it
 					if (modelRef.object.isDirty()) {
@@ -191,6 +192,14 @@ export class MainThreadCustomEditors extends Disposable implements extHostProtoc
 					modelRef.dispose();
 				});
 
+				// Also listen for when the input is disposed (e.g., during SaveAs when the webview is transferred to a new editor).
+				// In this case, webview.onDidDispose won't fire because the webview is reused.
+				const inputDisposeSub = webviewInput.onWillDispose(() => {
+					inputDisposeSub.dispose();
+					disposeSub.dispose();
+					modelRef.dispose();
+				});
+
 				if (capabilities.supportsMove) {
 					webviewInput.onMove(async (newResource: URI) => {
 						const oldModel = modelRef;
@@ -647,7 +656,9 @@ class MainThreadCustomEditorModel extends ResourceWorkingCopy implements ICustom
 			// TODO: handle cancellation
 			await createCancelablePromise(token => this._proxy.$onSaveAs(this._editorResource, this.viewType, targetResource, token));
 			this.change(() => {
+				this._isDirtyFromContentChange = false;
 				this._savePoint = this._currentEditIndex;
+				this._fromBackup = false;
 			});
 			return true;
 		} else {

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

@@ -104,8 +104,10 @@ class CustomDocumentStore {
 		return entry;
 	}
 
-	public delete(viewType: string, document: vscode.CustomDocument) {
-		const key = this.key(viewType, document.uri);
+	public delete(viewType: string, resource: vscode.Uri) {
+		// Use the resource parameter directly instead of document.uri, because the document's
+		// URI may have changed (e.g., after SaveAs from untitled to a file path).
+		const key = this.key(viewType, resource);
 		this._documents.delete(key);
 	}
 
@@ -242,7 +244,9 @@ export class ExtHostCustomEditors implements extHostProtocol.ExtHostCustomEditor
 
 		const revivedResource = URI.revive(resource);
 		const { document } = this.getCustomDocumentEntry(viewType, revivedResource);
-		this._documents.delete(viewType, document);
+		// Pass the resource we used to look up the document, not document.uri,
+		// because the document's URI may have changed (e.g., after SaveAs).
+		this._documents.delete(viewType, revivedResource);
 		document.dispose();
 	}
 

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

@@ -52,8 +52,15 @@ export class CustomEditorInput extends LazilyResolvedWebviewEditorInput {
 	): EditorInput {
 		return instantiationService.invokeFunction(accessor => {
 			// If it's an untitled file we must populate the untitledDocumentData
-			const untitledString = accessor.get(IUntitledTextEditorService).getValue(init.resource);
+			const untitledTextEditorService = accessor.get(IUntitledTextEditorService);
+			const untitledTextModel = untitledTextEditorService.get(init.resource);
+			const untitledString = untitledTextModel?.textEditorModel?.getValue();
 			const untitledDocumentData = untitledString ? VSBuffer.fromString(untitledString) : undefined;
+
+			// If we're taking over an untitled text editor, revert it so it's no longer
+			// tracked as a dirty working copy (fixes #125293).
+			untitledTextModel?.revert();
+
 			const webview = accessor.get(IWebviewService).createWebviewOverlay({
 				providedViewType: init.viewType,
 				title: init.webviewTitle,