chore: code review comments from src/kernels/deepnote/deepnoteToolkitInstaller.node.ts

Author: saltenasl

src/kernels/deepnote/deepnoteServerProvider.node.ts

@@ -9,6 +9,7 @@ import { IDisposableRegistry } from '../../platform/common/types';
 import { IJupyterServerProviderRegistry } from '../jupyter/types';
 import { JVSC_EXTENSION_ID } from '../../platform/common/constants';
 import { logger } from '../../platform/logging';
+import { DeepnoteServerNotFoundError } from '../../platform/errors/deepnoteServerNotFoundError';
 import { DeepnoteServerInfo } from './types';
 
 /**
@@ -79,7 +80,7 @@ export class DeepnoteServerProvider implements IExtensionSyncActivationService,
         const serverInfo = this.servers.get(server.id);
 
         if (!serverInfo) {
-            throw new Error(`Deepnote server not found: ${server.id}`);
+            throw new DeepnoteServerNotFoundError(server.id);
         }
 
         return {

src/kernels/deepnote/deepnoteServerStarter.node.ts

@@ -2,15 +2,15 @@
 // Licensed under the MIT License.
 
 import { inject, injectable, named } from 'inversify';
-import { Uri } from 'vscode';
+import { CancellationToken, Uri } from 'vscode';
 import { PythonEnvironment } from '../../platform/pythonEnvironments/info';
 import { IDeepnoteServerStarter, IDeepnoteToolkitInstaller, DeepnoteServerInfo, DEEPNOTE_DEFAULT_PORT } from './types';
 import { IProcessServiceFactory, ObservableExecutionResult } from '../../platform/common/process/types.node';
 import { logger } from '../../platform/logging';
-import { IOutputChannel, IDisposable } from '../../platform/common/types';
+import { IOutputChannel, IDisposable, IHttpClient } from '../../platform/common/types';
 import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
 import { sleep } from '../../platform/common/utils/async';
-import { HttpClient } from '../../platform/common/net/httpClient';
+import { Cancellation, raceCancellationError } from '../../platform/common/cancellation';
 import getPort from 'get-port';
 
 /**
@@ -21,31 +21,74 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter {
     private readonly serverProcesses: Map<string, ObservableExecutionResult<string>> = new Map();
     private readonly serverInfos: Map<string, DeepnoteServerInfo> = new Map();
     private readonly disposablesByFile: Map<string, IDisposable[]> = new Map();
-    private readonly httpClient = new HttpClient();
+    // Track in-flight operations per file to prevent concurrent start/stop
+    private readonly pendingOperations: Map<string, Promise<DeepnoteServerInfo | void>> = new Map();
 
     constructor(
         @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
         @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller,
-        @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel
+        @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel,
+        @inject(IHttpClient) private readonly httpClient: IHttpClient
     ) {}
 
-    public async getOrStartServer(interpreter: PythonEnvironment, deepnoteFileUri: Uri): Promise<DeepnoteServerInfo> {
+    public async getOrStartServer(
+        interpreter: PythonEnvironment,
+        deepnoteFileUri: Uri,
+        token?: CancellationToken
+    ): Promise<DeepnoteServerInfo> {
         const fileKey = deepnoteFileUri.fsPath;
 
+        // Wait for any pending operations on this file to complete
+        const pendingOp = this.pendingOperations.get(fileKey);
+        if (pendingOp) {
+            logger.info(`Waiting for pending operation on ${fileKey} to complete...`);
+            try {
+                await pendingOp;
+            } catch {
+                // Ignore errors from previous operations
+            }
+        }
+
         // If server is already running for this file, return existing info
         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 and track it
+        const operation = this.startServerImpl(interpreter, deepnoteFileUri, token);
+        this.pendingOperations.set(fileKey, operation);
+
+        try {
+            const result = await operation;
+            return result;
+        } finally {
+            // Remove from pending operations when done
+            if (this.pendingOperations.get(fileKey) === operation) {
+                this.pendingOperations.delete(fileKey);
+            }
+        }
+    }
+
+    private async startServerImpl(
+        interpreter: PythonEnvironment,
+        deepnoteFileUri: Uri,
+        token?: CancellationToken
+    ): Promise<DeepnoteServerInfo> {
+        const fileKey = deepnoteFileUri.fsPath;
+
+        Cancellation.throwIfCanceled(token);
+
         // Ensure toolkit is installed
         logger.info(`Ensuring deepnote-toolkit is installed for ${fileKey}...`);
-        const installed = await this.toolkitInstaller.ensureInstalled(interpreter, deepnoteFileUri);
+        const installed = await this.toolkitInstaller.ensureInstalled(interpreter, deepnoteFileUri, token);
         if (!installed) {
             throw new Error('Failed to install deepnote-toolkit. Please check the output for details.');
         }
 
+        Cancellation.throwIfCanceled(token);
+
         // Find available port
         const port = await getPort({ host: 'localhost', port: DEEPNOTE_DEFAULT_PORT });
         logger.info(`Starting deepnote-toolkit server on port ${port} for ${fileKey}`);
@@ -88,7 +131,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter {
         const serverInfo = { url, port };
         this.serverInfos.set(fileKey, serverInfo);
 
-        const serverReady = await this.waitForServer(serverInfo, 30000);
+        const serverReady = await this.waitForServer(serverInfo, 30000, token);
         if (!serverReady) {
             await this.stopServer(deepnoteFileUri);
             throw new Error('Deepnote server failed to start within timeout period');
@@ -102,6 +145,34 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter {
 
     public async stopServer(deepnoteFileUri: Uri): Promise<void> {
         const fileKey = deepnoteFileUri.fsPath;
+
+        // Wait for any pending operations on this file to complete
+        const pendingOp = this.pendingOperations.get(fileKey);
+        if (pendingOp) {
+            logger.info(`Waiting for pending operation on ${fileKey} before stopping...`);
+            try {
+                await pendingOp;
+            } catch {
+                // Ignore errors from previous operations
+            }
+        }
+
+        // Start the stop operation and track it
+        const operation = this.stopServerImpl(deepnoteFileUri);
+        this.pendingOperations.set(fileKey, operation);
+
+        try {
+            await operation;
+        } finally {
+            // Remove from pending operations when done
+            if (this.pendingOperations.get(fileKey) === operation) {
+                this.pendingOperations.delete(fileKey);
+            }
+        }
+    }
+
+    private async stopServerImpl(deepnoteFileUri: Uri): Promise<void> {
+        const fileKey = deepnoteFileUri.fsPath;
         const serverProcess = this.serverProcesses.get(fileKey);
 
         if (serverProcess) {
@@ -123,13 +194,18 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter {
         }
     }
 
-    private async waitForServer(serverInfo: DeepnoteServerInfo, timeout: number): Promise<boolean> {
+    private async waitForServer(
+        serverInfo: DeepnoteServerInfo,
+        timeout: number,
+        token?: CancellationToken
+    ): Promise<boolean> {
         const startTime = Date.now();
         while (Date.now() - startTime < timeout) {
+            Cancellation.throwIfCanceled(token);
             if (await this.isServerRunning(serverInfo)) {
                 return true;
             }
-            await sleep(500);
+            await raceCancellationError(token, sleep(500));
         }
         return false;
     }
@@ -143,4 +219,35 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter {
             return false;
         }
     }
+
+    public dispose(): void {
+        logger.info('Disposing DeepnoteServerStarter - stopping all servers...');
+
+        // Stop all server processes
+        for (const [fileKey, serverProcess] of this.serverProcesses.entries()) {
+            try {
+                logger.info(`Stopping Deepnote server for ${fileKey}...`);
+                serverProcess.proc?.kill();
+            } catch (ex) {
+                logger.error(`Error stopping Deepnote server for ${fileKey}: ${ex}`);
+            }
+        }
+
+        // Dispose all tracked disposables
+        for (const [fileKey, disposables] of this.disposablesByFile.entries()) {
+            try {
+                disposables.forEach((d) => d.dispose());
+            } catch (ex) {
+                logger.error(`Error disposing resources for ${fileKey}: ${ex}`);
+            }
+        }
+
+        // Clear all maps
+        this.serverProcesses.clear();
+        this.serverInfos.clear();
+        this.disposablesByFile.clear();
+        this.pendingOperations.clear();
+
+        logger.info('DeepnoteServerStarter disposed successfully');
+    }
 }

src/kernels/deepnote/deepnoteToolkitInstaller.node.ts

@@ -2,21 +2,24 @@
 // Licensed under the MIT License.
 
 import { inject, injectable, named } from 'inversify';
-import { Uri } from 'vscode';
+import { CancellationToken, Uri } from 'vscode';
 import { PythonEnvironment } from '../../platform/pythonEnvironments/info';
 import { IDeepnoteToolkitInstaller, DEEPNOTE_TOOLKIT_WHEEL_URL } from './types';
 import { IProcessServiceFactory } from '../../platform/common/process/types.node';
 import { logger } from '../../platform/logging';
 import { IOutputChannel, IExtensionContext } from '../../platform/common/types';
 import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
 import { IFileSystem } from '../../platform/common/platform/types';
+import { Cancellation } from '../../platform/common/cancellation';
 
 /**
  * Handles installation of the deepnote-toolkit Python package.
  */
 @injectable()
 export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
     private readonly venvPythonPaths: Map<string, Uri> = new Map();
