chore: second part of code review comments

saltenasl · saltenasl · commit 28c1f2d864bc · 2025-10-02T09:59:46.000+03:00
diff --git a/src/kernels/deepnote/deepnoteServerProvider.node.ts b/src/kernels/deepnote/deepnoteServerProvider.node.ts
@@ -10,14 +10,16 @@ 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';
+import { DeepnoteServerInfo, IDeepnoteServerProvider } from './types';
 
 /**
  * Jupyter Server Provider for Deepnote kernels.
  * This provider resolves server connections for Deepnote kernels.
  */
 @injectable()
-export class DeepnoteServerProvider implements IExtensionSyncActivationService, JupyterServerProvider {
+export class DeepnoteServerProvider
+    implements IDeepnoteServerProvider, IExtensionSyncActivationService, JupyterServerProvider
+{
     public readonly id = 'deepnote-server';
     private readonly _onDidChangeServers = new EventEmitter<void>();
     public readonly onDidChangeServers: Event<void> = this._onDidChangeServers.event;
@@ -53,6 +55,28 @@ export class DeepnoteServerProvider implements IExtensionSyncActivationService,
         this._onDidChangeServers.fire();
     }
 
+    /**
+     * Unregister a server for a specific handle.
+     * Called when the server is no longer needed or notebook is closed.
+     * No-op if the handle doesn't exist.
+     */
+    public unregisterServer(handle: string): void {
+        if (this.servers.has(handle)) {
+            logger.info(`Unregistering Deepnote server: ${handle}`);
+            this.servers.delete(handle);
+            this._onDidChangeServers.fire();
+        }
+    }
+
+    /**
+     * Dispose of all servers and resources.
+     */
+    public dispose(): void {
+        logger.info('Disposing Deepnote server provider, clearing all registered servers');
+        this.servers.clear();
+        this._onDidChangeServers.dispose();
+    }
+
     /**
      * Provides the list of available Deepnote servers.
      */
diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts
@@ -84,6 +84,18 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
             return existingVenv;
         }
 
+        // Double-check for race condition: another caller might have started installation
+        // while we were checking the venv
+        const pendingAfterCheck = this.pendingInstallations.get(venvKey);
+        if (pendingAfterCheck) {
+            logger.info(`Another installation started for ${venvKey} while checking, waiting for it...`);
+            try {
+                return await pendingAfterCheck;
+            } catch {
+                logger.info(`Concurrent installation for ${venvKey} failed, retrying...`);
+            }
+        }
+
         // Start the installation and track it
         const installation = this.installImpl(baseInterpreter, deepnoteFileUri, venvPath, token);
         this.pendingInstallations.set(venvKey, installation);
@@ -127,18 +139,21 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 throwOnStdErr: false
             });
 
