Support for Notebook CodeAction Kind (#183457)

* nb kind support -- wip

* allow notebook codeactions around single cell edit check

* move notebook code action type out of editor

---------

Co-authored-by: rebornix <penn.lv@gmail.com>

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

@@ -1525,7 +1525,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			InteractiveSessionCopyKind: extHostTypes.InteractiveSessionCopyKind,
 			InteractiveEditorResponseFeedbackKind: extHostTypes.InteractiveEditorResponseFeedbackKind,
 			StackFrameFocus: extHostTypes.StackFrameFocus,
-			ThreadFocus: extHostTypes.ThreadFocus
+			ThreadFocus: extHostTypes.ThreadFocus,
+			NotebookCodeActionKind: extHostTypes.NotebookCodeActionKind
 		};
 	};
 }

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

@@ -1358,7 +1358,6 @@ export class CodeAction {
 	}
 }
 
-
 @es5ClassCompat
 export class CodeActionKind {
 	private static readonly sep = '.';
@@ -1373,6 +1372,7 @@ export class CodeActionKind {
 	public static Source: CodeActionKind;
 	public static SourceOrganizeImports: CodeActionKind;
 	public static SourceFixAll: CodeActionKind;
+	public static Notebook: CodeActionKind;
 
 	constructor(
 		public readonly value: string
@@ -1390,6 +1390,17 @@ export class CodeActionKind {
 		return this.value === other.value || other.value.startsWith(this.value + CodeActionKind.sep);
 	}
 }
+
+export class NotebookCodeActionKind extends CodeActionKind {
+	public static override Notebook: CodeActionKind;
+
+	constructor(
+		public override readonly value: string
+	) {
+		super(value);
+	}
+}
+
 CodeActionKind.Empty = new CodeActionKind('');
 CodeActionKind.QuickFix = CodeActionKind.Empty.append('quickfix');
 CodeActionKind.Refactor = CodeActionKind.Empty.append('refactor');
@@ -1400,6 +1411,7 @@ CodeActionKind.RefactorRewrite = CodeActionKind.Refactor.append('rewrite');
 CodeActionKind.Source = CodeActionKind.Empty.append('source');
 CodeActionKind.SourceOrganizeImports = CodeActionKind.Source.append('organizeImports');
 CodeActionKind.SourceFixAll = CodeActionKind.Source.append('fixAll');
+CodeActionKind.Notebook = CodeActionKind.Empty.append('notebook');
 
 @es5ClassCompat
 export class SelectionRange {

src/vs/workbench/contrib/notebook/browser/contrib/saveParticipants/saveParticipants.ts

@@ -30,6 +30,8 @@ import { CodeActionTriggerType, CodeActionProvider, IWorkspaceTextEdit } from 'v
 import { applyCodeAction, ApplyCodeActionReason, getCodeActions } from 'vs/editor/contrib/codeAction/browser/codeAction';
 import { isEqual } from 'vs/base/common/resources';
 
+const NotebookCodeAction = new CodeActionKind('notebook');
+
 
 class FormatOnSaveParticipant implements IStoredFileWorkingCopySaveParticipant {
 	constructor(
@@ -133,9 +135,9 @@ class CodeActionOnSaveParticipant implements IStoredFileWorkingCopySaveParticipa
 			return undefined;
 		}
 
-		const codeActionsOnSave = this.createCodeActionsOnSave(settingItems);
+		const codeActionsOnSave = this.createCodeActionsOnSave(settingItems).filter(x => !NotebookCodeAction.contains(x));
+		const notebookCodeActionsOnSave = this.createCodeActionsOnSave(settingItems).filter(x => NotebookCodeAction.contains(x));
 
-		// TODO: potentially modify to account for new `Notebook` code action kind
 		// prioritize `source.fixAll` code actions
 		if (!Array.isArray(setting)) {
 			codeActionsOnSave.sort((a, b) => {
@@ -152,6 +154,9 @@ class CodeActionOnSaveParticipant implements IStoredFileWorkingCopySaveParticipa
 			});
 		}
 
+
+
+
 		if (!codeActionsOnSave.length) {
 			return undefined;
 		}
@@ -162,9 +167,28 @@ class CodeActionOnSaveParticipant implements IStoredFileWorkingCopySaveParticipa
 				.filter(x => setting[x] === false)
 				.map(x => new CodeActionKind(x));
 
+		const nbDisposable = new DisposableStore();
 
-		progress.report({ message: localize('notebookSaveParticipants.codeActions', "Running code actions") });
+		// run notebook code actions
+		progress.report({ message: localize('notebookSaveParticipants.notebookCodeActions', "Running 'Notebook' code actions") });
+		try {
+			const cell = notebookModel.cells[0];
+			const ref = await this.textModelService.createModelReference(cell.uri);
+			nbDisposable.add(ref);
+
+			const textEditorModel = ref.object.textEditorModel;
+
+			await this.applyOnSaveActions(textEditorModel, notebookCodeActionsOnSave, excludedActions, progress, token);
+		} catch {
+			this.logService.error('Failed to apply notebook code action on save');
+		} finally {
+			progress.report({ increment: 100 });
+			nbDisposable.dispose();
+		}
+
+		// run cell level code actions
 		const disposable = new DisposableStore();
+		progress.report({ message: localize('notebookSaveParticipants.cellCodeActions', "Running code actions") });
 		try {
 			await Promise.all(notebookModel.cells.map(async cell => {
 				const ref = await this.textModelService.createModelReference(cell.uri);
@@ -224,14 +248,16 @@ class CodeActionOnSaveParticipant implements IStoredFileWorkingCopySaveParticipa
 				for (const action of actionsToRun.validActions) {
 					const codeActionEdits = action.action.edit?.edits;
 					let breakFlag = false;
-					for (const edit of codeActionEdits ?? []) {
-						const workspaceTextEdit = edit as IWorkspaceTextEdit;
-						if (workspaceTextEdit.resource && isEqual(workspaceTextEdit.resource, model.uri)) {
-							continue;
-						} else {
-							// error -> applied to multiple resources
-							breakFlag = true;
-							break;
+					if (!action.action.kind?.includes('notebook')) {
+						for (const edit of codeActionEdits ?? []) {
+							const workspaceTextEdit = edit as IWorkspaceTextEdit;
+							if (workspaceTextEdit.resource && isEqual(workspaceTextEdit.resource, model.uri)) {
+								continue;
+							} else {
+								// error -> applied to multiple resources
+								breakFlag = true;
+								break;
+							}
 						}
 					}
 					if (breakFlag) {

src/vs/workbench/services/extensions/common/extensionsApiProposals.ts

@@ -55,6 +55,7 @@ export const allApiProposals = Object.freeze({
 	interactiveWindow: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.interactiveWindow.d.ts',
 	ipc: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.ipc.d.ts',
 	notebookCellExecutionState: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.notebookCellExecutionState.d.ts',
+	notebookCodeActions: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.notebookCodeActions.d.ts',
 	notebookControllerAffinityHidden: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.notebookControllerAffinityHidden.d.ts',
 	notebookDeprecated: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.notebookDeprecated.d.ts',
 	notebookExecution: 'https://raw.githubusercontent.com/microsoft/vscode/main/src/vscode-dts/vscode.proposed.notebookExecution.d.ts',

src/vscode-dts/vscode.proposed.notebookCodeActions.d.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+declare module 'vscode' {
+
+	// https://github.com/microsoft/vscode/issues/179213
+
+	export class NotebookCodeActionKind {
+		// can only return MULTI CELL workspaceEdits
+		// ex: notebook.organizeImprots
+		static readonly Notebook: CodeActionKind;
+
+		constructor(value: string);
+	}
+}