refactor: Update kernel selection and environment handling in Deepnote

- Modified the `ensureKernelSelected` method to include progress reporting and cancellation token parameters for improved user feedback during kernel selection.
- Enhanced the `selectEnvironmentForNotebook` method to accept a notebook parameter, ensuring the correct notebook context is used when selecting environments.
- Updated the environment selection command to handle active notebooks more robustly, providing warnings when no Deepnote notebook is found.
- Refactored logging to provide clearer insights into notebook state changes and environment selection processes.

These changes improve the user experience by providing better feedback and ensuring that the correct notebook context is maintained during operations.

Signed-off-by: Tomas Kislan <[email protected]>

Author: tkislan

src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts

@@ -189,6 +189,7 @@ export class DeepnoteEnvironmentManager implements IExtensionSyncActivationServi
 
         config.lastUsedAt = new Date();
         await this.persistEnvironments();
+        this._onDidChangeEnvironments.fire();
     }
 
     /**

src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts

@@ -1,5 +1,15 @@
 import { inject, injectable } from 'inversify';
-import { commands, Disposable, l10n, ProgressLocation, QuickPickItem, TreeView, window, workspace } from 'vscode';
+import {
+    commands,
+    Disposable,
+    l10n,
+    NotebookDocument,
+    ProgressLocation,
+    QuickPickItem,
+    TreeView,
+    window,
+    workspace
+} from 'vscode';
 import { IDisposableRegistry } from '../../../platform/common/types';
 import { logger } from '../../../platform/logging';
 import { IPythonApiProvider } from '../../../platform/api/types';
@@ -243,9 +253,19 @@ export class DeepnoteEnvironmentsView implements Disposable {
 
         // Switch environment for notebook command
         this.disposables.push(
-            commands.registerCommand('deepnote.environments.selectForNotebook', async () => {
-                await this.selectEnvironmentForNotebook();
-            })
+            commands.registerCommand(
+                'deepnote.environments.selectForNotebook',
+                async (options?: { notebook?: NotebookDocument }) => {
+                    // Get the active notebook
+                    const activeNotebook = options?.notebook ?? window.activeNotebookEditor?.notebook;
+                    if (!activeNotebook || activeNotebook.notebookType !== 'deepnote') {
+                        void window.showWarningMessage(l10n.t('No active Deepnote notebook found'));
+                        return;
+                    }
+
+                    await this.selectEnvironmentForNotebook({ notebook: activeNotebook });
+                }
+            )
         );
     }
 
@@ -344,16 +364,11 @@ export class DeepnoteEnvironmentsView implements Disposable {
         }
     }
 
-    public async selectEnvironmentForNotebook(): Promise<void> {
-        // Get the active notebook
-        const activeNotebook = window.activeNotebookEditor?.notebook;
-        if (!activeNotebook || activeNotebook.notebookType !== 'deepnote') {
-            void window.showWarningMessage(l10n.t('No active Deepnote notebook found'));
-            return;
-        }
+    public async selectEnvironmentForNotebook({ notebook }: { notebook: NotebookDocument }): Promise<void> {
+        logger.info('Selecting environment for notebook:', notebook);
 
         // Get base file URI (without query/fragment)
-        const baseFileUri = activeNotebook.uri.with({ query: '', fragment: '' });
+        const baseFileUri = notebook.uri.with({ query: '', fragment: '' });
 
         // Get current environment selection
         const currentEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri);
@@ -421,7 +436,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
 
         // Check if any cells are currently executing using the kernel execution state
         // This is more reliable than checking executionSummary
-        const kernel = this.kernelProvider.get(activeNotebook);
+        const kernel = this.kernelProvider.get(notebook);
         const hasExecutingCells = kernel
             ? this.kernelProvider.getKernelExecution(kernel).pendingCells.length > 0
             : false;
@@ -443,7 +458,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
         }
 
         // User selected a different environment - switch to it
-        logger.info(`Switching notebook ${getDisplayPath(activeNotebook.uri)} to environment ${selectedEnvironmentId}`);
+        logger.info(`Switching notebook ${getDisplayPath(notebook.uri)} to environment ${selectedEnvironmentId}`);
 
         try {
             await window.withProgress(
@@ -459,7 +474,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
                     // Force rebuild the controller with the new environment
                     // This clears cached metadata and creates a fresh controller.
                     // await this.kernelAutoSelector.ensureKernelSelected(activeNotebook);
-                    await this.kernelAutoSelector.rebuildController(activeNotebook);
+                    await this.kernelAutoSelector.rebuildController(notebook);
 
                     logger.info(`Successfully switched to environment ${selectedEnvironmentId}`);
                 }

src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts

@@ -1,7 +1,7 @@
 import { assert } from 'chai';
 import * as sinon from 'sinon';
 import { anything, capture, instance, mock, when, verify, deepEqual, resetCalls } from 'ts-mockito';
-import { CancellationToken, Disposable, ProgressOptions, Uri } from 'vscode';
+import { CancellationToken, Disposable, NotebookDocument, ProgressOptions, Uri } from 'vscode';
 import { DeepnoteEnvironmentsView } from './deepnoteEnvironmentsView.node';
 import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteNotebookEnvironmentMapper } from '../types';
 import { IPythonApiProvider } from '../../../platform/api/types';
@@ -660,7 +660,7 @@ suite('DeepnoteEnvironmentsView', () => {
             when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(undefined);
 
             // Execute the command
-            await view.selectEnvironmentForNotebook();
+            await view.selectEnvironmentForNotebook({ notebook: mockNotebook as NotebookDocument });
 
             // Verify API calls
             verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).once();

src/kernels/deepnote/types.ts

@@ -189,7 +189,11 @@ export interface IDeepnoteKernelAutoSelector {
      * @param notebook The notebook document
      * @param token Cancellation token to cancel the operation
      */
-    ensureKernelSelected(notebook: vscode.NotebookDocument, token?: vscode.CancellationToken): Promise<boolean>;
+    ensureKernelSelected(
+        notebook: vscode.NotebookDocument,
+        progress: { report(value: { message?: string; increment?: number }): void },
+        token: vscode.CancellationToken
+    ): Promise<boolean>;
 
     /**
      * Force rebuild the controller for a notebook by clearing cached controller and metadata.