From 1253a11a502c4fe5a7d482299fcab6f3da9fb0f3 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 12:47:15 +0100 Subject: [PATCH 01/11] fix: Prompt the user to select an environment when running one or more blocks. --- src/kernels/deepnote/types.ts | 10 + .../controllers/vscodeNotebookController.ts | 22 ++ .../deepnoteKernelAutoSelector.node.ts | 273 +++++++++++++++--- ...epnoteKernelAutoSelector.node.unit.test.ts | 15 +- 4 files changed, 274 insertions(+), 46 deletions(-) diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index 8fd384bc7c..0737491f46 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -221,6 +221,16 @@ export interface IDeepnoteKernelAutoSelector { token: vscode.CancellationToken ): Promise; + /** + * Ensure an environment is configured for the notebook before execution. + * If not configured, shows picker and sets up the kernel. + * @returns true if environment is ready, false if user cancelled + */ + ensureEnvironmentConfiguredBeforeExecution( + notebook: vscode.NotebookDocument, + token: vscode.CancellationToken + ): Promise; + /** * Force rebuild the controller for a notebook by clearing cached controller and metadata. * This is used when switching environments to ensure a new controller is created. diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 49a97791a9..6e954fe9b5 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -6,6 +6,7 @@ import type { KernelMessage } from '@jupyterlab/services'; import type { IAnyMessageArgs } from '@jupyterlab/services/lib/kernel/kernel'; import { CancellationError, + CancellationTokenSource, commands, Disposable, EventEmitter, @@ -88,6 +89,7 @@ import { KernelConnector } from './kernelConnector'; import { RemoteKernelReconnectBusyIndicator } from './remoteKernelReconnectBusyIndicator'; import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; import { notebookPathToDeepnoteProjectFilePath } from '../../platform/deepnote/deepnoteProjectUtils'; +import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteKernelAutoSelector } from '../../kernels/deepnote/types'; /** * Our implementation of the VSCode Notebook Controller. Called by VS code to execute cells in a notebook. Also displayed @@ -435,6 +437,26 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont if (!workspace.isTrusted) { return; } + + // For Deepnote notebooks, ensure environment is configured before execution + if (notebook.notebookType === DEEPNOTE_NOTEBOOK_TYPE) { + const kernelAutoSelector = + this.serviceContainer.tryGet(IDeepnoteKernelAutoSelector); + + if (kernelAutoSelector) { + const hasEnvironment = await kernelAutoSelector.ensureEnvironmentConfiguredBeforeExecution( + notebook, + new CancellationTokenSource().token + ); + + if (!hasEnvironment) { + // User cancelled - do not execute + logger.info(`Execution cancelled: no environment selected for ${getDisplayPath(notebook.uri)}`); + return; + } + } + } + logger.debug(`Handle Execution of Cells ${cells.map((c) => c.index)} for ${getDisplayPath(notebook.uri)}`); await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(notebook.uri, this.connection); telemetryTracker?.stop(); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 60fe624822..9f594a9616 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -11,6 +11,7 @@ import { NotebookController, NotebookControllerAffinity, NotebookDocument, + NotebookEditor, ProgressLocation, QuickPickItem, Uri, @@ -78,7 +79,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, { notebook: NotebookDocument; project: DeepnoteFile } >(); - private readonly deepnoteLoadingKernelController: NotebookController; + // Track per-notebook placeholder controllers for notebooks without configured environments + private readonly placeholderControllers = new Map(); constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @@ -101,9 +103,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, private readonly notebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller - ) { - this.deepnoteLoadingKernelController = DeepnoteKernelAutoSelector.createDeepnoteLoadingKernelController(); - } + ) {} public activate() { // Listen to notebook open events @@ -300,30 +300,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, logger.info(`Deepnote notebook closed: ${getDisplayPath(notebook.uri)}`); - return; - - // // Extract the base file URI to match what we used when registering - // const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); - // const notebookKey = baseFileUri.fsPath; - - // // Note: We intentionally don't clean up controllers, connection metadata, or servers here. - // // This allows the kernel to be reused if the user reopens the same .deepnote file. - // // The server will continue running and can be reused for better performance. - // // Cleanup will happen when the extension is disposed or when explicitly requested. - - // // However, we do unregister the server from the provider to keep it clean - // const serverHandle = this.notebookServerHandles.get(notebookKey); - // if (serverHandle) { - // logger.info(`Unregistering server for closed notebook: ${serverHandle}`); - // this.serverProvider.unregisterServer(serverHandle); - // this.notebookServerHandles.delete(notebookKey); - // } + // Clean up placeholder controller if it exists + const notebookKey = notebook.uri.toString(); + const placeholder = this.placeholderControllers.get(notebookKey); - // // Clean up pending init notebook tracking - // const projectId = notebook.metadata?.deepnoteProjectId; - // if (projectId) { - // this.projectsPendingInitNotebook.delete(projectId); - // } + if (placeholder) { + logger.info(`Disposing placeholder controller for closed notebook: ${getDisplayPath(notebook.uri)}`); + placeholder.dispose(); + this.placeholderControllers.delete(notebookKey); + } } public async onKernelStarted(kernel: IKernel) { @@ -456,13 +441,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, progress: { report(value: { message?: string; increment?: number }): void }, token: CancellationToken ): Promise { - this.deepnoteLoadingKernelController.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); - await commands.executeCommand('notebook.selectKernel', { - notebookEditor: notebook, - id: this.deepnoteLoadingKernelController.id, - extension: JVSC_EXTENSION_ID - }); - // baseFileUri identifies the PROJECT (without query/fragment) const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); // notebookKey uniquely identifies THIS NOTEBOOK (includes query with notebook ID) @@ -473,13 +451,53 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); if (environmentId == null) { + // No environment - create and select placeholder controller + const placeholder = this.createPlaceholderController(notebook); + placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); + + // Find the NotebookEditor for this document to properly select the kernel + const notebookEditor = await this.findNotebookEditor(notebook); + + if (notebookEditor) { + await commands.executeCommand('notebook.selectKernel', { + notebookEditor: notebookEditor, + id: placeholder.id, + extension: JVSC_EXTENSION_ID + }); + } else { + logger.warn( + `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` + ); + } + return false; } const environment = environmentId ? this.environmentManager.getEnvironment(environmentId) : undefined; + if (environment == null) { logger.info(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); + + // No environment - create and select placeholder controller + const placeholder = this.createPlaceholderController(notebook); + placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); + + // Find the NotebookEditor for this document to properly select the kernel + const notebookEditor = await this.findNotebookEditor(notebook); + + if (notebookEditor) { + await commands.executeCommand('notebook.selectKernel', { + notebookEditor: notebookEditor, + id: placeholder.id, + extension: JVSC_EXTENSION_ID + }); + } else { + logger.warn( + `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` + ); + } + return false; } @@ -505,6 +523,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, progress: { report(value: { message?: string; increment?: number }): void }, progressToken: CancellationToken ): Promise { + // Dispose placeholder controller if it exists (real controller is taking over) + const placeholder = this.placeholderControllers.get(notebookKey); + + if (placeholder) { + logger.info(`Disposing placeholder controller for ${getDisplayPath(notebook.uri)}`); + placeholder.dispose(); + this.placeholderControllers.delete(notebookKey); + } + logger.info(`Setting up kernel using configuration: ${configuration.name} (${configuration.id})`); progress.report({ message: `Using ${configuration.name}...` }); @@ -782,6 +809,78 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return kernelSpec; } + /** + * Ensure an environment is configured for the notebook before execution. + * If not configured, shows picker and sets up the kernel. + * @returns true if environment is ready, false if user cancelled + */ + public async ensureEnvironmentConfiguredBeforeExecution( + notebook: NotebookDocument, + token: CancellationToken + ): Promise { + Cancellation.throwIfCanceled(token); + + const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); + const notebookKey = notebook.uri.toString(); + const projectKey = baseFileUri.fsPath; + + // Check if environment is already configured + const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); + + if (existingEnvironmentId) { + // Validate that the environment still exists + const environment = this.environmentManager.getEnvironment(existingEnvironmentId); + + if (environment) { + logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); + return true; + } + + // Environment no longer exists, remove the stale mapping + logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); + await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); + } + + Cancellation.throwIfCanceled(token); + + // No environment configured, show the picker + logger.info(`No environment configured for ${getDisplayPath(notebook.uri)}, showing picker`); + const selectedEnvironment = await this.pickEnvironment(notebook.uri); + + if (!selectedEnvironment) { + logger.info(`User cancelled environment selection for ${getDisplayPath(notebook.uri)}`); + return false; + } + + Cancellation.throwIfCanceled(token); + + // Save the selection + await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironment.id); + + // Set up the kernel with the selected environment + await window.withProgress( + { + location: ProgressLocation.Notification, + title: l10n.t('Setting up Deepnote kernel...'), + cancellable: true + }, + async (progress, progressToken) => { + await this.ensureKernelSelectedWithConfiguration( + notebook, + selectedEnvironment, + baseFileUri, + notebookKey, + projectKey, + progress, + progressToken + ); + } + ); + + logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); + return true; + } + /** * Clear the controller selection for a notebook using a specific environment. * This is used when deleting an environment to unselect its controller from any open notebooks. @@ -809,6 +908,33 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, : Uri.joinPath(venvPath, 'bin', 'python'); } + /** + * Find the NotebookEditor for a given NotebookDocument. + * Required for properly selecting a kernel with the notebook.selectKernel command. + * Includes retry logic since the editor might not be visible immediately when the document opens. + */ + private async findNotebookEditor(notebook: NotebookDocument): Promise { + // Try to find immediately + let editor = window.visibleNotebookEditors.find((e) => e.notebook.uri.toString() === notebook.uri.toString()); + + if (editor) { + return editor; + } + + // If not found, wait briefly and retry (editor might not be visible yet) + for (let i = 0; i < 10; i++) { + await new Promise((resolve) => setTimeout(resolve, 100)); + + editor = window.visibleNotebookEditors.find((e) => e.notebook.uri.toString() === notebook.uri.toString()); + + if (editor) { + return editor; + } + } + + return undefined; + } + /** * Handle kernel selection errors with user-friendly messages and actions */ @@ -953,22 +1079,83 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } } - public static createDeepnoteLoadingKernelController(): NotebookController { + /** + * Create a placeholder controller for a notebook without a configured environment. + * Each notebook gets its own placeholder with a unique ID. + * The placeholder's executeHandler shows the environment picker when user tries to run cells. + */ + private createPlaceholderController(notebook: NotebookDocument): NotebookController { + const notebookKey = notebook.uri.toString(); + + // Check if we already have one + const existing = this.placeholderControllers.get(notebookKey); + + if (existing) { + return existing; + } + const controller = notebooks.createNotebookController( - `deepnote-loading-kernel`, + `deepnote-placeholder-${notebookKey}`, DEEPNOTE_NOTEBOOK_TYPE, - l10n.t('Loading Deepnote Kernel...') + l10n.t('Deepnote: Select Environment') ); - // Set it as the preferred controller immediately - controller.supportsExecutionOrder = false; - controller.supportedLanguages = ['python']; + controller.supportsExecutionOrder = true; + controller.supportedLanguages = ['python', 'sql', 'markdown']; + + // Execution handler that shows environment picker when user tries to run without an environment + controller.executeHandler = async (cells, doc) => { + logger.info( + `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ + cells.length + } cells` + ); + + const cts = new CancellationTokenSource(); + + try { + const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, cts.token); + + if (!hasEnvironment) { + logger.info(`User cancelled environment selection, not executing cells`); + return; + } + + // Environment is now configured, execute the cells through the kernel + const docNotebookKey = doc.uri.toString(); + const realController = this.notebookControllers.get(docNotebookKey); + + if (!realController) { + logger.error(`No controller found after environment configuration for ${docNotebookKey}`); + return; + } + + logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); - // Execution handler that does nothing - cells will just sit there until real kernel is ready - controller.executeHandler = () => { - // No-op: execution is blocked until the real controller takes over + // Get or create a kernel for this notebook with the new connection + const kernel = this.kernelProvider.getOrCreate(doc, { + metadata: realController.connection, + controller: realController.controller, + resourceUri: doc.uri + }); + + // Execute cells through the kernel + const kernelExecution = this.kernelProvider.getKernelExecution(kernel); + + for (const cell of cells) { + await kernelExecution.executeCell(cell); + } + + logger.info(`Finished executing ${cells.length} cells`); + } catch (error) { + logger.error(`Error in placeholder controller execute handler`, error); + } finally { + cts.dispose(); + } }; + this.placeholderControllers.set(notebookKey, controller); + return controller; } } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index a99508aca5..8141e32093 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -107,9 +107,18 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Mock disposable registry - push returns the index when(mockDisposableRegistry.push(anything())).thenReturn(0); - sandbox - .stub(DeepnoteKernelAutoSelector, 'createDeepnoteLoadingKernelController') - .returns(mock()); + // Mock notebooks.createNotebookController to return a mock controller for the loading kernel + const mockLoadingController = { + id: 'deepnote-loading-kernel', + supportsExecutionOrder: false, + supportedLanguages: ['python'], + executeHandler: undefined as unknown, + updateNotebookAffinity: sandbox.stub(), + dispose: sandbox.stub() + } as unknown as NotebookController; + when(mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything())).thenReturn( + mockLoadingController + ); // Create selector instance selector = new DeepnoteKernelAutoSelector( From de059fc0493c672bc982771e48b7f94e6e639ff4 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 15:10:34 +0100 Subject: [PATCH 02/11] CR Feedback --- src/kernels/deepnote/types.ts | 38 ++++++------ .../controllers/vscodeNotebookController.ts | 22 ++++--- .../deepnoteKernelAutoSelector.node.ts | 61 +++++++++++-------- 3 files changed, 69 insertions(+), 52 deletions(-) diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index 0737491f46..cd9669b28d 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -211,15 +211,12 @@ export interface IDeepnoteServerProvider { export const IDeepnoteKernelAutoSelector = Symbol('IDeepnoteKernelAutoSelector'); export interface IDeepnoteKernelAutoSelector { /** - * Automatically selects and starts a Deepnote kernel for the given notebook. + * Clear the controller selection for a notebook using a specific environment. + * This is used when deleting an environment to unselect its controller from any open notebooks. * @param notebook The notebook document - * @param token Cancellation token to cancel the operation + * @param environmentId The environment ID */ - ensureKernelSelected( - notebook: vscode.NotebookDocument, - progress: { report(value: { message?: string; increment?: number }): void }, - token: vscode.CancellationToken - ): Promise; + clearControllerForEnvironment(notebook: vscode.NotebookDocument, environmentId: string): void; /** * Ensure an environment is configured for the notebook before execution. @@ -232,24 +229,15 @@ export interface IDeepnoteKernelAutoSelector { ): Promise; /** - * Force rebuild the controller for a notebook by clearing cached controller and metadata. - * This is used when switching environments to ensure a new controller is created. + * Automatically selects and starts a Deepnote kernel for the given notebook. * @param notebook The notebook document * @param token Cancellation token to cancel the operation */ - rebuildController( + ensureKernelSelected( notebook: vscode.NotebookDocument, progress: { report(value: { message?: string; increment?: number }): void }, token: vscode.CancellationToken - ): Promise; - - /** - * Clear the controller selection for a notebook using a specific environment. - * This is used when deleting an environment to unselect its controller from any open notebooks. - * @param notebook The notebook document - * @param environmentId The environment ID - */ - clearControllerForEnvironment(notebook: vscode.NotebookDocument, environmentId: string): void; + ): Promise; /** * Handle kernel selection errors with user-friendly messages and actions @@ -257,6 +245,18 @@ export interface IDeepnoteKernelAutoSelector { * @param notebook The notebook document associated with the error */ handleKernelSelectionError(error: unknown, notebook: vscode.NotebookDocument): Promise; + + /** + * Force rebuild the controller for a notebook by clearing cached controller and metadata. + * This is used when switching environments to ensure a new controller is created. + * @param notebook The notebook document + * @param token Cancellation token to cancel the operation + */ + rebuildController( + notebook: vscode.NotebookDocument, + progress: { report(value: { message?: string; increment?: number }): void }, + token: vscode.CancellationToken + ): Promise; } export const IDeepnoteEnvironmentManager = Symbol('IDeepnoteEnvironmentManager'); diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 6e954fe9b5..eeeb2614e8 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -444,15 +444,21 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont this.serviceContainer.tryGet(IDeepnoteKernelAutoSelector); if (kernelAutoSelector) { - const hasEnvironment = await kernelAutoSelector.ensureEnvironmentConfiguredBeforeExecution( - notebook, - new CancellationTokenSource().token - ); + const cts = new CancellationTokenSource(); - if (!hasEnvironment) { - // User cancelled - do not execute - logger.info(`Execution cancelled: no environment selected for ${getDisplayPath(notebook.uri)}`); - return; + try { + const hasEnvironment = await kernelAutoSelector.ensureEnvironmentConfiguredBeforeExecution( + notebook, + cts.token + ); + + if (!hasEnvironment) { + // User cancelled - do not execute + logger.info(`Execution cancelled: no environment selected for ${getDisplayPath(notebook.uri)}`); + return; + } + } finally { + cts.dispose(); } } } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 9f594a9616..e3578fc512 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -45,7 +45,7 @@ import { import { IJupyterKernelSpec, IKernel, IKernelProvider } from '../../kernels/types'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IPythonExtensionChecker } from '../../platform/api/types'; -import { Cancellation } from '../../platform/common/cancellation'; +import { Cancellation, isCancellationError } from '../../platform/common/cancellation'; import { JVSC_EXTENSION_ID, STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants'; import { getDisplayPath } from '../../platform/common/platform/fs-paths.node'; import { IConfigurationService, IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; @@ -65,23 +65,22 @@ import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; */ @injectable() export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, IExtensionSyncActivationService { - // Track server handles per PROJECT (baseFileUri) - one server per project - private readonly projectServerHandles = new Map(); + // Track connection metadata per NOTEBOOK for reuse + private readonly notebookConnectionMetadata = new Map(); // Track registered controllers per NOTEBOOK (full URI with query) - one controller per notebook private readonly notebookControllers = new Map(); // Track environment for each notebook private readonly notebookEnvironmentsIds = new Map(); - // Track connection metadata per NOTEBOOK for reuse - private readonly notebookConnectionMetadata = new Map(); + // Track per-notebook placeholder controllers for notebooks without configured environments + private readonly placeholderControllers = new Map(); + // Track server handles per PROJECT (baseFileUri) - one server per project + private readonly projectServerHandles = new Map(); // Track projects where we need to run init notebook (set during controller setup) private readonly projectsPendingInitNotebook = new Map< string, { notebook: NotebookDocument; project: DeepnoteFile } >(); - // Track per-notebook placeholder controllers for notebooks without configured environments - private readonly placeholderControllers = new Map(); - constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration, @@ -833,6 +832,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (environment) { logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); + return true; } @@ -849,6 +849,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (!selectedEnvironment) { logger.info(`User cancelled environment selection for ${getDisplayPath(notebook.uri)}`); + return false; } @@ -858,26 +859,36 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironment.id); // Set up the kernel with the selected environment - await window.withProgress( - { - location: ProgressLocation.Notification, - title: l10n.t('Setting up Deepnote kernel...'), - cancellable: true - }, - async (progress, progressToken) => { - await this.ensureKernelSelectedWithConfiguration( - notebook, - selectedEnvironment, - baseFileUri, - notebookKey, - projectKey, - progress, - progressToken - ); + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: l10n.t('Setting up Deepnote kernel...'), + cancellable: true + }, + async (progress, progressToken) => { + await this.ensureKernelSelectedWithConfiguration( + notebook, + selectedEnvironment, + baseFileUri, + notebookKey, + projectKey, + progress, + progressToken + ); + } + ); + } catch (error) { + if (token.isCancellationRequested || isCancellationError(error as Error)) { + logger.info(`Kernel setup cancelled for ${getDisplayPath(notebook.uri)}`); + + return false; } - ); + throw error; + } logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); + return true; } From 506fc538124edaf527da184bc9b4feae76997a60 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 15:18:18 +0100 Subject: [PATCH 03/11] refactor: Address additional CodeRabbit feedback - Extract selectPlaceholderController helper to eliminate duplicate code - Add NOTEBOOK_EDITOR_RETRY_COUNT and NOTEBOOK_EDITOR_RETRY_DELAY_MS constants - Add per-cell error handling in placeholder controller execute handler --- .../deepnoteKernelAutoSelector.node.ts | 76 +++++++++---------- 1 file changed, 36 insertions(+), 40 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index e3578fc512..dd93f6c3f7 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -60,6 +60,10 @@ import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; +// Constants for NotebookEditor retry logic +const NOTEBOOK_EDITOR_RETRY_COUNT = 10; +const NOTEBOOK_EDITOR_RETRY_DELAY_MS = 100; + /** * Automatically selects and starts Deepnote kernel for .deepnote notebooks */ @@ -450,24 +454,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); if (environmentId == null) { - // No environment - create and select placeholder controller - const placeholder = this.createPlaceholderController(notebook); - placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); - - // Find the NotebookEditor for this document to properly select the kernel - const notebookEditor = await this.findNotebookEditor(notebook); - - if (notebookEditor) { - await commands.executeCommand('notebook.selectKernel', { - notebookEditor: notebookEditor, - id: placeholder.id, - extension: JVSC_EXTENSION_ID - }); - } else { - logger.warn( - `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` - ); - } + await this.selectPlaceholderController(notebook); return false; } @@ -477,25 +464,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (environment == null) { logger.info(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); - - // No environment - create and select placeholder controller - const placeholder = this.createPlaceholderController(notebook); - placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); - - // Find the NotebookEditor for this document to properly select the kernel - const notebookEditor = await this.findNotebookEditor(notebook); - - if (notebookEditor) { - await commands.executeCommand('notebook.selectKernel', { - notebookEditor: notebookEditor, - id: placeholder.id, - extension: JVSC_EXTENSION_ID - }); - } else { - logger.warn( - `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` - ); - } + await this.selectPlaceholderController(notebook); return false; } @@ -933,8 +902,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } // If not found, wait briefly and retry (editor might not be visible yet) - for (let i = 0; i < 10; i++) { - await new Promise((resolve) => setTimeout(resolve, 100)); + for (let i = 0; i < NOTEBOOK_EDITOR_RETRY_COUNT; i++) { + await new Promise((resolve) => setTimeout(resolve, NOTEBOOK_EDITOR_RETRY_DELAY_MS)); editor = window.visibleNotebookEditors.find((e) => e.notebook.uri.toString() === notebook.uri.toString()); @@ -946,6 +915,28 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return undefined; } + /** + * Create and select a placeholder controller for a notebook without a configured environment. + */ + private async selectPlaceholderController(notebook: NotebookDocument): Promise { + const placeholder = this.createPlaceholderController(notebook); + placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); + + const notebookEditor = await this.findNotebookEditor(notebook); + + if (notebookEditor) { + await commands.executeCommand('notebook.selectKernel', { + notebookEditor: notebookEditor, + id: placeholder.id, + extension: JVSC_EXTENSION_ID + }); + } else { + logger.warn( + `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` + ); + } + } + /** * Handle kernel selection errors with user-friendly messages and actions */ @@ -1154,7 +1145,12 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const kernelExecution = this.kernelProvider.getKernelExecution(kernel); for (const cell of cells) { - await kernelExecution.executeCell(cell); + try { + await kernelExecution.executeCell(cell); + } catch (cellError) { + logger.error(`Error executing cell ${cell.index}`, cellError); + // Continue with remaining cells + } } logger.info(`Finished executing ${cells.length} cells`); From b8af711e09e15aee02fd064a5866dd12b641e4a3 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 16:03:43 +0100 Subject: [PATCH 04/11] CR Feedback --- .../deepnote/deepnoteKernelAutoSelector.node.ts | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index dd93f6c3f7..ce88a34c72 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1113,10 +1113,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); - const cts = new CancellationTokenSource(); + // No external cancellation source for execute handler - use a never-cancelled token + const neverCancelledToken: CancellationToken = { + isCancellationRequested: false, + // eslint-disable-next-line @typescript-eslint/no-empty-function + onCancellationRequested: () => ({ dispose: () => {} }) + }; try { - const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, cts.token); + const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, neverCancelledToken); if (!hasEnvironment) { logger.info(`User cancelled environment selection, not executing cells`); @@ -1156,8 +1161,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, logger.info(`Finished executing ${cells.length} cells`); } catch (error) { logger.error(`Error in placeholder controller execute handler`, error); - } finally { - cts.dispose(); } }; From c6eb7fe2b225aa8a3c0e9c6e54e23f7c69048ae9 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 16:13:04 +0100 Subject: [PATCH 05/11] CR --- .../deepnoteKernelAutoSelector.node.ts | 49 ++++++++++++++++++- 1 file changed, 47 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index ce88a34c72..7d69cf4ba8 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -800,9 +800,54 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const environment = this.environmentManager.getEnvironment(existingEnvironmentId); if (environment) { - logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); + // Also verify that a controller exists for this notebook + const existingController = this.notebookControllers.get(notebookKey); - return true; + if (existingController) { + logger.info( + `Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}` + ); + + return true; + } + + // Controller is missing (e.g., after VS Code restart) - need to set up kernel + logger.info( + `Environment "${environment.name}" configured but controller missing for ${getDisplayPath( + notebook.uri + )}, triggering setup` + ); + + // Set up the kernel with the existing environment + try { + await window.withProgress( + { + location: ProgressLocation.Notification, + title: l10n.t('Setting up Deepnote kernel...'), + cancellable: true + }, + async (progress, progressToken) => { + await this.ensureKernelSelectedWithConfiguration( + notebook, + environment, + baseFileUri, + notebookKey, + projectKey, + progress, + progressToken + ); + } + ); + + return true; + } catch (error) { + if (token.isCancellationRequested || isCancellationError(error as Error)) { + logger.info(`Kernel setup cancelled for ${getDisplayPath(notebook.uri)}`); + + return false; + } + throw error; + } } // Environment no longer exists, remove the stale mapping From e9f404722d330c5124cccf2747f22984139e5f97 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 16:35:18 +0100 Subject: [PATCH 06/11] CR --- .../deepnoteKernelAutoSelector.node.ts | 56 +++++++++++++++---- 1 file changed, 45 insertions(+), 11 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 7d69cf4ba8..049410896f 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -208,10 +208,11 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, public async pickEnvironment(notebookUri: Uri): Promise { logger.info(`Picking environment for notebook ${getDisplayPath(notebookUri)}`); + // Wait for environment manager to finish loading environments from storage await this.environmentManager.waitForInitialization(); + const environments = this.environmentManager.listEnvironments(); - // Build quick pick items const items: (QuickPickItem & { environment?: DeepnoteEnvironment })[] = environments.map((env) => { return { label: env.name, @@ -222,12 +223,13 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment: env }; }); - // Add "Create new" option at the end + items.push({ label: '$(add) Create New Environment', description: 'Set up a new kernel environment', alwaysShow: true }); + const selected = await window.showQuickPick(items, { placeHolder: `Select an environment for ${getDisplayPath(notebookUri)}`, matchOnDescription: true, @@ -236,25 +238,29 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (!selected) { logger.info('User cancelled environment selection'); - return undefined; // User cancelled + return; // User cancelled } if (!selected.environment) { - // User chose "Create new" - execute the create command and retry logger.info('User chose to create new environment - triggering create command'); + await commands.executeCommand('deepnote.environments.create'); - // After creation, refresh the list and show picker again + const newEnvironments = this.environmentManager.listEnvironments(); + if (newEnvironments.length > environments.length) { - // A new environment was created, show the picker again logger.info('Environment created, showing picker again'); + return this.pickEnvironment(notebookUri); } - // User cancelled creation + logger.info('No new environment created'); - return undefined; + + return; } + logger.info(`Selected environment "${selected.environment.name}" for notebook ${getDisplayPath(notebookUri)}`); + return selected.environment; } @@ -359,8 +365,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Check if this is a cancellation error - if so, just log and continue if (error instanceof Error && error.message === 'Cancelled') { logger.info(`Init notebook cancelled for project ${projectId}`); + return; } + logger.error('Error running init notebook', error); // Continue anyway - don't block user if init fails } finally { @@ -838,8 +846,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ); } ); - - return true; } catch (error) { if (token.isCancellationRequested || isCancellationError(error as Error)) { logger.info(`Kernel setup cancelled for ${getDisplayPath(notebook.uri)}`); @@ -848,6 +854,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } throw error; } + + // Verify controller was actually created + const createdController = this.notebookControllers.get(notebookKey); + + if (!createdController) { + logger.warn( + `Controller not created for "${environment.name}" on ${getDisplayPath( + notebook.uri + )} after setup` + ); + + return false; + } + + return true; } // Environment no longer exists, remove the stale mapping @@ -901,6 +922,19 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, throw error; } + // Verify controller was actually created + const createdController = this.notebookControllers.get(notebookKey); + + if (!createdController) { + logger.warn( + `Controller not created for "${selectedEnvironment.name}" on ${getDisplayPath( + notebook.uri + )} after setup` + ); + + return false; + } + logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); return true; @@ -957,7 +991,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } } - return undefined; + return; } /** From 49c9b4d8bf9b233f167b1cb5e8e2ca071c8c8a9d Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 16:38:00 +0100 Subject: [PATCH 07/11] Remove some comments --- src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 049410896f..89713926b9 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -800,15 +800,12 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const notebookKey = notebook.uri.toString(); const projectKey = baseFileUri.fsPath; - // Check if environment is already configured const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); if (existingEnvironmentId) { - // Validate that the environment still exists const environment = this.environmentManager.getEnvironment(existingEnvironmentId); if (environment) { - // Also verify that a controller exists for this notebook const existingController = this.notebookControllers.get(notebookKey); if (existingController) { @@ -819,14 +816,12 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return true; } - // Controller is missing (e.g., after VS Code restart) - need to set up kernel logger.info( `Environment "${environment.name}" configured but controller missing for ${getDisplayPath( notebook.uri )}, triggering setup` ); - // Set up the kernel with the existing environment try { await window.withProgress( { @@ -873,6 +868,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Environment no longer exists, remove the stale mapping logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); + await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); } @@ -890,7 +886,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, Cancellation.throwIfCanceled(token); - // Save the selection await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironment.id); // Set up the kernel with the selected environment From b64bf9aac70e9554ba8e917925b3e89bcc8eb0c5 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 17:10:56 +0100 Subject: [PATCH 08/11] refactor --- .../deepnoteKernelAutoSelector.node.ts | 134 +++++++++--------- 1 file changed, 64 insertions(+), 70 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 89713926b9..ca16482e2b 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -802,79 +802,52 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); - if (existingEnvironmentId) { - const environment = this.environmentManager.getEnvironment(existingEnvironmentId); - - if (environment) { - const existingController = this.notebookControllers.get(notebookKey); - - if (existingController) { - logger.info( - `Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}` - ); - - return true; - } - - logger.info( - `Environment "${environment.name}" configured but controller missing for ${getDisplayPath( - notebook.uri - )}, triggering setup` - ); + // No environment configured - need to pick one + if (!existingEnvironmentId) { + return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, token); + } - try { - await window.withProgress( - { - location: ProgressLocation.Notification, - title: l10n.t('Setting up Deepnote kernel...'), - cancellable: true - }, - async (progress, progressToken) => { - await this.ensureKernelSelectedWithConfiguration( - notebook, - environment, - baseFileUri, - notebookKey, - projectKey, - progress, - progressToken - ); - } - ); - } catch (error) { - if (token.isCancellationRequested || isCancellationError(error as Error)) { - logger.info(`Kernel setup cancelled for ${getDisplayPath(notebook.uri)}`); + const environment = this.environmentManager.getEnvironment(existingEnvironmentId); - return false; - } - throw error; - } + // Environment no longer exists - remove stale mapping and pick a new one + if (!environment) { + logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); + await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); - // Verify controller was actually created - const createdController = this.notebookControllers.get(notebookKey); + return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, token); + } - if (!createdController) { - logger.warn( - `Controller not created for "${environment.name}" on ${getDisplayPath( - notebook.uri - )} after setup` - ); + const existingController = this.notebookControllers.get(notebookKey); - return false; - } + // Environment and controller already configured + if (existingController) { + logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); - return true; - } + return true; + } - // Environment no longer exists, remove the stale mapping - logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); + // Environment exists but controller is missing - set it up + logger.info( + `Environment "${environment.name}" configured but controller missing for ${getDisplayPath( + notebook.uri + )}, triggering setup` + ); - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); - } + return this.setupKernelForEnvironment(notebook, environment, baseFileUri, notebookKey, projectKey, token); + } + /** + * Pick an environment and set up the kernel for a notebook. + */ + private async pickAndSetupEnvironment( + notebook: NotebookDocument, + baseFileUri: Uri, + notebookKey: string, + projectKey: string, + token: CancellationToken + ): Promise { Cancellation.throwIfCanceled(token); - // No environment configured, show the picker logger.info(`No environment configured for ${getDisplayPath(notebook.uri)}, showing picker`); const selectedEnvironment = await this.pickEnvironment(notebook.uri); @@ -888,7 +861,33 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironment.id); - // Set up the kernel with the selected environment + const result = await this.setupKernelForEnvironment( + notebook, + selectedEnvironment, + baseFileUri, + notebookKey, + projectKey, + token + ); + + if (result) { + logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); + } + + return result; + } + + /** + * Set up the kernel for a given environment. + */ + private async setupKernelForEnvironment( + notebook: NotebookDocument, + environment: DeepnoteEnvironment, + baseFileUri: Uri, + notebookKey: string, + projectKey: string, + token: CancellationToken + ): Promise { try { await window.withProgress( { @@ -899,7 +898,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, async (progress, progressToken) => { await this.ensureKernelSelectedWithConfiguration( notebook, - selectedEnvironment, + environment, baseFileUri, notebookKey, projectKey, @@ -917,21 +916,16 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, throw error; } - // Verify controller was actually created const createdController = this.notebookControllers.get(notebookKey); if (!createdController) { logger.warn( - `Controller not created for "${selectedEnvironment.name}" on ${getDisplayPath( - notebook.uri - )} after setup` + `Controller not created for "${environment.name}" on ${getDisplayPath(notebook.uri)} after setup` ); return false; } - logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); - return true; } From 30b951d8bf58f0be51bf0fb73cec3f48a7616001 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 17:29:25 +0100 Subject: [PATCH 09/11] cr --- .../deepnoteKernelAutoSelector.node.ts | 56 ++++++++++++++++--- 1 file changed, 47 insertions(+), 9 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index ca16482e2b..64d109be6e 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -819,8 +819,35 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const existingController = this.notebookControllers.get(notebookKey); - // Environment and controller already configured + // Environment and controller already configured - but verify interpreter path still matches if (existingController) { + const existingInterpreter = existingController.connection.interpreter; + + if (existingInterpreter) { + const expectedInterpreter = + process.platform === 'win32' + ? Uri.joinPath(environment.venvPath, 'Scripts', 'python.exe') + : Uri.joinPath(environment.venvPath, 'bin', 'python'); + + if (existingInterpreter.uri.fsPath !== expectedInterpreter.fsPath) { + logger.warn( + `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingInterpreter.uri.fsPath}. Recreating controller.` + ); + + existingController.dispose(); + this.notebookControllers.delete(notebookKey); + + return this.setupKernelForEnvironment( + notebook, + environment, + baseFileUri, + notebookKey, + projectKey, + token + ); + } + } + logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); return true; @@ -1181,18 +1208,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); - // No external cancellation source for execute handler - use a never-cancelled token - const neverCancelledToken: CancellationToken = { - isCancellationRequested: false, - // eslint-disable-next-line @typescript-eslint/no-empty-function - onCancellationRequested: () => ({ dispose: () => {} }) - }; + // Create a cancellation token that cancels when the notebook is closed + const cts = new CancellationTokenSource(); + const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { + if (closedDoc.uri.toString() === doc.uri.toString()) { + logger.info(`Notebook closed during environment setup, cancelling operation`); + cts.cancel(); + } + }); try { - const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, neverCancelledToken); + const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, cts.token); if (!hasEnvironment) { logger.info(`User cancelled environment selection, not executing cells`); + return; } @@ -1202,6 +1232,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (!realController) { logger.error(`No controller found after environment configuration for ${docNotebookKey}`); + return; } @@ -1228,7 +1259,14 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, logger.info(`Finished executing ${cells.length} cells`); } catch (error) { - logger.error(`Error in placeholder controller execute handler`, error); + if (isCancellationError(error)) { + logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); + } else { + logger.error(`Error in placeholder controller execute handler`, error); + } + } finally { + closeListener.dispose(); + cts.dispose(); } }; From 721642169b75ce0fd1e766d71d67b043ca9fccfa Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 17:38:29 +0100 Subject: [PATCH 10/11] CR --- .../deepnoteKernelAutoSelector.node.ts | 90 ++++++++++--------- 1 file changed, 48 insertions(+), 42 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 64d109be6e..63e7129200 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -526,26 +526,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Verify the controller's interpreter path matches the expected venv path // This handles cases where notebooks were used in VS Code and now opened in Cursor - const existingInterpreter = existingController.connection.interpreter; - if (existingInterpreter) { - const expectedInterpreter = - process.platform === 'win32' - ? Uri.joinPath(configuration.venvPath, 'Scripts', 'python.exe') - : Uri.joinPath(configuration.venvPath, 'bin', 'python'); - - if (existingInterpreter.uri.fsPath !== expectedInterpreter.fsPath) { - logger.warn( - `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingInterpreter.uri.fsPath}. Recreating controller.` - ); - // Dispose old controller and recreate it - existingController.dispose(); - this.notebookControllers.delete(notebookKey); - } else { - logger.info(`Controller verified, selecting it`); - await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken); - return; - } + if (this.isControllerInterpreterValid(existingController, configuration.venvPath)) { + logger.info(`Controller verified, selecting it`); + await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken); + + return; } + + const expectedInterpreter = this.getVenvInterpreterUri(configuration.venvPath); + logger.warn( + `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingController.connection.interpreter?.uri.fsPath}. Recreating controller.` + ); + + // Dispose old controller and recreate it + existingController.dispose(); + this.notebookControllers.delete(notebookKey); } // Ensure server is running (startServer is idempotent - returns early if already running) @@ -821,31 +816,23 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Environment and controller already configured - but verify interpreter path still matches if (existingController) { - const existingInterpreter = existingController.connection.interpreter; - - if (existingInterpreter) { - const expectedInterpreter = - process.platform === 'win32' - ? Uri.joinPath(environment.venvPath, 'Scripts', 'python.exe') - : Uri.joinPath(environment.venvPath, 'bin', 'python'); - - if (existingInterpreter.uri.fsPath !== expectedInterpreter.fsPath) { - logger.warn( - `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingInterpreter.uri.fsPath}. Recreating controller.` - ); + if (!this.isControllerInterpreterValid(existingController, environment.venvPath)) { + const expectedInterpreter = this.getVenvInterpreterUri(environment.venvPath); + logger.warn( + `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingController.connection.interpreter?.uri.fsPath}. Recreating controller.` + ); - existingController.dispose(); - this.notebookControllers.delete(notebookKey); + existingController.dispose(); + this.notebookControllers.delete(notebookKey); - return this.setupKernelForEnvironment( - notebook, - environment, - baseFileUri, - notebookKey, - projectKey, - token - ); - } + return this.setupKernelForEnvironment( + notebook, + environment, + baseFileUri, + notebookKey, + projectKey, + token + ); } logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); @@ -983,6 +970,25 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, : Uri.joinPath(venvPath, 'bin', 'python'); } + /** + * Check if a controller's interpreter path matches the expected venv path. + * Returns true when no interpreter is present (nothing to validate) or when paths match. + */ + private isControllerInterpreterValid( + controller: { connection: { interpreter?: { uri: Uri } } }, + venvPath: Uri + ): boolean { + const existingInterpreter = controller.connection.interpreter; + + if (!existingInterpreter) { + return true; + } + + const expectedInterpreter = this.getVenvInterpreterUri(venvPath); + + return existingInterpreter.uri.fsPath === expectedInterpreter.fsPath; + } + /** * Find the NotebookEditor for a given NotebookDocument. * Required for properly selecting a kernel with the notebook.selectKernel command. From d347cbe6f5b43e8b50f53cd165021e4334b91074 Mon Sep 17 00:00:00 2001 From: Christoffer Artmann Date: Tue, 3 Feb 2026 17:47:16 +0100 Subject: [PATCH 11/11] Update claude --- CLAUDE.md | 41 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 86bc821f1b..6862398bca 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,7 +5,6 @@ - Use `Uri.joinPath()` for constructing file paths to ensure platform-correct path separators (e.g., `Uri.joinPath(venvPath, 'share', 'jupyter', 'kernels')` instead of string concatenation with `/`) - Follow established patterns, especially when importing new packages (e.g. instead of importing uuid directly, use the helper `import { generateUuid } from '../platform/common/uuid';`) - ## Code conventions - Always run `npx prettier` before committing @@ -19,7 +18,6 @@ - Tests run against compiled JavaScript files in `out/` directory - Use `assert.deepStrictEqual()` for object comparisons instead of checking individual properties - ## Project Structure - VSCode extension for Jupyter notebooks @@ -40,3 +38,42 @@ - Whitespace is good for readability, add a blank line after const groups and before return statements - Separate third-party and local file imports - How the extension works is described in @specs/architecture.md + +## Best Practices + +### Resource Cleanup + +- Always dispose `CancellationTokenSource` - never create inline without storing/disposing +- Use try/finally to ensure cleanup: + ```typescript + const cts = new CancellationTokenSource(); + try { + await fn(cts.token); + } finally { + cts.dispose(); + } + ``` + +### DRY Principle + +- Extract duplicate logic into helper methods to prevent drift +- When similar logic appears in multiple places (e.g., placeholder controller setup, interpreter validation), consolidate it + +### Magic Numbers + +- Extract magic numbers (retry counts, delays, timeouts) as named constants near the top of the module + +### Error Handling + +- Use per-iteration error handling in loops - wrap each iteration in try/catch so one failure doesn't stop the rest +- Handle `withProgress` cancellation gracefully - it throws when user cancels, so wrap in try/catch and return appropriate value + +### State Validation + +- Verify state after async setup operations - methods can return early without throwing, so check expected state was created +- Validate cached state before early returns - before returning "already configured", verify the state is still valid (e.g., interpreter paths match, controllers aren't stale) + +### Cancellation Tokens + +- Use real cancellation tokens tied to lifecycle events instead of fake/never-cancelled tokens +- Create `CancellationTokenSource` tied to relevant events (e.g., notebook close, cell cancel)