Explicitly register for save and saveAs

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

@@ -280,8 +280,6 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 						this._proxy.$applyEdits(handle, editsToApply);
 					}
 				});
-				model.onWillSave(e => { e.waitUntil(this._proxy.$onSave(handle)); });
-				model.onWillSaveAs(e => { e.waitUntil(this._proxy.$onSaveAs(handle, e.resource.toJSON(), e.targetResource.toJSON())); });
 
 				webviewInput.onDisposeWebview(() => {
 					this._customEditorService.models.disposeModel(model);
@@ -315,6 +313,20 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		this._editorProviders.delete(viewType);
 	}
 
+	public async $registerCapabilities(handle: extHostProtocol.WebviewPanelHandle, capabilities: readonly extHostProtocol.WebviewEditorCapabilities[]): Promise<void> {
+		const webviewInput = this.getWebviewInput(handle);
+		const model = await this._customEditorService.models.loadOrCreate(webviewInput.getResource(), webviewInput.viewType);
+
+		const capabilitiesSet = new Set(capabilities);
+
+		if (capabilitiesSet.has(extHostProtocol.WebviewEditorCapabilities.Save)) {
+			model.onWillSave(e => { e.waitUntil(this._proxy.$onSave(handle)); });
+		}
+		if (capabilitiesSet.has(extHostProtocol.WebviewEditorCapabilities.SaveAs)) {
+			model.onWillSaveAs(e => { e.waitUntil(this._proxy.$onSaveAs(handle, e.resource.toJSON(), e.targetResource.toJSON())); });
+		}
+	}
+
 	public $onEdit(handle: extHostProtocol.WebviewPanelHandle, editData: any): void {
 		const webview = this.getWebviewInput(handle);
 		if (!(webview instanceof CustomFileEditorInput)) {

src/vs/workbench/api/common/extHost.protocol.ts

@@ -556,6 +556,11 @@ export interface WebviewExtensionDescription {
 	readonly location: UriComponents;
 }
 
+export enum WebviewEditorCapabilities {
+	Save,
+	SaveAs,
+}
+
 export interface MainThreadWebviewsShape extends IDisposable {
 	$createWebviewPanel(extension: WebviewExtensionDescription, handle: WebviewPanelHandle, viewType: string, title: string, showOptions: WebviewPanelShowOptions, options: modes.IWebviewPanelOptions & modes.IWebviewOptions): void;
 	$disposeWebview(handle: WebviewPanelHandle): void;
@@ -573,6 +578,7 @@ export interface MainThreadWebviewsShape extends IDisposable {
 
 	$registerEditorProvider(extension: WebviewExtensionDescription, viewType: string, options: modes.IWebviewPanelOptions): void;
 	$unregisterEditorProvider(viewType: string): void;
+	$registerCapabilities(handle: WebviewPanelHandle, capabilities: readonly WebviewEditorCapabilities[]): void;
 
 	$onEdit(handle: WebviewPanelHandle, editJson: any): void;
 }

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

@@ -16,7 +16,7 @@ import { IExtHostWorkspace } from 'vs/workbench/api/common/extHostWorkspace';
 import { EditorViewColumn } from 'vs/workbench/api/common/shared/editor';
 import { asWebviewUri, WebviewInitData } from 'vs/workbench/api/common/shared/webview';
 import * as vscode from 'vscode';
-import { ExtHostWebviewsShape, IMainContext, MainContext, MainThreadWebviewsShape, WebviewPanelHandle, WebviewPanelViewStateData } from './extHost.protocol';
+import { ExtHostWebviewsShape, IMainContext, MainContext, MainThreadWebviewsShape, WebviewPanelHandle, WebviewPanelViewStateData, WebviewEditorCapabilities } from './extHost.protocol';
 import { Disposable as VSCodeDisposable } from './extHostTypes';
 
 type IconPath = URI | { light: URI, dark: URI };
@@ -257,12 +257,11 @@ export class ExtHostWebviewEditor extends Disposable implements vscode.WebviewPa
 	}
 
 	async _onSave(): Promise<void> {
-		await assertIsDefined(this._capabilities).editingCapability?.save();
+		await assertIsDefined(this._capabilities?.editingCapability)?.save();
 	}
 
-
 	async _onSaveAs(resource: vscode.Uri, targetResource: vscode.Uri): Promise<void> {
-		await assertIsDefined(this._capabilities).editingCapability?.saveAs(resource, targetResource);
+		await assertIsDefined(this._capabilities?.editingCapability)?.saveAs(resource, targetResource);
 	}
 
 	private assertNotDisposed() {
@@ -450,6 +449,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		this._webviewPanels.set(handle, revivedPanel);
 		const capabilities = await provider.resolveWebviewEditor({ resource: URI.revive(input.resource) }, revivedPanel);
 		revivedPanel._setCapabilities(capabilities);
+		this.registerCapabilites(handle, capabilities);
 
 		// TODO: the first set of edits should likely be passed when resolving
 		if (input.edits.length) {
@@ -480,6 +480,17 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	private getWebviewPanel(handle: WebviewPanelHandle): ExtHostWebviewEditor | undefined {
 		return this._webviewPanels.get(handle);
 	}
+
+	private registerCapabilites(handle: WebviewPanelHandle, capabilities: vscode.WebviewEditorCapabilities) {
+		const declaredCapabilites: WebviewEditorCapabilities[] = [];
+		if (capabilities.editingCapability?.save) {
+			declaredCapabilites.push(WebviewEditorCapabilities.Save);
+		}
+		if (capabilities.editingCapability?.saveAs) {
+			declaredCapabilites.push(WebviewEditorCapabilities.SaveAs);
+		}
+		this._proxy.$registerCapabilities(handle, declaredCapabilites);
+	}
 }
 
 function convertWebviewOptions(