+    // Track in-flight installations per venv path to prevent concurrent installs
+    private readonly pendingInstallations: Map<string, Promise<PythonEnvironment | undefined>> = new Map();
 
     constructor(
         @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
@@ -56,17 +59,54 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
 
     public async ensureInstalled(
         baseInterpreter: PythonEnvironment,
-        deepnoteFileUri: Uri
+        deepnoteFileUri: Uri,
+        token?: CancellationToken
     ): Promise<PythonEnvironment | undefined> {
         const venvPath = this.getVenvPath(deepnoteFileUri);
+        const venvKey = venvPath.fsPath;
+
+        // Wait for any pending installation for this venv to complete
+        const pendingInstall = this.pendingInstallations.get(venvKey);
+        if (pendingInstall) {
+            logger.info(`Waiting for pending installation for ${venvKey} to complete...`);
+            try {
+                return await pendingInstall;
+            } catch {
+                // If the previous installation failed, continue to retry
+                logger.info(`Previous installation for ${venvKey} failed, retrying...`);
+            }
+        }
+
+        // Check if venv already exists with toolkit installed
+        const existingVenv = await this.getVenvInterpreter(deepnoteFileUri);
+        if (existingVenv && (await this.isToolkitInstalled(existingVenv))) {
+            logger.info(`deepnote-toolkit venv already exists and is ready for ${deepnoteFileUri.fsPath}`);
+            return existingVenv;
+        }
+
+        // Start the installation and track it
+        const installation = this.installImpl(baseInterpreter, deepnoteFileUri, venvPath, token);
+        this.pendingInstallations.set(venvKey, installation);
 
         try {
-            // Check if venv already exists with toolkit installed
-            const existingVenv = await this.getVenvInterpreter(deepnoteFileUri);
-            if (existingVenv && (await this.isToolkitInstalled(existingVenv))) {
-                logger.info(`deepnote-toolkit venv already exists and is ready for ${deepnoteFileUri.fsPath}`);
-                return existingVenv;
+            const result = await installation;
+            return result;
+        } finally {
+            // Remove from pending installations when done
+            if (this.pendingInstallations.get(venvKey) === installation) {
+                this.pendingInstallations.delete(venvKey);
             }
+        }
+    }
+
+    private async installImpl(
+        baseInterpreter: PythonEnvironment,
+        deepnoteFileUri: Uri,
+        venvPath: Uri,
+        token?: CancellationToken
+    ): Promise<PythonEnvironment | undefined> {
+        try {
+            Cancellation.throwIfCanceled(token);
 
             logger.info(`Creating virtual environment at ${venvPath.fsPath} for ${deepnoteFileUri.fsPath}`);
             this.outputChannel.appendLine(`Setting up Deepnote toolkit environment for ${deepnoteFileUri.fsPath}...`);
@@ -93,6 +133,8 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 return undefined;
             }
 
+            Cancellation.throwIfCanceled(token);
+
             // Get venv Python interpreter
             const venvInterpreter = await this.getVenvInterpreter(deepnoteFileUri);
             if (!venvInterpreter) {
@@ -111,6 +153,8 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 { throwOnStdErr: false }
             );
 
+            Cancellation.throwIfCanceled(token);
+
             if (installResult.stdout) {
                 this.outputChannel.appendLine(installResult.stdout);
             }

src/kernels/deepnote/types.ts

@@ -71,11 +71,13 @@ export interface IDeepnoteToolkitInstaller {
      * Ensures deepnote-toolkit is installed in a dedicated virtual environment.
      * @param baseInterpreter The base Python interpreter to use for creating the venv
      * @param deepnoteFileUri The URI of the .deepnote file (used to create a unique venv per file)
+     * @param token Cancellation token to cancel the operation
      * @returns The Python interpreter from the venv if installed successfully, undefined otherwise
      */
     ensureInstalled(
         baseInterpreter: PythonEnvironment,
-        deepnoteFileUri: vscode.Uri
+        deepnoteFileUri: vscode.Uri,
+        token?: vscode.CancellationToken
     ): Promise<PythonEnvironment | undefined>;
 
     /**
@@ -86,14 +88,19 @@ export interface IDeepnoteToolkitInstaller {
 }
 
 export const IDeepnoteServerStarter = Symbol('IDeepnoteServerStarter');
-export interface IDeepnoteServerStarter {
+export interface IDeepnoteServerStarter extends vscode.Disposable {
     /**
      * Starts or gets an existing deepnote-toolkit Jupyter server.
      * @param interpreter The Python interpreter to use
      * @param deepnoteFileUri The URI of the .deepnote file (for server management per file)
+     * @param token Cancellation token to cancel the operation
      * @returns Connection information (URL, port, etc.)
      */
-    getOrStartServer(interpreter: PythonEnvironment, deepnoteFileUri: vscode.Uri): Promise<DeepnoteServerInfo>;
+    getOrStartServer(
+        interpreter: PythonEnvironment,
+        deepnoteFileUri: vscode.Uri,
+        token?: vscode.CancellationToken
+    ): Promise<DeepnoteServerInfo>;
 
     /**
      * Stops the deepnote-toolkit server if running.
@@ -112,8 +119,10 @@ export const IDeepnoteKernelAutoSelector = Symbol('IDeepnoteKernelAutoSelector')
 export interface IDeepnoteKernelAutoSelector {
     /**
      * Automatically selects and starts a Deepnote kernel for the given notebook.
+     * @param notebook The notebook document
+     * @param token Cancellation token to cancel the operation
      */
-    ensureKernelSelected(notebook: vscode.NotebookDocument): Promise<void>;
+    ensureKernelSelected(notebook: vscode.NotebookDocument, token?: vscode.CancellationToken): Promise<void>;
 }
 
 export const DEEPNOTE_TOOLKIT_WHEEL_URL =