Auto select kernel using Positron foreground session (#2505)

* auto select kernel using positron foreground session

Extract a notebook controller manager, which updates a notebook's
affinity for each controller on notebook open.

Fix another bug in the notebook editor widget, causing new untitled
notebooks to store the incorrect `selectedKernelId` and therefore auto
select the incorrect kernel for a notebook.

* logging

* imports

* address review comments

extensions/positron-notebook-controllers/src/extension.ts

@@ -1,33 +1,39 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2023 Posit Software, PBC. All rights reserved.
+ *  Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
 import * as positron from 'positron';
 import * as vscode from 'vscode';
-import { initializeLogging, trace } from './logging';
-import { NotebookController } from './notebookController';
+import { initializeLogging } from './logging';
+import { NotebookControllerManager } from './notebookControllerManager';
 
 export async function activate(context: vscode.ExtensionContext): Promise<void> {
 	initializeLogging();
 
-	const knownLanguageIds = new Set<string>();
-	// Ensure that a notebook controller is registered for a given language.
-	const ensureNotebookController = (languageId: string) => {
-		if (!knownLanguageIds.has(languageId)) {
-			const controller = new NotebookController(languageId);
-			knownLanguageIds.add(languageId);
-			context.subscriptions.push(controller);
-			trace(`Registered notebook controller for language: ${languageId}`);
-		}
-	};
+	const manager = new NotebookControllerManager();
+	context.subscriptions.push(manager);
 
 	// Register notebook controllers for newly registered runtimes.
 	context.subscriptions.push(positron.runtime.onDidRegisterRuntime((runtime) => {
-		ensureNotebookController(runtime.languageId);
+		if (!manager.controllers.has(runtime.languageId)) {
+			manager.createNotebookController(runtime.languageId);
+		}
 	}));
 
 	// Register notebook controllers for existing runtimes.
 	for (const runtime of await positron.runtime.getRegisteredRuntimes()) {
-		ensureNotebookController(runtime.languageId);
+		if (!manager.controllers.has(runtime.languageId)) {
+			manager.createNotebookController(runtime.languageId);
+		}
+	}
+
+	// Update notebook affinity when a notebook is opened.
+	context.subscriptions.push(vscode.workspace.onDidOpenNotebookDocument(async (notebook) => {
+		manager.updateNotebookAffinity(notebook);
+	}));
+
+	// Update notebook affinity for notebooks that are already opened.
+	for (const notebook of vscode.workspace.notebookDocuments) {
+		manager.updateNotebookAffinity(notebook);
 	}
 }

extensions/positron-notebook-controllers/src/notebookController.ts

@@ -13,21 +13,21 @@ import { PromiseHandles, delay, noop } from './util';
  */
 export class NotebookController implements vscode.Disposable {
 
-	private disposables: vscode.Disposable[] = [];
+	private readonly disposables: vscode.Disposable[] = [];
 
 	/** The wrapped VSCode notebook controller. */
-	private controller: vscode.NotebookController;
+	public readonly controller: vscode.NotebookController;
 
 	/** Deferred notebook runtime data objects keyed by notebook. */
-	private notebookRuntimes: Map<vscode.NotebookDocument, PromiseHandles<NotebookRuntimeData>> = new Map();
+	private readonly notebookRuntimes: Map<vscode.NotebookDocument, PromiseHandles<NotebookRuntimeData>> = new Map();
 
 	/** Incremented for each cell we create to give it a unique ID. */
 	private static CELL_COUNTER = 0;
 
 	/**
 	 * @param languageId The language ID for which this controller is responsible.
 	 */
-	constructor(private languageId: string) {
+	constructor(private readonly languageId: string) {
 		// Create a VSCode notebook controller for this language.
 		this.controller = vscode.notebooks.createNotebookController(
 			`positron-${languageId}`,
@@ -71,7 +71,7 @@ export class NotebookController implements vscode.Disposable {
 				// already selected.
 
 				// Start updating the notebook's language info as soon as possible, but only await it at the end.
-				const setNotebookLanguagePromise = setNotebookLanguage(e.notebook, this.languageId);
+				const setNotebookLanguagePromise = updateNotebookLanguage(e.notebook, this.languageId);
 
 				// Set the notebook's deferred runtime data. This needs to be set before any awaits.
 				// When a user executes code without a controller selected, they will be presented
@@ -291,8 +291,8 @@ export class NotebookController implements vscode.Disposable {
  * @param notebook Notebook whose language to set.
  * @param languageId The VSCode-compatible language ID compatible.
  * @returns Promise that resolves when the language has been set.
- * */
-async function setNotebookLanguage(notebook: vscode.NotebookDocument, languageId: string): Promise<void> {
+ */
+async function updateNotebookLanguage(notebook: vscode.NotebookDocument, languageId: string): Promise<void> {
 	// Set the language in the notebook's metadata.
 	// This follows the approach from the vscode-jupyter extension.
 	if (notebook.metadata?.custom?.metadata?.language_info?.name !== languageId) {

extensions/positron-notebook-controllers/src/notebookControllerManager.ts

@@ -0,0 +1,88 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import { trace } from './logging';
+import { NotebookController } from './notebookController';
+
+/**
+ * Manages notebook controllers.
+ */
+export class NotebookControllerManager implements vscode.Disposable {
+	/** Notebook controllers keyed by languageId. */
+	public readonly controllers = new Map<string, NotebookController>();
+
+	/**
+	 * Create a notebook controller for a language.
+	 *
+	 * @param languageId The language ID for which to create a notebook controller.
+	 */
+	public createNotebookController(languageId: string): void {
+		if (!this.controllers.has(languageId)) {
+			const controller = new NotebookController(languageId);
+			this.controllers.set(languageId, controller);
+			trace(`Registered notebook controller for language: ${languageId}`);
+		}
+	}
+
+	/**
+	 * Update a notebook's affinity for all known controllers.
+	 *
+	 * Positron automates certain decisions if a notebook only has a single 'preferred' controller.
+	 *
+	 * @param notebook The notebook whose affinity to update.
+	 * @returns Promise that resolves when the notebook's affinity has been updated across all controllers.
+	 */
+	public async updateNotebookAffinity(notebook: vscode.NotebookDocument): Promise<void> {
+		if (notebook.uri.scheme === 'untitled') {
+			// If the notebook is untitled, wait for its data to be updated. This works around the fact
+			// that `vscode.openNotebookDocument(notebookType, content)` first creates a notebook
+			// (triggering `onDidOpenNotebookDocument`), and later updates its content (triggering
+			// `onDidChangeNotebookDocument`).
+			const event = await new Promise<vscode.NotebookDocumentChangeEvent | undefined>((resolve) => {
+				// Apply a short timeout to avoid waiting indefinitely.
+				const timeout = setTimeout(() => {
+					disposable.dispose();
+					resolve(undefined);
+				}, 50);
+
+				const disposable = vscode.workspace.onDidChangeNotebookDocument((e) => {
+					if (e.notebook === notebook) {
+						clearTimeout(timeout);
+						disposable.dispose();
+						resolve(e);
+					}
+				});
+			});
+
+			if (event) {
+				notebook = event.notebook;
+			}
+		}
+
+		// Detect the notebook's language.
+		// First try the notebook metadata.
+		const metadata = notebook.metadata?.custom?.metadata;
+		const languageId = metadata?.language_info?.name
+			?? metadata?.kernelspec?.language
+			// Fall back to the first cell's language.
+			?? notebook.getCells()?.[0].document.languageId;
+
+		// Get the preferred controller for the language.
+		const preferredController = languageId && this.controllers.get(languageId);
+
+		// Set the affinity across all known controllers.
+		for (const controller of this.controllers.values()) {
+			const affinity = controller === preferredController
+				? vscode.NotebookControllerAffinity.Preferred
+				: vscode.NotebookControllerAffinity.Default;
+			controller.controller.updateNotebookAffinity(notebook, affinity);
+			trace(`Updated notebook affinity for language: ${languageId}, notebook: ${notebook.uri.path}, affinity: ${affinity}`);
+		}
+	}
+
+	dispose(): void {
+		this.controllers.forEach(c => c.dispose());
+	}
+}

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

@@ -1772,6 +1772,15 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 				// --- End Positron ---
 				this.notebookKernelService.selectKernelForNotebook(kernel, this.textModel);
 			}
+			// --- Start Positron ---
+		} else if (this.textModel) {
+			// If we still haven't selected a kernel, try to select the suggested one.
+			const matching = this.notebookKernelService.getMatchingKernel(this.textModel);
+			const kernel = matching.suggestions.length === 1 ? matching.suggestions[0] : undefined;
+			if (kernel) {
+				this.notebookKernelService.selectKernelForNotebook(kernel, this.textModel);
+			}
+			// --- End Positron ---
 		}
 	}
 
@@ -1819,7 +1828,15 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		}
 		state.contributionsState = contributionsState;
 		if (this.textModel?.uri.scheme === Schemas.untitled) {
-			state.selectedKernelId = this.activeKernel?.id;
+			// --- Start Positron ---
+			// Upstream, this set selectedKernelId using notebookKernelService.getSelectedOrSuggestedKernel,
+			// which caused a bug if there is a single kernel with 'preferred' affinity for the
+			// notebook. In that case, when closing an untitled notebook, after the kernel is deselected
+			// the suggested kernel would still be stored in the editor view state. The next untitled
+			// notebook would then restart that kernel even if a different kernel now has 'preferred'
+			// affinity.
+			state.selectedKernelId = this.textModel && this.notebookKernelService.getMatchingKernel(this.textModel).selected?.id;
+			// --- End Positron ---
 		}
 
 		return state;