Allow invoking kernel picker for inactive notebook given `NotebookEditor` in...

Allow invoking kernel picker for inactive notebook given `NotebookEditor` in `notebook.selectKernel` command args (#132879)

* Allow invoking kernel picker for inactive notebook

`notebook.selectKernel` is now an API command that takes target `NotebookEditor` as an arg

* move select kernel command to extHost Kernel.
Co-authored-by: Nrebornix <penn.lv@gmail.com>

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

@@ -151,7 +151,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 	const extHostNotebook = rpcProtocol.set(ExtHostContext.ExtHostNotebook, new ExtHostNotebookController(rpcProtocol, extHostCommands, extHostDocumentsAndEditors, extHostDocuments, extensionStoragePaths));
 	const extHostNotebookDocuments = rpcProtocol.set(ExtHostContext.ExtHostNotebookDocuments, new ExtHostNotebookDocuments(extHostLogService, extHostNotebook));
 	const extHostNotebookEditors = rpcProtocol.set(ExtHostContext.ExtHostNotebookEditors, new ExtHostNotebookEditors(extHostLogService, rpcProtocol, extHostNotebook));
-	const extHostNotebookKernels = rpcProtocol.set(ExtHostContext.ExtHostNotebookKernels, new ExtHostNotebookKernels(rpcProtocol, initData, extHostNotebook, extHostLogService));
+	const extHostNotebookKernels = rpcProtocol.set(ExtHostContext.ExtHostNotebookKernels, new ExtHostNotebookKernels(rpcProtocol, initData, extHostNotebook, extHostCommands, extHostLogService));
 	const extHostNotebookRenderers = rpcProtocol.set(ExtHostContext.ExtHostNotebookRenderers, new ExtHostNotebookRenderers(rpcProtocol, extHostNotebook));
 	const extHostEditors = rpcProtocol.set(ExtHostContext.ExtHostEditors, new ExtHostEditors(rpcProtocol, extHostDocumentsAndEditors));
 	const extHostTreeViews = rpcProtocol.set(ExtHostContext.ExtHostTreeViews, new ExtHostTreeViews(rpcProtocol.getProxy(MainContext.MainThreadTreeViews), extHostCommands, extHostLogService));

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

@@ -17,7 +17,7 @@ export class ExtHostInteractive implements ExtHostInteractiveShape {
 		private _textDocumentsAndEditors: ExtHostDocumentsAndEditors,
 		private _commands: ExtHostCommands
 	) {
-		const apiCommand = new ApiCommand(
+		const openApiCommand = new ApiCommand(
 			'interactive.open',
 			'_interactive.open',
 			'Open interactive window and return notebook editor and input URI',
@@ -35,7 +35,7 @@ export class ExtHostInteractive implements ExtHostInteractiveShape {
 				return { notebookUri: URI.revive(v.notebookUri), inputUri: URI.revive(v.inputUri) };
 			})
 		);
-		this._commands.registerApiCommand(apiCommand);
+		this._commands.registerApiCommand(openApiCommand);
 	}
 
 	$willAddInteractiveDocument(uri: UriComponents, eol: string, modeId: string, notebookUri: UriComponents) {

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

@@ -13,6 +13,7 @@ import { URI, UriComponents } from 'vs/base/common/uri';
 import { ExtensionIdentifier, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ExtHostNotebookKernelsShape, ICellExecuteUpdateDto, IMainContext, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape, NotebookOutputDto } from 'vs/workbench/api/common/extHost.protocol';
+import { ApiCommand, ApiCommandArgument, ApiCommandResult, ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
 import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitDataService';
 import { ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
 import { ExtHostCell, ExtHostNotebookDocument } from 'vs/workbench/api/common/extHostNotebookDocument';
@@ -32,6 +33,11 @@ interface IKernelData {
 	associatedNotebooks: ResourceMap<boolean>;
 }
 
+type ExtHostSelectKernelArgs = ControllerInfo | { notebookEditor: vscode.NotebookEditor } | ControllerInfo & { notebookEditor: vscode.NotebookEditor } | undefined;
+export type SelectKernelReturnArgs = ControllerInfo | { notebookEditorId: string } | ControllerInfo & { notebookEditorId: string } | undefined;
+type ControllerInfo = { id: string, extension: string };
+
+
 export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 
 	private readonly _proxy: MainThreadNotebookKernelsShape;
@@ -44,9 +50,35 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		mainContext: IMainContext,
 		private readonly _initData: IExtHostInitDataService,
 		private readonly _extHostNotebook: ExtHostNotebookController,
+		private _commands: ExtHostCommands,
 		@ILogService private readonly _logService: ILogService,
 	) {
 		this._proxy = mainContext.getProxy(MainContext.MainThreadNotebookKernels);
+
+		// todo@rebornix @joyceerhl: move to APICommands once stablized.
+		const selectKernelApiCommand = new ApiCommand(
+			'notebook.selectKernel',
+			'_notebook.selectKernel',
+			'Trigger kernel picker for specified notebook editor widget',
+			[
+				new ApiCommandArgument<ExtHostSelectKernelArgs, SelectKernelReturnArgs>('options', 'Select kernel options', v => true, (v: ExtHostSelectKernelArgs) => {
+					if (v && 'notebookEditor' in v && 'id' in v) {
+						const notebookEditorId = this._extHostNotebook.getIdByEditor(v.notebookEditor);
+						return {
+							id: v.id, extension: v.extension, notebookEditorId
+						};
+					} else if (v && 'notebookEditor' in v) {
+						const notebookEditorId = this._extHostNotebook.getIdByEditor(v.notebookEditor);
+						if (notebookEditorId === undefined) {
+							throw new Error(`Cannot invoke 'notebook.selectKernel' for unrecognized notebook editor ${v.notebookEditor.document.uri.toString()}`);
+						}
+						return { notebookEditorId };
+					}
+					return v;
+				})
+			],
+			ApiCommandResult.Void);
+		this._commands.registerApiCommand(selectKernelApiCommand);
 	}
 
 	createNotebookController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable<void>, preloads?: vscode.NotebookRendererScript[]): vscode.NotebookController {

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

@@ -16,6 +16,7 @@ import { ILogService } from 'vs/platform/log/common/log';
 import { IQuickInputButton, IQuickInputService, IQuickPickItem } from 'vs/platform/quickinput/common/quickInput';
 import { Registry } from 'vs/platform/registry/common/platform';
 import { ThemeIcon } from 'vs/platform/theme/common/themeService';
+import type { SelectKernelReturnArgs } from 'vs/workbench/api/common/extHostNotebookKernels';
 import { Extensions as WorkbenchExtensions, IWorkbenchContribution, IWorkbenchContributionsRegistry } from 'vs/workbench/common/contributions';
 import { IExtensionsViewPaneContainer, VIEWLET_ID as EXTENSION_VIEWLET_ID } from 'vs/workbench/contrib/extensions/common/extensions';
 import { NOTEBOOK_ACTIONS_CATEGORY, SELECT_KERNEL_ID } from 'vs/workbench/contrib/notebook/browser/controller/coreActions';
@@ -33,10 +34,10 @@ import { IViewletService } from 'vs/workbench/services/viewlet/browser/viewlet';
 registerAction2(class extends Action2 {
 	constructor() {
 		super({
-			id: SELECT_KERNEL_ID,
+			id: '_notebook.selectKernel',
 			category: NOTEBOOK_ACTIONS_CATEGORY,
 			title: { value: nls.localize('notebookActions.selectKernel', "Select Notebook Kernel"), original: 'Select Notebook Kernel' },
-			precondition: NOTEBOOK_IS_ACTIVE_EDITOR,
+			// precondition: NOTEBOOK_IS_ACTIVE_EDITOR,
 			icon: selectKernelIcon,
 			f1: true,
 			menu: [{
@@ -77,6 +78,9 @@ registerAction2(class extends Action2 {
 								},
 								'extension': {
 									'type': 'string'
+								},
+								'notebookEditorId': {
+									'type': 'string'
 								}
 							}
 						}
@@ -86,7 +90,7 @@ registerAction2(class extends Action2 {
 		});
 	}
 
-	async run(accessor: ServicesAccessor, context?: { id: string, extension: string, ui?: boolean, notebookEditor?: NotebookEditorWidget }): Promise<boolean> {
+	async run(accessor: ServicesAccessor, context?: SelectKernelReturnArgs | { ui?: boolean, notebookEditor?: NotebookEditorWidget }): Promise<boolean> {
 		const notebookKernelService = accessor.get(INotebookKernelService);
 		const editorService = accessor.get(IEditorService);
 		const quickInputService = accessor.get(IQuickInputService);
@@ -94,27 +98,45 @@ registerAction2(class extends Action2 {
 		const logService = accessor.get(ILogService);
 		const viewletService = accessor.get(IViewletService);
 
-		const editor = context?.notebookEditor ?? getNotebookEditorFromEditorPane(editorService.activeEditorPane);
+		let editor;
+		if (context !== undefined) {
+			if ('notebookEditorId' in context) {
+				const editorId = context.notebookEditorId;
+				const matchingEditor = editorService.visibleEditorPanes.find((editorPane) => {
+					const notebookEditor = getNotebookEditorFromEditorPane(editorPane);
+					return notebookEditor?.getId() === editorId;
+				});
+				editor = getNotebookEditorFromEditorPane(matchingEditor);
+			} else if ('notebookEditor' in context) {
+				editor = context?.notebookEditor;
+			} else {
+				editor = getNotebookEditorFromEditorPane(editorService.activeEditorPane);
+			}
+		}
+
 		if (!editor || !editor.hasModel()) {
 			return false;
 		}
+		let controllerId = context && 'id' in context ? context.id : undefined;
+		let extensionId = context && 'extension' in context ? context.extension : undefined;
 
-		if (context && (typeof context.id !== 'string' || typeof context.extension !== 'string')) {
+		if (controllerId && (typeof controllerId !== 'string' || typeof extensionId !== 'string')) {
 			// validate context: id & extension MUST be strings
-			context = undefined;
+			controllerId = undefined;
+			extensionId = undefined;
 		}
 
 		const notebook = editor.textModel;
 		const { selected, all } = notebookKernelService.getMatchingKernel(notebook);
 
-		if (selected && context && selected.id === context.id && ExtensionIdentifier.equals(selected.extension, context.extension)) {
+		if (selected && controllerId && selected.id === controllerId && ExtensionIdentifier.equals(selected.extension, extensionId)) {
 			// current kernel is wanted kernel -> done
 			return true;
 		}
 
 		let newKernel: INotebookKernel | undefined;
-		if (context) {
-			const wantedId = `${context.extension}/${context.id}`;
+		if (controllerId) {
+			const wantedId = `${extensionId}/${controllerId}`;
 			for (let candidate of all) {
 				if (candidate.id === wantedId) {
 					newKernel = candidate;
@@ -323,7 +345,7 @@ export class KernelStatus extends Disposable implements IWorkbenchContribution {
 					tooltip: isSuggested ? nls.localize('tooltop', "{0} (suggestion)", tooltip) : tooltip,
 					command: SELECT_KERNEL_ID,
 				},
-				'notebook.selectKernel',
+				'_notebook.selectKernel',
 				StatusbarAlignment.RIGHT,
 				10
 			));
@@ -341,7 +363,7 @@ export class KernelStatus extends Disposable implements IWorkbenchContribution {
 					command: SELECT_KERNEL_ID,
 					backgroundColor: { id: 'statusBarItem.prominentBackground' }
 				},
-				'notebook.selectKernel',
+				'_notebook.selectKernel',
 				StatusbarAlignment.RIGHT,
 				10
 			));

src/vs/workbench/contrib/notebook/browser/controller/coreActions.ts

@@ -22,7 +22,7 @@ import { IJSONSchema } from 'vs/base/common/jsonSchema';
 import { MarshalledId } from 'vs/base/common/marshalling';
 
 // Kernel Command
-export const SELECT_KERNEL_ID = 'notebook.selectKernel';
+export const SELECT_KERNEL_ID = '_notebook.selectKernel';
 export const NOTEBOOK_ACTIONS_CATEGORY = { value: localize('notebookActions.category', "Notebook"), original: 'Notebook' };
 
 export const CELL_TITLE_CELL_GROUP_ID = 'inline/cell';

src/vs/workbench/test/browser/api/extHostNotebookKernel.test.ts

@@ -37,6 +37,7 @@ suite('NotebookKernel', function () {
 	let extHostDocuments: ExtHostDocuments;
 	let extHostNotebooks: ExtHostNotebookController;
 	let extHostNotebookDocuments: ExtHostNotebookDocuments;
+	let extHostCommands: ExtHostCommands;
 
 	const notebookUri = URI.parse('test:///notebook.file');
 	const kernelData = new Map<number, INotebookKernelDto2>();
@@ -84,7 +85,8 @@ suite('NotebookKernel', function () {
 				return URI.from({ scheme: 'test', path: generateUuid() });
 			}
 		};
-		extHostNotebooks = new ExtHostNotebookController(rpcProtocol, new ExtHostCommands(rpcProtocol, new NullLogService()), extHostDocumentsAndEditors, extHostDocuments, extHostStoragePaths);
+		extHostCommands = new ExtHostCommands(rpcProtocol, new NullLogService());
+		extHostNotebooks = new ExtHostNotebookController(rpcProtocol, extHostCommands, extHostDocumentsAndEditors, extHostDocuments, extHostStoragePaths);
 
 		extHostNotebookDocuments = new ExtHostNotebookDocuments(new NullLogService(), extHostNotebooks);
 
@@ -130,6 +132,7 @@ suite('NotebookKernel', function () {
 			rpcProtocol,
 			new class extends mock<IExtHostInitDataService>() { },
 			extHostNotebooks,
+			extHostCommands,
 			new NullLogService()
 		);
 	});