Preserve Jupyter kernels during notebook file rename

src/kernels/kernelProvider.base.ts

@@ -150,6 +150,32 @@
     protected disposeOldKernel(notebook: NotebookDocument, reason: 'notebookClosed' | 'createNewKernel') {
         const kernelToDispose = this.kernelsByNotebook.get(notebook);
         if (kernelToDispose) {
+            // Check if this kernel is marked for migration due to a file rename
+            const migrationTarget = (kernelToDispose.kernel as any)._migrationTarget;
+            if (migrationTarget && reason === 'notebookClosed') {
+                logger.debug(
+                    `Kernel ${kernelToDispose.kernel.id} associated with ${getDisplayPath(
+                        notebook.uri
+                    )} is marked for migration to ${migrationTarget}, attempting to migrate instead of disposing`
+                );
+
+                // Attempt to find the new notebook document
+                const targetNotebook = workspace.notebookDocuments.find(
+                    (doc) => doc.uri.toString() === migrationTarget
+                );
+                if (targetNotebook) {
+                    logger.debug(`Found target notebook for migration: ${getDisplayPath(targetNotebook.uri)}`);
+
+                    // Migrate the kernel to the new notebook
+                    this.migrateKernel(notebook, targetNotebook, kernelToDispose.kernel, kernelToDispose.options);
+                    return;
+                } else {
+                    logger.debug(`Target notebook not found for migration, proceeding with disposal`);
+                    // Clear the migration flag and proceed with normal disposal
+                    delete (kernelToDispose.kernel as any)._migrationTarget;
+                }
+            }
+
             logger.debug(
                 `Disposing kernel associated with ${getDisplayPath(notebook.uri)}, isClosed=${
                     notebook.isClosed
@@ -166,6 +192,42 @@
         this.kernelsByNotebook.delete(notebook);
     }
 
+    /**
+     * Migrates a kernel from one notebook to another during file rename operations.
+     * This preserves the kernel session instead of disposing and recreating it.
+     */
+    protected migrateKernel(
+        oldNotebook: NotebookDocument,
+        newNotebook: NotebookDocument,
+        kernel: IKernel,
+        options: KernelOptions
+    ) {
+        logger.debug(
+            `Migrating kernel ${kernel.id} from ${getDisplayPath(oldNotebook.uri)} to ${getDisplayPath(
+                newNotebook.uri
+            )}`
+        );
+
+        // Remove the old mappings
+        this.kernelsByNotebook.delete(oldNotebook);
+
+        // Update the kernel's internal notebook reference if possible
+        if ((kernel as any).notebook) {
+            (kernel as any).notebook = newNotebook;
+        }
+        if ((kernel as any).uri) {
+            (kernel as any).uri = newNotebook.uri;
+        }
+
+        // Create new mappings for the target notebook
+        this.kernelsByNotebook.set(newNotebook, { options, kernel });
+
+        // Clear the migration flag
+        delete (kernel as any)._migrationTarget;
+
+        logger.debug(`Successfully migrated kernel ${kernel.id} to ${getDisplayPath(newNotebook.uri)}`);
+    }
+
     protected handleServerRemoval(servers: JupyterServerProviderHandle[]) {
         workspace.notebookDocuments.forEach((document) => {
             const kernel = this.kernelsByNotebook.get(document);

src/kernels/kernelProvider.base.unit.test.ts

@@ -0,0 +1,160 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { assert } from 'chai';
+import { anything, instance, mock, when } from 'ts-mockito';
+import { NotebookDocument, Uri } from 'vscode';
+import { IAsyncDisposableRegistry, IDisposableRegistry } from '../platform/common/types';
+import { IKernel, KernelOptions } from './types';
+import { BaseCoreKernelProvider } from './kernelProvider.base';
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+// Test interface to add migration properties
+interface ITestKernel extends IKernel {
+    _migrationTarget?: string;
+}
+
+// Create a concrete test class since BaseCoreKernelProvider is abstract
+class TestKernelProvider extends BaseCoreKernelProvider {
+    public getOrCreate(_notebook: NotebookDocument, _options: KernelOptions): IKernel {
+        throw new Error('Not implemented for test');
+    }
+
+    // Expose protected methods for testing
+    public testDisposeOldKernel(notebook: NotebookDocument, reason: 'notebookClosed' | 'createNewKernel') {
+        return this.disposeOldKernel(notebook, reason);
+    }
+
+    public testMigrateKernel(
+        oldNotebook: NotebookDocument,
+        newNotebook: NotebookDocument,
+        kernel: IKernel,
+        options: KernelOptions
+    ) {
+        return this.migrateKernel(oldNotebook, newNotebook, kernel, options);
+    }
+
+    // Override storeKernel to make it accessible for testing
+    public testStoreKernel(notebook: NotebookDocument, options: KernelOptions, kernel: IKernel) {
+        return this.storeKernel(notebook, options, kernel);
+    }
+}
+
+suite('BaseCoreKernelProvider Kernel Migration', () => {
+    let kernelProvider: TestKernelProvider;
+    let asyncDisposables: IAsyncDisposableRegistry;
+    let disposables: IDisposableRegistry;
+
+    setup(() => {
+        asyncDisposables = mock<IAsyncDisposableRegistry>();
+        disposables = mock<IDisposableRegistry>();
+
+        when(asyncDisposables.push(anything())).thenReturn();
+        when(disposables.push(anything())).thenReturn();
+
+        kernelProvider = new TestKernelProvider(instance(asyncDisposables), instance(disposables));
+    });
+
+    test('Should migrate kernel when marked for migration', () => {
+        const oldNotebook = mock<NotebookDocument>();
+        const newNotebook = mock<NotebookDocument>();
+        // Create a simple mock kernel object
+        const kernel = {
+            id: 'test-kernel-id',
+            _migrationTarget: undefined
+        } as ITestKernel;
+        const options = {} as KernelOptions;
+
+        const oldUri = Uri.file('/test/old.ipynb');
+        const newUri = Uri.file('/test/new.ipynb');
+
+        when(oldNotebook.uri).thenReturn(oldUri);
+        when(newNotebook.uri).thenReturn(newUri);
+
+        // Store a kernel first
+        kernelProvider.testStoreKernel(instance(oldNotebook), options, kernel);
+
+        // Mark the kernel for migration
+        kernel._migrationTarget = newUri.toString();
+
+        // Test the migration directly by calling migrateKernel
+        kernelProvider.testMigrateKernel(instance(oldNotebook), instance(newNotebook), kernel, options);
+
+        // Assert - kernel should be accessible via new notebook
+        const migratedKernel = kernelProvider.get(instance(newNotebook));
+        assert.strictEqual(migratedKernel, kernel);
+
+        // Assert - migration flag should be cleared
+        assert.isUndefined(kernel._migrationTarget, 'Migration flag should be cleared');
+    });
+
+    test('Should not migrate kernel when not marked for migration', () => {
+        const oldNotebook = mock<NotebookDocument>();
+        const kernel = {
+            id: 'test-kernel-id',
+            dispose: () => Promise.resolve()
+        } as IKernel;
+        const options = {} as KernelOptions;
+
+        const oldUri = Uri.file('/test/old.ipynb');
+        when(oldNotebook.uri).thenReturn(oldUri);
+        when(oldNotebook.isClosed).thenReturn(true);
+
+        // Store a kernel first
+        kernelProvider.testStoreKernel(instance(oldNotebook), options, kernel);
+
+        // Don't mark for migration
+
+        // Act - dispose old kernel (should dispose normally)
+        kernelProvider.testDisposeOldKernel(instance(oldNotebook), 'notebookClosed');
+
+        // Assert - kernel should no longer be in the provider
+        const disposedKernel = kernelProvider.get(instance(oldNotebook));
+        assert.isUndefined(disposedKernel);
+    });
+
+    test('Should directly migrate kernel when migrateKernel is called', () => {
+        const oldNotebook = mock<NotebookDocument>();
+        const newNotebook = mock<NotebookDocument>();
+        const kernel = {
+            id: 'test-kernel-id',
+            _migrationTarget: undefined,
+            notebook: undefined as any,
+            uri: undefined as any
+        } as ITestKernel & { notebook: any; uri: any };
+        const options = {} as KernelOptions;
+
+        const oldUri = Uri.file('/test/old.ipynb');
+        const newUri = Uri.file('/test/new.ipynb');
+
+        when(oldNotebook.uri).thenReturn(oldUri);
+        when(newNotebook.uri).thenReturn(newUri);
+
+        // Store a kernel first
+        kernelProvider.testStoreKernel(instance(oldNotebook), options, kernel);
+
+        // Mark for migration and set up kernel properties
+        kernel._migrationTarget = newUri.toString();
+        kernel.notebook = instance(oldNotebook);
+        kernel.uri = oldUri;
+
+        // Act - directly call migrate
+        kernelProvider.testMigrateKernel(instance(oldNotebook), instance(newNotebook), kernel, options);
+
+        // Assert - kernel should be accessible via new notebook
+        const migratedKernel = kernelProvider.get(instance(newNotebook));
+        assert.strictEqual(migratedKernel, kernel);
+
+        // Assert - kernel should no longer be accessible via old notebook
+        const oldKernel = kernelProvider.get(instance(oldNotebook));
+        assert.isUndefined(oldKernel);
+
+        // Assert - migration flag should be cleared
+        assert.isUndefined(kernel._migrationTarget, 'Migration flag should be cleared');
+
+        // Assert - kernel properties should be updated
+        assert.strictEqual(kernel.notebook, instance(newNotebook));
+        assert.strictEqual(kernel.uri, newUri);
+    });
+});

src/kernels/notebookRenameHandler.ts

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { injectable } from 'inversify';
+import { Disposable, Uri, workspace } from 'vscode';
+import { logger } from '../platform/logging';
+import { getDisplayPath } from '../platform/common/platform/fs-paths';
+import { IDisposableRegistry } from '../platform/common/types';
+import { IExtensionSyncActivationService } from '../platform/activation/types';
+import { IServiceContainer } from '../platform/ioc/types';
+import { IKernelProvider } from './types';
+
+/**
+ * Handles notebook file renames to preserve kernel sessions.
+ * When a notebook is renamed, transfers the kernel from the old URI to the new URI
+ * instead of disposing it.
+ */
+@injectable()
+export class NotebookRenameHandler implements IExtensionSyncActivationService {
+    private readonly disposables: Disposable[] = [];
+    private kernelProvider: IKernelProvider | undefined;
+
+    constructor(
+        private readonly serviceContainer: IServiceContainer,
+        private readonly disposableRegistry: IDisposableRegistry
+    ) {
+        this.disposableRegistry.push(this);
+    }
+
+    public activate(): void {
+        // Listen for file rename events
+        workspace.onWillRenameFiles(
+            async (event) => {
+                for (const file of event.files) {
+                    // Check if this is a notebook file rename
+                    if (file.oldUri.fsPath.endsWith('.ipynb') && file.newUri.fsPath.endsWith('.ipynb')) {
+                        await this.handleNotebookRename(file.oldUri, file.newUri);
+                    }
+                }
+            },
+            this,
+            this.disposables
+        );
+
+        this.disposableRegistry.push(...this.disposables);
+    }
+
+    private async handleNotebookRename(oldUri: Uri, newUri: Uri): Promise<void> {
+        try {
+            logger.debug(`Handling notebook rename from ${getDisplayPath(oldUri)} to ${getDisplayPath(newUri)}`);
+
+            // Get the kernel provider
+            if (!this.kernelProvider) {
+                this.kernelProvider = this.serviceContainer.get<IKernelProvider>(IKernelProvider);
+            }
+
+            // Get the existing kernel for the old URI
+            const existingKernel = this.kernelProvider.get(oldUri);
+            if (!existingKernel) {
+                logger.debug(`No kernel found for ${getDisplayPath(oldUri)}, nothing to migrate`);
+                return;
+            }
+
+            logger.debug(`Found kernel ${existingKernel.id} for ${getDisplayPath(oldUri)}, preparing to migrate`);
+
+            // Store kernel information for migration
+            const kernelInfo = (this.kernelProvider as any).getInternal?.(
+                workspace.notebookDocuments.find((doc) => doc.uri.toString() === oldUri.toString())
+            );
+            if (!kernelInfo) {
+                logger.debug(`No internal kernel info found for ${getDisplayPath(oldUri)}`);
+                return;
+            }
+
+            // Flag this kernel for migration (we'll handle the actual migration in the onDidRenameFiles event)
+            (existingKernel as any)._migrationTarget = newUri.toString();
+            logger.debug(`Marked kernel ${existingKernel.id} for migration to ${getDisplayPath(newUri)}`);
+        } catch (error) {
+            logger.error(
+                `Error handling notebook rename from ${getDisplayPath(oldUri)} to ${getDisplayPath(newUri)}`,
+                error
+            );
+        }
+    }
+
+    public dispose(): void {
+        this.disposables.forEach((d) => d.dispose());
+    }
+}