feat: Add clearControllerForEnvironment method and disposeKernelsUsingEnvironment logic

- Introduced clearControllerForEnvironment method in IDeepnoteKernelAutoSelector to unselect the controller for a notebook when an environment is deleted.
- Implemented disposeKernelsUsingEnvironment method in DeepnoteEnvironmentsView to dispose of kernels from open notebooks using the deleted environment.
- Enhanced unit tests for DeepnoteEnvironmentsView to verify kernel disposal behavior when an environment is deleted.

Signed-off-by: Tomas Kislan <tomas@kislan.sk>

Author: tkislan

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

@@ -3,12 +3,15 @@ import { IDeepnoteEnvironmentManager } from '../types';
 import { EnvironmentTreeItemType, DeepnoteEnvironmentTreeItem } from './deepnoteEnvironmentTreeItem.node';
 import { EnvironmentStatus } from './deepnoteEnvironment';
 import { inject, injectable } from 'inversify';
+import { IExtensionSyncActivationService } from '../../../platform/activation/types';
 
 /**
  * Tree data provider for the Deepnote kernel environments view
  */
 @injectable()
-export class DeepnoteEnvironmentTreeDataProvider implements TreeDataProvider<DeepnoteEnvironmentTreeItem>, Disposable {
+export class DeepnoteEnvironmentTreeDataProvider
+    implements TreeDataProvider<DeepnoteEnvironmentTreeItem>, IExtensionSyncActivationService, Disposable
+{
     private readonly _onDidChangeTreeData = new EventEmitter<DeepnoteEnvironmentTreeItem | undefined | void>();
     private readonly disposables: Disposable[] = [];
 
@@ -21,6 +24,10 @@ export class DeepnoteEnvironmentTreeDataProvider implements TreeDataProvider<Dee
         );
     }
 
+    public activate(): void {
+        this.refresh();
+    }
+
     public get onDidChangeTreeData(): Event<DeepnoteEnvironmentTreeItem | undefined | void> {
         return this._onDidChangeTreeData.event;
     }

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

@@ -1,9 +1,14 @@
 import { inject, injectable } from 'inversify';
-import { commands, Disposable, l10n, ProgressLocation, QuickPickItem, TreeView, window } from 'vscode';
+import { commands, Disposable, l10n, ProgressLocation, QuickPickItem, TreeView, window, workspace } from 'vscode';
 import { IDisposableRegistry } from '../../../platform/common/types';
 import { logger } from '../../../platform/logging';
 import { IPythonApiProvider } from '../../../platform/api/types';
-import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteNotebookEnvironmentMapper } from '../types';
+import {
+    DeepnoteKernelConnectionMetadata,
+    IDeepnoteEnvironmentManager,
+    IDeepnoteKernelAutoSelector,
+    IDeepnoteNotebookEnvironmentMapper
+} from '../types';
 import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node';
 import { DeepnoteEnvironmentTreeItem } from './deepnoteEnvironmentTreeItem.node';
 import { CreateDeepnoteEnvironmentOptions, EnvironmentStatus } from './deepnoteEnvironment';
@@ -303,6 +308,9 @@ export class DeepnoteEnvironmentsView implements Disposable {
                         await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(nb);
                     }
 
+                    // Dispose kernels from any open notebooks using this environment
+                    await this.disposeKernelsUsingEnvironment(environmentId);
+
                     await this.environmentManager.deleteEnvironment(environmentId, token);
                     logger.info(`Deleted environment: ${environmentId}`);
                 }
@@ -315,6 +323,52 @@ export class DeepnoteEnvironmentsView implements Disposable {
         }
     }
 