-            if (venvResult.stderr && !venvResult.stderr.includes('WARNING')) {
-                logger.error(`Failed to create venv: ${venvResult.stderr}`);
-                this.outputChannel.appendLine(`Error creating venv: ${venvResult.stderr}`);
-                return undefined;
+            // Log any stderr output (warnings, etc.) but don't fail on it
+            if (venvResult.stderr) {
+                logger.info(`venv creation stderr: ${venvResult.stderr}`);
             }
 
             Cancellation.throwIfCanceled(token);
 
-            // Get venv Python interpreter
+            // Verify venv was created successfully by checking for the Python interpreter
             const venvInterpreter = await this.getVenvInterpreter(deepnoteFileUri);
             if (!venvInterpreter) {
-                logger.error('Failed to locate venv Python interpreter');
+                logger.error('Failed to create venv: Python interpreter not found after venv creation');
+                if (venvResult.stderr) {
+                    logger.error(`venv stderr: ${venvResult.stderr}`);
+                }
+                this.outputChannel.appendLine('Error: Failed to create virtual environment');
                 return undefined;
             }
 
diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts
@@ -115,6 +115,22 @@ export interface DeepnoteServerInfo {
     token?: string;
 }
 
+export const IDeepnoteServerProvider = Symbol('IDeepnoteServerProvider');
+export interface IDeepnoteServerProvider {
+    /**
+     * Register a server for a specific handle.
+     * Called by DeepnoteKernelAutoSelector when a server is started.
+     */
+    registerServer(handle: string, serverInfo: DeepnoteServerInfo): void;
+
+    /**
+     * Unregister a server for a specific handle.
+     * Called when the server is no longer needed or notebook is closed.
+     * No-op if the handle doesn't exist.
+     */
+    unregisterServer(handle: string): void;
+}
+
 export const IDeepnoteKernelAutoSelector = Symbol('IDeepnoteKernelAutoSelector');
 export interface IDeepnoteKernelAutoSelector {
     /**
diff --git a/src/notebooks/controllers/controllerRegistration.ts b/src/notebooks/controllers/controllerRegistration.ts
@@ -30,6 +30,7 @@ import {
 } from './types';
 import { VSCodeNotebookController } from './vscodeNotebookController';
 import { IJupyterVariablesProvider } from '../../kernels/variables/types';
+import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types';
 
 /**
  * Keeps track of registered controllers and available KernelConnectionMetadatas.
@@ -236,14 +237,14 @@ export class ControllerRegistration implements IControllerRegistration, IExtensi
     }
     addOrUpdate(
         metadata: KernelConnectionMetadata,
-        types: ('jupyter-notebook' | 'interactive' | 'deepnote')[]
+        types: (typeof JupyterNotebookView | typeof InteractiveWindowView | typeof DEEPNOTE_NOTEBOOK_TYPE)[]
     ): IVSCodeNotebookController[] {
         const { added, existing } = this.addImpl(metadata, types, true);
         return added.concat(existing);
     }
     addImpl(
         metadata: KernelConnectionMetadata,
-        types: ('jupyter-notebook' | 'interactive' | 'deepnote')[],
+        types: (typeof JupyterNotebookView | typeof InteractiveWindowView | typeof DEEPNOTE_NOTEBOOK_TYPE)[],
         triggerChangeEvent: boolean
     ): { added: IVSCodeNotebookController[]; existing: IVSCodeNotebookController[] } {
         const added: IVSCodeNotebookController[] = [];
diff --git a/src/notebooks/controllers/types.ts b/src/notebooks/controllers/types.ts
@@ -16,6 +16,7 @@ import { IDisposable } from '../../platform/common/types';
 import { JupyterServerCollection } from '../../api';
 import { EnvironmentPath } from '@vscode/python-extension';
 import type { VSCodeNotebookController } from './vscodeNotebookController';
+import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types';
 
 export const InteractiveControllerIdSuffix = ' (Interactive)';
 
@@ -83,7 +84,7 @@ export interface IControllerRegistration {
      */
     addOrUpdate(
         metadata: KernelConnectionMetadata,
-        types: (typeof JupyterNotebookView | typeof InteractiveWindowView | 'deepnote')[]
+        types: (typeof JupyterNotebookView | typeof InteractiveWindowView | typeof DEEPNOTE_NOTEBOOK_TYPE)[]
     ): IVSCodeNotebookController[];
     /**
      * Gets the controller for a particular connection
@@ -92,7 +93,7 @@ export interface IControllerRegistration {
      */
     get(
         connection: KernelConnectionMetadata,
-        notebookType: typeof JupyterNotebookView | typeof InteractiveWindowView | 'deepnote'
+        notebookType: typeof JupyterNotebookView | typeof InteractiveWindowView | typeof DEEPNOTE_NOTEBOOK_TYPE
     ): IVSCodeNotebookController | undefined;
 }
 
diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts
@@ -11,6 +11,7 @@ import {
     IDeepnoteKernelAutoSelector,
     IDeepnoteServerStarter,
     IDeepnoteToolkitInstaller,
+    IDeepnoteServerProvider,
     DEEPNOTE_NOTEBOOK_TYPE,
     DeepnoteKernelConnectionMetadata
 } from '../../kernels/deepnote/types';
@@ -19,7 +20,6 @@ import { JVSC_EXTENSION_ID } from '../../platform/common/constants';
 import { getDisplayPath } from '../../platform/common/platform/fs-paths';
 import { JupyterServerProviderHandle } from '../../kernels/jupyter/types';
 import { IPythonExtensionChecker } from '../../platform/api/types';
-import { DeepnoteServerProvider } from '../../kernels/deepnote/deepnoteServerProvider.node';
 import { JupyterLabHelper } from '../../kernels/jupyter/session/jupyterLabHelper';
 import { createJupyterConnectionInfo } from '../../kernels/jupyter/jupyterUtils';
 import { IJupyterRequestCreator, IJupyterRequestAgentCreator } from '../../kernels/jupyter/types';
@@ -31,14 +31,17 @@ import { disposeAsync } from '../../platform/common/utils';
  */
 @injectable()
 export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, IExtensionSyncActivationService {
+    // Track server handles per notebook URI for cleanup
+    private readonly notebookServerHandles = new Map<string, string>();
+
     constructor(
         @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
         @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration,
         @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
         @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller,
         @inject(IDeepnoteServerStarter) private readonly serverStarter: IDeepnoteServerStarter,
         @inject(IPythonExtensionChecker) private readonly pythonExtensionChecker: IPythonExtensionChecker,
-        @inject(DeepnoteServerProvider) private readonly serverProvider: DeepnoteServerProvider,
+        @inject(IDeepnoteServerProvider) private readonly serverProvider: IDeepnoteServerProvider,
         @inject(IJupyterRequestCreator) private readonly requestCreator: IJupyterRequestCreator,
         @inject(IJupyterRequestAgentCreator)
         @optional()
@@ -50,8 +53,13 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         // Listen to notebook open events
         workspace.onDidOpenNotebookDocument(this.onDidOpenNotebook, this, this.disposables);
 
-        // Handle currently open notebooks
-        workspace.notebookDocuments.forEach((d) => this.onDidOpenNotebook(d));
+        // Listen to notebook close events for cleanup
+        workspace.onDidCloseNotebookDocument(this.onDidCloseNotebook, this, this.disposables);
+
+        // Handle currently open notebooks - await all async operations
+        Promise.all(workspace.notebookDocuments.map((d) => this.onDidOpenNotebook(d))).catch((error) => {
+            logger.error(`Error handling open notebooks during activation: ${error}`);
+        });
     }
 
     private async onDidOpenNotebook(notebook: NotebookDocument) {
@@ -69,8 +77,34 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
             return;
         }
 
-        // Auto-select Deepnote kernel
-        await this.ensureKernelSelected(notebook);
+        // Auto-select Deepnote kernel with error handling
+        try {
+            await this.ensureKernelSelected(notebook);
+        } catch (error) {
+            logger.error(`Failed to auto-select Deepnote kernel for ${getDisplayPath(notebook.uri)}: ${error}`);
+            // Don't rethrow - we want activation to continue even if one notebook fails
+        }
+    }
+
+    private onDidCloseNotebook(notebook: NotebookDocument) {
+        // Only handle deepnote notebooks
+        if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) {
+            return;
+        }
+
+        logger.info(`Deepnote notebook closed: ${getDisplayPath(notebook.uri)}`);
+
+        // Extract the base file URI to match what we used when registering
+        const baseFileUri = notebook.uri.with({ query: '', fragment: '' });
+        const notebookKey = baseFileUri.fsPath;
+
+        // Unregister the server if we have a handle for this notebook
+        const serverHandle = this.notebookServerHandles.get(notebookKey);
+        if (serverHandle) {
+            logger.info(`Unregistering server for closed notebook: ${serverHandle}`);
+            this.serverProvider.unregisterServer(serverHandle);
+            this.notebookServerHandles.delete(notebookKey);
+        }
     }
 
     public async ensureKernelSelected(notebook: NotebookDocument, token?: CancellationToken): Promise<void> {
@@ -121,6 +155,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
             // Register the server with the provider so it can be resolved
             this.serverProvider.registerServer(serverProviderHandle.handle, serverInfo);
 
+            // Track the server handle for cleanup when notebook is closed
+            this.notebookServerHandles.set(baseFileUri.fsPath, serverProviderHandle.handle);
+
             // Connect to the server and get available kernel specs
             const connectionInfo = createJupyterConnectionInfo(
                 serverProviderHandle,
diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts
@@ -46,7 +46,8 @@ import { IDeepnoteNotebookManager } from './types';
 import {
     IDeepnoteToolkitInstaller,
     IDeepnoteServerStarter,
-    IDeepnoteKernelAutoSelector
+    IDeepnoteKernelAutoSelector,
+    IDeepnoteServerProvider
 } from '../kernels/deepnote/types';
 import { DeepnoteToolkitInstaller } from '../kernels/deepnote/deepnoteToolkitInstaller.node';
 import { DeepnoteServerStarter } from '../kernels/deepnote/deepnoteServerStarter.node';
@@ -129,8 +130,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
     // Deepnote kernel services
     serviceManager.addSingleton<IDeepnoteToolkitInstaller>(IDeepnoteToolkitInstaller, DeepnoteToolkitInstaller);
     serviceManager.addSingleton<IDeepnoteServerStarter>(IDeepnoteServerStarter, DeepnoteServerStarter);
-    serviceManager.addSingleton<DeepnoteServerProvider>(DeepnoteServerProvider, DeepnoteServerProvider);
-    serviceManager.addBinding(DeepnoteServerProvider, IExtensionSyncActivationService);
+    serviceManager.addSingleton<IDeepnoteServerProvider>(IDeepnoteServerProvider, DeepnoteServerProvider);
+    serviceManager.addBinding(IDeepnoteServerProvider, IExtensionSyncActivationService);
     serviceManager.addSingleton<IDeepnoteKernelAutoSelector>(IDeepnoteKernelAutoSelector, DeepnoteKernelAutoSelector);
     serviceManager.addBinding(IDeepnoteKernelAutoSelector, IExtensionSyncActivationService);
 
