refactor: Simplify environment selection logic in DeepnoteEnvironmentPicker and DeepnoteEnvironmentsView

tkislan · tkislan · commit a09ba4009662 · 2025-10-31T12:28:32.000Z
- Removed redundant checks for existing environments and streamlined the user prompts for creating new environments.
- Updated the createEnvironmentCommand to return the created environment for better handling in the environment selection process.
- Enhanced logging for environment switching and user interactions.

This refactor improves the clarity and efficiency of the environment management workflow.

Signed-off-by: Tomas Kislan &lt;tomas@kislan.sk&gt;
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentPicker.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentPicker.ts
@@ -26,34 +26,6 @@ export class DeepnoteEnvironmentPicker {
 
         const environments = this.environmentManager.listEnvironments();
 
-        if (environments.length === 0) {
-            // No environments exist - prompt user to create one
-            const createLabel = l10n.t('Create Environment');
-            const cancelLabel = l10n.t('Cancel');
-
-            const choice = await window.showInformationMessage(
-                l10n.t('No environments found. Create one to use with {0}?', getDisplayPath(notebookUri)),
-                createLabel,
-                cancelLabel
-            );
-
-            if (choice === createLabel) {
-                // Trigger the create command
-                logger.info('Triggering create environment command from picker');
-                await commands.executeCommand('deepnote.environments.create');
-
-                // Check if an environment was created
-                const newEnvironments = this.environmentManager.listEnvironments();
-                if (newEnvironments.length > 0) {
-                    // Environment created, show picker again
-                    logger.info('Environment created, showing picker again');
-                    return this.pickEnvironment(notebookUri);
-                }
-            }
-
-            return undefined;
-        }
-
         // Build quick pick items
         const items: (QuickPickItem & { environment?: DeepnoteEnvironment })[] = environments.map((env) => {
             const envWithStatus = this.environmentManager.getEnvironmentWithStatus(env.id);
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts
@@ -11,7 +11,7 @@ import {
 } from '../types';
 import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node';
 import { DeepnoteEnvironmentTreeItem } from './deepnoteEnvironmentTreeItem.node';
-import { CreateDeepnoteEnvironmentOptions, EnvironmentStatus } from './deepnoteEnvironment';
+import { CreateDeepnoteEnvironmentOptions, DeepnoteEnvironment, EnvironmentStatus } from './deepnoteEnvironment';
 import {
     getCachedEnvironment,
     resolvedPythonEnvToJupyterEnv,
@@ -58,7 +58,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
         disposableRegistry.push(this);
     }
 
-    public async createEnvironmentCommand(): Promise<void> {
+    public async createEnvironmentCommand(): Promise<DeepnoteEnvironment | undefined> {
         try {
             // Step 1: Select Python interpreter
             const api = await this.pythonApiProvider.getNewApi();
@@ -124,7 +124,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
             // Step 3: Enter packages (optional)
             const packagesInput = await window.showInputBox({
                 prompt: l10n.t('Enter additional packages to install (comma-separated, optional)'),
-                placeHolder: l10n.t('e.g., pandas, numpy, matplotlib'),
+                placeHolder: l10n.t('e.g., matplotlib, terraform'),
                 validateInput: (value: string) => {
                     if (!value || value.trim().length === 0) {
                         return undefined; // Empty is OK
@@ -160,7 +160,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
             });
 
             // Create environment with progress
-            await window.withProgress(
+            return await window.withProgress(
                 {
                     location: ProgressLocation.Notification,
                     title: l10n.t('Creating environment "{0}"...', name),
@@ -183,6 +183,8 @@ export class DeepnoteEnvironmentsView implements Disposable {
                         void window.showInformationMessage(
                             l10n.t('Environment "{0}" created successfully!', config.name)
                         );
+
+                        return config;
                     } catch (error) {
                         logger.error('Failed to create environment', error);
                         throw error;
@@ -389,19 +391,6 @@ export class DeepnoteEnvironmentsView implements Disposable {
         // Get all environments
         const environments = this.environmentManager.listEnvironments();
 
-        if (environments.length === 0) {
-            const choice = await window.showInformationMessage(
-                l10n.t('No environments found. Create one first?'),
-                l10n.t('Create Environment'),
-                l10n.t('Cancel')
-            );
-
-            if (choice === l10n.t('Create Environment')) {
-                await commands.executeCommand('deepnote.environments.create');
-            }
-            return;
-        }
-
         // Build quick pick items
         const items: (QuickPickItem & { environmentId?: string })[] = environments.map((env) => {
             const envWithStatus = this.environmentManager.getEnvironmentWithStatus(env.id);
@@ -422,9 +411,11 @@ export class DeepnoteEnvironmentsView implements Disposable {
             };
         });
 
+        const createNewLabel = l10n.t('$(add) Create New Environment');
+
         // Add "Create new" option at the end
         items.push({
-            label: l10n.t('$(add) Create New Environment'),
+            label: createNewLabel,
             description: l10n.t('Set up a new kernel environment'),
             alwaysShow: true
         });
@@ -439,16 +430,26 @@ export class DeepnoteEnvironmentsView implements Disposable {
             return; // User cancelled
         }
 
-        if (!selected.environmentId) {
-            // User chose "Create new"
-            await commands.executeCommand('deepnote.environments.create');
-            return;
+        let selectedEnvironmentId: string | undefined;
+
+        if (selected.label === createNewLabel) {
+            const newEnvironment = await this.createEnvironmentCommand();
+            if (newEnvironment == null) {
+                return;
+            }
+            // return;
+            selectedEnvironmentId = newEnvironment.id;
+        } else {
+            selectedEnvironmentId = selected.environmentId;
         }
 
         // Check if user selected the same environment
-        if (selected.environmentId === currentEnvironmentId) {
+        if (selectedEnvironmentId === currentEnvironmentId) {
             logger.info(`User selected the same environment - no changes needed`);
             return;
+        } else if (selectedEnvironmentId == null) {
+            logger.info('User cancelled environment selection');
+            return;
         }
 
         // Check if any cells are currently executing using the kernel execution state
@@ -475,9 +476,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
         }
 
         // User selected a different environment - switch to it
-        logger.info(
-            `Switching notebook ${getDisplayPath(activeNotebook.uri)} to environment ${selected.environmentId}`
-        );
+        logger.info(`Switching notebook ${getDisplayPath(activeNotebook.uri)} to environment ${selectedEnvironmentId}`);
 
         try {
             await window.withProgress(
@@ -488,16 +487,13 @@ export class DeepnoteEnvironmentsView implements Disposable {
                 },
                 async () => {
                     // Update the notebook-to-environment mapping
-                    await this.notebookEnvironmentMapper.setEnvironmentForNotebook(
-                        baseFileUri,
-                        selected.environmentId!
-                    );
+                    await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironmentId);
 
                     // Force rebuild the controller with the new environment
                     // This clears cached metadata and creates a fresh controller.
                     await this.kernelAutoSelector.rebuildController(activeNotebook);
 
-                    logger.info(`Successfully switched to environment ${selected.environmentId}`);
+                    logger.info(`Successfully switched to environment ${selectedEnvironmentId}`);
                 }
             );
 
diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts
@@ -9,13 +9,13 @@ import {
     NotebookControllerAffinity,
     window,
     ProgressLocation,
-    notebooks,
     NotebookController,
     CancellationTokenSource,
     Disposable,
     Uri,
     l10n,
-    env
+    env,
+    commands
 } from 'vscode';
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
 import { IDisposableRegistry } from '../../platform/common/types';
@@ -122,17 +122,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
 
         logger.info(`Deepnote notebook opened: ${getDisplayPath(notebook.uri)}`);
 
-        // Check if we already have a controller ready for this notebook
-        const baseFileUri = notebook.uri.with({ query: '', fragment: '' });
-        const notebookKey = baseFileUri.fsPath;
-        const hasExistingController = this.notebookControllers.has(notebookKey);
-
-        // If no existing controller, create a temporary "Loading" controller immediately
-        // This prevents the kernel selector from appearing when user clicks Run All
-        if (!hasExistingController) {
-            this.createLoadingController(notebook, notebookKey);
-        }
-
         // Always try to ensure kernel is selected (this will reuse existing controllers)
         // Don't await - let it happen in background so notebook opens quickly
         void this.ensureKernelSelected(notebook).catch((error) => {
@@ -301,20 +290,22 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
     }
 
     public async ensureKernelSelected(notebook: NotebookDocument, _token?: CancellationToken): Promise<void> {
-        return window.withProgress(
+        const baseFileUri = notebook.uri.with({ query: '', fragment: '' });
+        const notebookKey = baseFileUri.fsPath;
+
+        const kernelSelected = await window.withProgress(
             {
                 location: ProgressLocation.Notification,
                 title: l10n.t('Loading Deepnote Kernel'),
                 cancellable: true
             },
-            async (progress, progressToken) => {
+            async (progress, progressToken): Promise<boolean> => {
                 try {
                     logger.info(`Ensuring Deepnote kernel is selected for ${getDisplayPath(notebook.uri)}`);
 
                     // Extract the base file URI (without query parameters)
                     // Notebooks from the same .deepnote file have different URIs with ?notebook=id query params
-                    const baseFileUri = notebook.uri.with({ query: '', fragment: '' });
-                    const notebookKey = baseFileUri.fsPath;
+
                     logger.info(`Base Deepnote file: ${getDisplayPath(baseFileUri)}`);
 
                     // Check if we already have a controller for this notebook file
@@ -345,7 +336,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                         const selectedController = this.controllerRegistration.getSelected(notebook);
                         if (selectedController && selectedController.id === existingController.id) {
                             logger.info(`Controller already selected for ${getDisplayPath(notebook.uri)}`);
-                            return;
+                            return true;
                         }
 
                         // Auto-select the existing controller for this notebook
@@ -363,7 +354,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                             logger.info(`Disposed loading controller for ${notebookKey}`);
                         }
 
-                        return;
+                        return true;
                     }
 
                     // No existing controller - check if user has selected a configuration for this notebook
@@ -387,10 +378,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                         selectedConfig = await this.configurationPicker.pickEnvironment(baseFileUri);
 
                         if (!selectedConfig) {
-                            logger.info(`User cancelled configuration selection - no kernel will be loaded`);
-                            throw new Error(
-                                'No environment selected. Please create an environment using the Deepnote Environments view.'
+                            logger.info(
+                                `User cancelled configuration selection or no environment found - no kernel will be loaded`
                             );
+                            return false;
                         }
 
                         // Save the selection
@@ -401,20 +392,61 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                     }
 
                     // Use the selected configuration
-                    return this.ensureKernelSelectedWithConfiguration(
+                    await this.ensureKernelSelectedWithConfiguration(
                         notebook,
                         selectedConfig,
                         baseFileUri,
                         notebookKey,
                         progress,
                         progressToken
                     );
+
+                    return true;
                 } catch (ex) {
                     logger.error('Failed to auto-select Deepnote kernel', ex);
                     throw ex;
                 }
             }
         );
+
+        if (!kernelSelected) {
+            const createLabel = l10n.t('Create Environment');
+            const cancelLabel = l10n.t('Cancel');
+
+            const choice = await window.showInformationMessage(
+                l10n.t('No environments found. Create one to use with {0}?', getDisplayPath(baseFileUri)),
+                createLabel,
+                cancelLabel
+            );
+
+            if (choice === createLabel) {
+                // Trigger the create command
+                logger.info('Triggering create environment command from picker');
+                await commands.executeCommand('deepnote.environments.create');
+
+                const selectedConfig = await this.configurationPicker.pickEnvironment(baseFileUri);
+                if (!selectedConfig) {
+                    return;
+                }
+
+                const tmpCancellation = new CancellationTokenSource();
+                const tmpCancellationToken = tmpCancellation.token;
+
+                // Use the selected configuration
+                await this.ensureKernelSelectedWithConfiguration(
+                    notebook,
+                    selectedConfig,
+                    baseFileUri,
+                    notebookKey,
+                    {
+                        report: () => {
+                            logger.info('Progress report');
+                        }
+                    },
+                    tmpCancellationToken
+                );
+            }
+        }
     }
 
     private async ensureKernelSelectedWithConfiguration(
@@ -627,31 +659,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         return kernelSpec;
     }
 
-    private createLoadingController(notebook: NotebookDocument, notebookKey: string): void {
-        // Create a temporary controller that shows "Loading..." and prevents kernel selection prompt
-        const loadingController = notebooks.createNotebookController(
-            `deepnote-loading-${notebookKey}`,
-            DEEPNOTE_NOTEBOOK_TYPE,
-            l10n.t('Loading Deepnote Kernel...')
-        );
-
-        // Set it as the preferred controller immediately
-        loadingController.supportsExecutionOrder = false;
-        loadingController.supportedLanguages = ['python'];
-
-        // Execution handler that does nothing - cells will just sit there until real kernel is ready
-        loadingController.executeHandler = () => {
-            // No-op: execution is blocked until the real controller takes over
-        };
-
-        // Select this controller for the notebook
-        loadingController.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred);
-
-        // Store it so we can dispose it later
-        this.loadingControllers.set(notebookKey, loadingController);
-        logger.info(`Created loading controller for ${notebookKey}`);
-    }
-
     /**
      * 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.