+    /**
+     * Dispose kernels from any open notebooks that are using the specified environment.
+     * This ensures the UI reflects that the kernel is no longer available.
+     */
+    private async disposeKernelsUsingEnvironment(environmentId: string): Promise<void> {
+        const openNotebooks = workspace.notebookDocuments;
+
+        for (const notebook of openNotebooks) {
+            // Only check Deepnote notebooks
+            if (notebook.notebookType !== 'deepnote') {
+                continue;
+            }
+
+            // Get the kernel for this notebook
+            const kernel = this.kernelProvider.get(notebook);
+            if (!kernel) {
+                continue;
+            }
+
+            // Check if this kernel is using the environment being deleted
+            const connectionMetadata = kernel.kernelConnectionMetadata;
+            if (connectionMetadata.kind === 'startUsingDeepnoteKernel') {
+                const deepnoteMetadata = connectionMetadata as DeepnoteKernelConnectionMetadata;
+                const expectedHandle = `deepnote-config-server-${environmentId}`;
+
+                if (deepnoteMetadata.serverProviderHandle.handle === expectedHandle) {
+                    logger.info(
+                        `Disposing kernel for notebook ${getDisplayPath(
+                            notebook.uri
+                        )} as it uses deleted environment ${environmentId}`
+                    );
+
+                    try {
+                        // First, unselect the controller from the notebook UI
+                        this.kernelAutoSelector.clearControllerForEnvironment(notebook, environmentId);
+
+                        // Then dispose the kernel
+                        await kernel.dispose();
+                    } catch (error) {
+                        logger.error(`Failed to dispose kernel for ${getDisplayPath(notebook.uri)}`, error);
+                    }
+                }
+            }
+        }
+    }
+
     public async selectEnvironmentForNotebook(): Promise<void> {
         // Get the active notebook
         const activeNotebook = window.activeNotebookEditor?.notebook;
@@ -359,7 +413,7 @@ export class DeepnoteEnvironmentsView implements Disposable {
             const isCurrent = currentEnvironment?.id === env.id;
 
             return {
-                label: `${icon} ${env.name} [${text}]${isCurrent ? ' $(check)' : ''}`,
+                label: `$(${icon}) ${env.name} [${text}]${isCurrent ? ' $(check)' : ''}`,
                 description: getDisplayPath(env.pythonInterpreter.uri),
                 detail: env.packages?.length
                     ? l10n.t('Packages: {0}', env.packages.join(', '))

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

@@ -456,6 +456,107 @@ suite('DeepnoteEnvironmentsView', () => {
             // Verify success message was shown
             verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once();
         });
+
+        test('should dispose kernels from open notebooks using the deleted environment', async () => {
+            // Mock environment exists
+            when(mockConfigManager.getEnvironment(testEnvironmentId)).thenReturn(testEnvironment);
+
+            // Mock user confirmation
+            when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn(
+                Promise.resolve('Delete')
+            );
+
+            // Mock notebooks using this environment
+            when(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testEnvironmentId)).thenReturn([]);
+            when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve();
+
+            // Mock open notebooks with kernels
+            const openNotebook1 = {
+                uri: Uri.file('/workspace/open-notebook1.deepnote'),
+                notebookType: 'deepnote',
+                isClosed: false
+            } as any;
+
+            const openNotebook2 = {
+                uri: Uri.file('/workspace/open-notebook2.deepnote'),
+                notebookType: 'jupyter-notebook',
+                isClosed: false
+            } as any;
+
+            const openNotebook3 = {
+                uri: Uri.file('/workspace/open-notebook3.deepnote'),
+                notebookType: 'deepnote',
+                isClosed: false
+            } as any;
+
+            // Mock workspace.notebookDocuments
+            when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([
+                openNotebook1,
+                openNotebook2,
+                openNotebook3
+            ]);
+
+            // Mock kernels
+            const mockKernel1 = {
+                kernelConnectionMetadata: {
+                    kind: 'startUsingDeepnoteKernel',
+                    serverProviderHandle: {
+                        handle: `deepnote-config-server-${testEnvironmentId}`
+                    }
+                },
+                dispose: sinon.stub().resolves()
+            };
+
+            const mockKernel3 = {
+                kernelConnectionMetadata: {
+                    kind: 'startUsingDeepnoteKernel',
+                    serverProviderHandle: {
+                        handle: 'deepnote-config-server-different-env'
+                    }
+                },
+                dispose: sinon.stub().resolves()
+            };
+
+            // Mock kernelProvider.get()
+            when(mockKernelProvider.get(openNotebook1)).thenReturn(mockKernel1 as any);
+            when(mockKernelProvider.get(openNotebook2)).thenReturn(undefined); // No kernel for jupyter notebook
+            when(mockKernelProvider.get(openNotebook3)).thenReturn(mockKernel3 as any);
+
+            // Mock window.withProgress
+            when(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).thenCall(
+                (_options: ProgressOptions, callback: Function) => {
+                    const mockProgress = {
+                        report: () => {
+                            // Mock progress reporting
+                        }
+                    };
+                    const mockToken: CancellationToken = {
+                        isCancellationRequested: false,
+                        onCancellationRequested: () => ({
+                            dispose: () => {
+                                // Mock disposable
+                            }
+                        })
+                    };
+                    return callback(mockProgress, mockToken);
+                }
+            );
+
+            // Mock environment deletion
+            when(mockConfigManager.deleteEnvironment(testEnvironmentId, anything())).thenResolve();
+            when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(undefined);
+
+            // Execute the command
+            await view.deleteEnvironmentCommand(testEnvironmentId);
+
+            // Verify that only kernel1 (using the deleted environment) was disposed
+            assert.strictEqual(mockKernel1.dispose.callCount, 1, 'Kernel using deleted environment should be disposed');
+            assert.strictEqual(
+                mockKernel3.dispose.callCount,
+                0,
+                'Kernel using different environment should not be disposed'
+            );
+        });
     });
 
     suite('selectEnvironmentForNotebook', () => {

src/kernels/deepnote/types.ts

@@ -194,6 +194,14 @@ export interface IDeepnoteKernelAutoSelector {
      * @param token Cancellation token to cancel the operation
      */
     rebuildController(notebook: vscode.NotebookDocument, token?: vscode.CancellationToken): Promise<void>;
+
+    /**
+     * 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.
+     * @param notebook The notebook document
+     * @param environmentId The environment ID
+     */
+    clearControllerForEnvironment(notebook: vscode.NotebookDocument, environmentId: string): void;
 }
 
 export const IDeepnoteEnvironmentManager = Symbol('IDeepnoteEnvironmentManager');

src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts

@@ -419,7 +419,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
 
     private async ensureKernelSelectedWithConfiguration(
         notebook: NotebookDocument,
-        // configuration: import('./../../kernels/deepnote/environments/deepnoteEnvironment').DeepnoteEnvironment,
         configuration: DeepnoteEnvironment,
         baseFileUri: Uri,
         notebookKey: string,
@@ -653,6 +652,28 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         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.
+     */
+    public clearControllerForEnvironment(notebook: NotebookDocument, environmentId: string): void {
+        const selectedController = this.controllerRegistration.getSelected(notebook);
+        if (!selectedController || selectedController.connection.kind !== 'startUsingDeepnoteKernel') {
+            return;
+        }
+
+        const deepnoteMetadata = selectedController.connection as DeepnoteKernelConnectionMetadata;
+        const expectedHandle = `deepnote-config-server-${environmentId}`;
+
+        if (deepnoteMetadata.serverProviderHandle.handle === expectedHandle) {
+            // Unselect the controller by setting affinity to Default
+            selectedController.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity.Default);
+            logger.info(
+                `Cleared controller for notebook ${getDisplayPath(notebook.uri)} (environment ${environmentId})`
+            );
+        }
+    }
+
     /**
      * Handle kernel selection errors with user-friendly messages and actions
      */