feat: Enhance Deepnote server management with additional package support

tkislan · tkislan · commit c4d241b17704 · 2025-11-02T14:35:07.000Z
- Added `additionalPackages` parameter to `startServer` and `startServerForEnvironment` methods in `DeepnoteServerStarter` to allow installation of extra packages.
- Updated the `createDeepnoteServerConfigHandle` utility function for consistent server handle creation.
- Refactored kernel connection metadata handling in `DeepnoteKernelAutoSelector` to utilize the new server config handle.
- Cleaned up test cases to reflect changes in method signatures and ensure proper functionality.

This commit improves the flexibility of the Deepnote server setup by enabling the installation of additional packages during server startup.

Signed-off-by: Tomas Kislan &lt;tomas@kislan.sk&gt;
diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts
@@ -102,6 +102,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
     public async startServer(
         interpreter: PythonEnvironment,
         venvPath: Uri,
+        additionalPackages: string[],
         environmentId: string,
         deepnoteFileUri: Uri,
         token?: CancellationToken
@@ -110,9 +111,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
         const serverKey = `${fileKey}-${environmentId}`;
 
         // Wait for any pending operations on this environment to complete
-        let pendingOp = this.pendingOperations.get(fileKey);
+        let pendingOp = this.pendingOperations.get(serverKey);
         if (pendingOp) {
-            logger.info(`Waiting for pending operation on ${fileKey} to complete...`);
+            logger.info(`Waiting for pending operation on ${serverKey} to complete...`);
             try {
                 await pendingOp.promise;
             } catch {
@@ -131,7 +132,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
                 }
 
                 // Start the operation if not already pending
-                pendingOp = this.pendingOperations.get(fileKey);
+                pendingOp = this.pendingOperations.get(serverKey);
 
                 if (pendingOp && pendingOp.type === 'start') {
                     // TODO - check pending operation environment id ?
@@ -140,7 +141,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
             } else {
                 // Stop the existing server
                 logger.info(
-                    `Stopping existing server for ${fileKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...`
+                    `Stopping existing server for ${serverKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...`
                 );
                 await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token);
                 // TODO - Clear controllers for the notebook ?
@@ -157,39 +158,20 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
             existingContext = newContext;
         }
 
-        // if (existingContext == null) {
-        //     // TODO - solve with better typing
-        //     throw new Error('Invariant violation: existingContext should not be null here');
-        // }
-
-        // // If server is already running for this environment, return existing info
-        // // const existingServerInfo = this.serverInfos.get(environmentId);
-        // const existingServerInfo = this.serverInfos.get(fileKey);
-        // if (existingServerInfo && (await this.isServerRunning(existingServerInfo))) {
-        //     logger.info(`Deepnote server already running at ${existingServerInfo.url} for ${fileKey}`);
-        //     return existingServerInfo;
-        // }
-
-        // // Start the operation if not already pending
-        // pendingOp = this.pendingOperations.get(fileKey);
-
-        // if (pendingOp && pendingOp.type === 'start') {
-        //     return await pendingOp.promise;
-        // }
-
         // Start the operation and track it
         const operation = {
             type: 'start' as const,
             promise: this.startServerForEnvironment(
                 existingContext,
                 interpreter,
                 venvPath,
+                additionalPackages,
                 environmentId,
                 deepnoteFileUri,
                 token
             )
         };
-        this.pendingOperations.set(fileKey, operation);
+        this.pendingOperations.set(serverKey, operation);
 
         try {
             const result = await operation.promise;
@@ -199,8 +181,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
             return result;
         } finally {
             // Remove from pending operations when done
-            if (this.pendingOperations.get(fileKey) === operation) {
-                this.pendingOperations.delete(fileKey);
+            if (this.pendingOperations.get(serverKey) === operation) {
+                this.pendingOperations.delete(serverKey);
             }
         }
     }
@@ -253,6 +235,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
         projectContext: ProjectContext,
         interpreter: PythonEnvironment,
         venvPath: Uri,
+        additionalPackages: string[],
         environmentId: string,
         deepnoteFileUri: Uri,
         token?: CancellationToken
@@ -272,6 +255,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
 
         Cancellation.throwIfCanceled(token);
 
+        await this.toolkitInstaller.installAdditionalPackages(venvPath, additionalPackages, token);
+
+        Cancellation.throwIfCanceled(token);
+
         // Allocate both ports with global lock to prevent race conditions
         // Note: allocatePorts reserves both ports immediately in serverInfos
         // const { jupyterPort, lspPort } = await this.allocatePorts(environmentId);
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts
@@ -19,6 +19,7 @@ import {
 } from '../../../platform/interpreter/helpers';
 import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
 import { IKernelProvider } from '../../types';
+import { createDeepnoteServerConfigHandle } from '../../../platform/deepnote/deepnoteServerUtils.node';
 
 /**
  * View controller for the Deepnote kernel environments tree view.
@@ -320,7 +321,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
             const connectionMetadata = kernel.kernelConnectionMetadata;
             if (connectionMetadata.kind === 'startUsingDeepnoteKernel') {
                 const deepnoteMetadata = connectionMetadata as DeepnoteKernelConnectionMetadata;
-                const expectedHandle = `deepnote-config-server-${environmentId}-${notebook.uri.fsPath}`;
+                const expectedHandle = createDeepnoteServerConfigHandle(environmentId, notebook.uri);
 
                 if (deepnoteMetadata.serverProviderHandle.handle === expectedHandle) {
                     logger.info(
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts
@@ -12,6 +12,7 @@ import { PythonEnvironment } from '../../../platform/pythonEnvironments/info';
 import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock';
 import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node';
 import * as interpreterHelpers from '../../../platform/interpreter/helpers';
+import { createDeepnoteServerConfigHandle } from '../../../platform/deepnote/deepnoteServerUtils.node';
 
 suite('DeepnoteEnvironmentsView', () => {
     let view: DeepnoteEnvironmentsView;
@@ -501,7 +502,7 @@ suite('DeepnoteEnvironmentsView', () => {
                 kernelConnectionMetadata: {
                     kind: 'startUsingDeepnoteKernel',
                     serverProviderHandle: {
-                        handle: `deepnote-config-server-${testEnvironmentId}`
+                        handle: createDeepnoteServerConfigHandle(testEnvironmentId, openNotebook1.uri)
                     }
                 },
                 dispose: sinon.stub().resolves()
@@ -511,7 +512,7 @@ suite('DeepnoteEnvironmentsView', () => {
                 kernelConnectionMetadata: {
                     kind: 'startUsingDeepnoteKernel',
                     serverProviderHandle: {
-                        handle: 'deepnote-config-server-different-env'
+                        handle: createDeepnoteServerConfigHandle('different-env-id', openNotebook3.uri)
                     }
                 },
                 dispose: sinon.stub().resolves()
diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts
@@ -138,6 +138,7 @@ export interface IDeepnoteServerStarter {
     startServer(
         interpreter: PythonEnvironment,
         venvPath: vscode.Uri,
+        additionalPackages: string[],
         environmentId: string,
         deepnoteFileUri: vscode.Uri,
         token?: vscode.CancellationToken
diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts
@@ -8,13 +8,11 @@ import {
     workspace,
     NotebookControllerAffinity,
     window,
-    NotebookController,
     CancellationTokenSource,
     Disposable,
     Uri,
     l10n,
     env,
-    notebooks,
     ProgressLocation,
     QuickPickItem,
     commands
@@ -50,6 +48,7 @@ import { DeepnoteKernelError } from '../../platform/errors/deepnoteKernelErrors'
 import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment';
 import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
 import { IOutputChannel } from '../../platform/common/types';
+import { createDeepnoteServerConfigHandle } from '../../platform/deepnote/deepnoteServerUtils.node';
 
 /**
  * Automatically selects and starts Deepnote kernel for .deepnote notebooks
@@ -62,8 +61,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
     private readonly notebookControllers = new Map<string, IVSCodeNotebookController>();
     // Track connection metadata per notebook file for reuse
     private readonly notebookConnectionMetadata = new Map<string, DeepnoteKernelConnectionMetadata>();
-    // Track temporary loading controllers that get disposed when real controller is ready
-    private readonly loadingControllers = new Map<string, NotebookController>();
     // Track projects where we need to run init notebook (set during controller setup)
     private readonly projectsPendingInitNotebook = new Map<
         string,
@@ -133,25 +130,26 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                 title: l10n.t('Auto-selecting Deepnote kernel...'),
                 cancellable: false
             },
-            async (progress, token) => {
+            // async (progress, token) => {
+            async () => {
                 try {
                     const result = await this.ensureKernelSelected(notebook);
                     if (!result) {
-                        const selectedEnvironment = await this.pickEnvironment(notebook.uri);
-                        if (selectedEnvironment) {
-                            await this.notebookEnvironmentMapper.setEnvironmentForNotebook(
-                                notebook.uri,
-                                selectedEnvironment.id
-                            );
-                            await this.ensureKernelSelectedWithConfiguration(
-                                notebook,
-                                selectedEnvironment,
-                                notebook.uri,
-                                notebook.uri.fsPath,
-                                progress,
-                                token
-                            );
-                        }
+                        // const selectedEnvironment = await this.pickEnvironment(notebook.uri);
+                        // if (selectedEnvironment) {
+                        //     await this.notebookEnvironmentMapper.setEnvironmentForNotebook(
+                        //         notebook.uri,
+                        //         selectedEnvironment.id
+                        //     );
+                        //     await this.ensureKernelSelectedWithConfiguration(
+                        //         notebook,
+                        //         selectedEnvironment,
+                        //         notebook.uri,
+                        //         notebook.uri.fsPath,
+                        //         progress,
+                        //         token
+                        //     );
+                        // }
                     }
                 } catch (error) {
                     logger.error(`Failed to auto-select Deepnote kernel for ${getDisplayPath(notebook.uri)}`, error);
@@ -459,6 +457,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         const serverInfo = await this.serverStarter.startServer(
             configuration.pythonInterpreter,
             configuration.venvPath,
+            configuration.packages ?? [],
             configuration.id,
             baseFileUri,
             progressToken
@@ -473,7 +472,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         const serverProviderHandle: JupyterServerProviderHandle = {
             extensionId: JVSC_EXTENSION_ID,
             id: 'deepnote-server',
-            handle: `deepnote-config-server-${configuration.id}-${baseFileUri.fsPath}`
+            handle: createDeepnoteServerConfigHandle(configuration.id, baseFileUri)
         };
 
         // Register the server with the provider
@@ -643,7 +642,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
             return;
         }
 
-        const expectedHandle = `deepnote-config-server-${environmentId}`;
+        const expectedHandle = createDeepnoteServerConfigHandle(environmentId, notebook.uri);
 
         if (selectedController.connection.serverProviderHandle.handle === expectedHandle) {
             // Unselect the controller by setting affinity to Default
@@ -654,31 +653,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         }
     }
 
-    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}`);
-    }
-
     /**
      * Handle kernel selection errors with user-friendly messages and actions
      */
diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts
diff --git a/src/platform/deepnote/deepnoteServerUtils.node.ts b/src/platform/deepnote/deepnoteServerUtils.node.ts
