From ef657567eaf99a872df595892312f32632cab011 Mon Sep 17 00:00:00 2001 From: Wasim Lorgat Date: Mon, 25 Mar 2024 22:09:15 +0200 Subject: [PATCH] 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 --- .../src/extension.ts | 36 ++++---- .../src/notebookController.ts | 14 +-- .../src/notebookControllerManager.ts | 88 +++++++++++++++++++ .../notebook/browser/notebookEditorWidget.ts | 19 +++- 4 files changed, 134 insertions(+), 23 deletions(-) create mode 100644 extensions/positron-notebook-controllers/src/notebookControllerManager.ts diff --git a/extensions/positron-notebook-controllers/src/extension.ts b/extensions/positron-notebook-controllers/src/extension.ts index e8ba753197a..1c08cd948ca 100644 --- a/extensions/positron-notebook-controllers/src/extension.ts +++ b/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 { initializeLogging(); - const knownLanguageIds = new Set(); - // 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); } } diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts index ddb22af8c4d..5895dafea39 100644 --- a/extensions/positron-notebook-controllers/src/notebookController.ts +++ b/extensions/positron-notebook-controllers/src/notebookController.ts @@ -13,13 +13,13 @@ 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> = new Map(); + private readonly notebookRuntimes: Map> = new Map(); /** Incremented for each cell we create to give it a unique ID. */ private static CELL_COUNTER = 0; @@ -27,7 +27,7 @@ export class NotebookController implements vscode.Disposable { /** * @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 { + */ +async function updateNotebookLanguage(notebook: vscode.NotebookDocument, languageId: string): Promise { // 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) { diff --git a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts new file mode 100644 index 00000000000..5536fe2c01d --- /dev/null +++ b/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(); + + /** + * 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 { + 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((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()); + } +} diff --git a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts b/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts index da4337a0d06..47a455a84a1 100644 --- a/src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts +++ b/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;