fix: migrate venv paths for cross-editor compatibility and legacy naming (#188)

saltenasl · web-flow · commit db84c276233b · 2025-11-05T08:51:41.000-08:00
diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts
@@ -82,6 +82,18 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         const venvKey = venvPath.fsPath;
 
         logger.info(`Ensuring virtual environment at ${venvKey}`);
+        logger.info(`Base interpreter: ${baseInterpreter.uri.fsPath}`);
+
+        // Validate that venv path is in current globalStorage (not from a different editor like VS Code)
+        const expectedStoragePrefix = this.context.globalStorageUri.fsPath;
+        if (!venvKey.startsWith(expectedStoragePrefix)) {
+            const error = new Error(
+                `Venv path mismatch! Expected venv under ${expectedStoragePrefix} but got ${venvKey}. ` +
+                    `This might happen if the notebook was previously used in a different editor (VS Code vs Cursor).`
+            );
+            logger.error(error.message);
+            throw error;
+        }
 
         // Wait for any pending installation for this venv to complete
         const pendingInstall = this.pendingInstallations.get(venvKey);
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts
@@ -1,4 +1,5 @@
 import { inject, injectable, named } from 'inversify';
+import * as path from '../../../platform/vscode-path/path';
 import { CancellationToken, EventEmitter, l10n, Uri } from 'vscode';
 import { IExtensionSyncActivationService } from '../../../platform/activation/types';
 import { Cancellation } from '../../../platform/common/cancellation';
@@ -52,10 +53,40 @@ export class DeepnoteEnvironmentManager implements IExtensionSyncActivationServi
             const configs = await this.storage.loadEnvironments();
             this.environments.clear();
 
+            let needsMigration = false;
+
             for (const config of configs) {
+                const venvDirName = path.basename(config.venvPath.fsPath);
+
+                // Check if venv path is under current globalStorage
+                const expectedVenvParent = Uri.joinPath(this.context.globalStorageUri, 'deepnote-venvs').fsPath;
+                const actualVenvParent = path.dirname(config.venvPath.fsPath);
+                const isInCorrectStorage = actualVenvParent === expectedVenvParent;
+
+                // Check if directory name matches the environment ID and is in correct storage
+                const isExpectedPath = venvDirName === config.id && isInCorrectStorage;
+                const needsPathMigration = !isExpectedPath;
+
+                if (needsPathMigration) {
+                    logger.info(
+                        `Migrating environment "${config.name}" from ${config.venvPath.fsPath} to ID-based path`
+                    );
+
+                    config.venvPath = Uri.joinPath(this.context.globalStorageUri, 'deepnote-venvs', config.id);
+                    config.toolkitVersion = undefined;
+
+                    logger.info(`New venv path: ${config.venvPath.fsPath} (will be recreated on next use)`);
+                    needsMigration = true;
+                }
+
                 this.environments.set(config.id, config);
             }
 
+            if (needsMigration) {
+                logger.info('Saving migrated environments to storage');
+                await this.persistEnvironments();
+            }
+
             logger.info(`Initialized environment manager with ${this.environments.size} environments`);
 
             // Fire event to notify tree view of loaded environments
diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentManager.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentManager.unit.test.ts
@@ -266,4 +266,189 @@ suite('DeepnoteEnvironmentManager', () => {
             // Should not throw
         });
     });
+
+    suite('environment migration', () => {
+        test('should migrate hash-based venv paths to UUID-based paths', async () => {
+            const oldHashBasedConfig = {
+                id: 'abcd1234-5678-90ab-cdef-123456789012',
+                name: 'Old Hash Config',
+                pythonInterpreter: testInterpreter,
+                venvPath: Uri.file('/global/storage/deepnote-venvs/venv_7626587d-1.0.0'),
+                createdAt: new Date(),
+                lastUsedAt: new Date()
+            };
+
+            when(mockStorage.loadEnvironments()).thenResolve([oldHashBasedConfig]);
+            when(mockStorage.saveEnvironments(anything())).thenResolve();
+            when(mockContext.globalStorageUri).thenReturn(Uri.file('/global/storage'));
+
+            manager.activate();
+            await manager.waitForInitialization();
+
+            const configs = manager.listEnvironments();
+            assert.strictEqual(configs.length, 1);
+
+            // Should have migrated to UUID-based path
+            assert.strictEqual(
+                configs[0].venvPath.fsPath,
+                '/global/storage/deepnote-venvs/abcd1234-5678-90ab-cdef-123456789012'
+            );
+
+            // Should clear toolkit version to force reinstallation
+            assert.isUndefined(configs[0].toolkitVersion);
+
+            // Should have saved the migration
+            verify(mockStorage.saveEnvironments(anything())).once();
+        });
+
+        test('should migrate VS Code storage paths to Cursor storage paths', async () => {
+            const vsCodeConfig = {
+                id: 'cursor-env-id',
+                name: 'VS Code Environment',
+                pythonInterpreter: testInterpreter,
+                venvPath: Uri.file(
+                    '/Library/Application Support/Code/User/globalStorage/deepnote.vscode-deepnote/deepnote-venvs/cursor-env-id'
+                ),
+                createdAt: new Date(),
+                lastUsedAt: new Date(),
+                toolkitVersion: '1.0.0'
+            };
+
+            when(mockStorage.loadEnvironments()).thenResolve([vsCodeConfig]);
+            when(mockStorage.saveEnvironments(anything())).thenResolve();
+            when(mockContext.globalStorageUri).thenReturn(
+                Uri.file('/Library/Application Support/Cursor/User/globalStorage/deepnote.vscode-deepnote')
+            );
+
+            manager.activate();
+            await manager.waitForInitialization();
+
+            const configs = manager.listEnvironments();
+            assert.strictEqual(configs.length, 1);
+
+            // Should have migrated to Cursor storage
+            assert.match(configs[0].venvPath.fsPath, /Cursor.*deepnote-venvs\/cursor-env-id$/);
+
+            // Should clear toolkit version to force reinstallation
+            assert.isUndefined(configs[0].toolkitVersion);
+
+            verify(mockStorage.saveEnvironments(anything())).once();
+        });
+
+        test('should not migrate environments with correct ID-based paths in correct storage', async () => {
+            const testDate = new Date();
+            const correctConfig = {
+                id: '12345678-1234-1234-1234-123456789abc',
+                name: 'Correct Config',
+                pythonInterpreter: testInterpreter,
+                venvPath: Uri.file('/global/storage/deepnote-venvs/12345678-1234-1234-1234-123456789abc'),
+                createdAt: testDate,
+                lastUsedAt: testDate,
+                toolkitVersion: '1.0.0',
+                packages: []
+            };
+
+            when(mockStorage.loadEnvironments()).thenResolve([correctConfig]);
+            when(mockStorage.saveEnvironments(anything())).thenResolve();
+            when(mockContext.globalStorageUri).thenReturn(Uri.file('/global/storage'));
+
+            manager.activate();
+            await manager.waitForInitialization();
+
+            const configs = manager.listEnvironments();
+            assert.strictEqual(configs.length, 1);
+
+            // Path should remain unchanged
+            assert.strictEqual(
+                configs[0].venvPath.fsPath,
+                '/global/storage/deepnote-venvs/12345678-1234-1234-1234-123456789abc'
+            );
+
+            // ID and name should be preserved
+            assert.strictEqual(configs[0].id, '12345678-1234-1234-1234-123456789abc');
+            assert.strictEqual(configs[0].name, 'Correct Config');
+
+            // Should NOT have saved (no migration needed)
+            verify(mockStorage.saveEnvironments(anything())).never();
+        });
+
+        test('should not migrate environments with non-UUID IDs when path already matches', async () => {
+            const testDate = new Date();
+            const customIdConfig = {
+                id: 'my-custom-env-id',
+                name: 'Custom ID Environment',
+                pythonInterpreter: testInterpreter,
+                venvPath: Uri.file('/global/storage/deepnote-venvs/my-custom-env-id'),
+                createdAt: testDate,
+                lastUsedAt: testDate,
+                toolkitVersion: '1.0.0'
+            };
+
+            when(mockStorage.loadEnvironments()).thenResolve([customIdConfig]);
+            when(mockStorage.saveEnvironments(anything())).thenResolve();
+            when(mockContext.globalStorageUri).thenReturn(Uri.file('/global/storage'));
+
+            manager.activate();
+            await manager.waitForInitialization();
+
+            const configs = manager.listEnvironments();
+            assert.strictEqual(configs.length, 1);
+
+            // Path should remain unchanged
+            assert.strictEqual(configs[0].venvPath.fsPath, '/global/storage/deepnote-venvs/my-custom-env-id');
+
+            // Toolkit version should NOT be cleared
+            assert.strictEqual(configs[0].toolkitVersion, '1.0.0');
+
+            // Should NOT have saved (no migration needed)
+            verify(mockStorage.saveEnvironments(anything())).never();
+        });
+
+        test('should migrate multiple environments at once', async () => {
+            const configs = [
+                {
+                    id: 'uuid1',
+                    name: 'Hash Config',
+                    pythonInterpreter: testInterpreter,
+                    venvPath: Uri.file('/global/storage/deepnote-venvs/venv_abc123-1.0.0'),
+                    createdAt: new Date(),
+                    lastUsedAt: new Date()
+                },
+                {
+                    id: 'uuid2',
+                    name: 'VS Code Config',
+                    pythonInterpreter: testInterpreter,
+                    venvPath: Uri.file('/Code/globalStorage/deepnote-venvs/uuid2'),
+                    createdAt: new Date(),
+                    lastUsedAt: new Date()
+                },
+                {
+                    id: 'uuid3',
+                    name: 'Correct Config',
+                    pythonInterpreter: testInterpreter,
+                    venvPath: Uri.file('/global/storage/deepnote-venvs/uuid3'),
+                    createdAt: new Date(),
+                    lastUsedAt: new Date()
+                }
+            ];
+
+            when(mockStorage.loadEnvironments()).thenResolve(configs);
+            when(mockStorage.saveEnvironments(anything())).thenResolve();
+            when(mockContext.globalStorageUri).thenReturn(Uri.file('/global/storage'));
+
+            manager.activate();
+            await manager.waitForInitialization();
+
+            const loaded = manager.listEnvironments();
+            assert.strictEqual(loaded.length, 3);
+
+            // First two should be migrated
+            assert.strictEqual(loaded[0].venvPath.fsPath, '/global/storage/deepnote-venvs/uuid1');
+            assert.strictEqual(loaded[1].venvPath.fsPath, '/global/storage/deepnote-venvs/uuid2');
+            // Third should remain unchanged
+            assert.strictEqual(loaded[2].venvPath.fsPath, '/global/storage/deepnote-venvs/uuid3');
+
+            verify(mockStorage.saveEnvironments(anything())).once();
+        });
+    });
 });
diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts
@@ -508,9 +508,30 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
         const existingEnvironmentId = this.notebookEnvironmentsIds.get(notebookKey);
 
         if (existingEnvironmentId != null && existingController != null && existingEnvironmentId === configuration.id) {
-            logger.info(`Existing controller found for notebook ${getDisplayPath(notebook.uri)}, selecting it`);
-            await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken);
-            return;
+            logger.info(`Existing controller found for notebook ${getDisplayPath(notebook.uri)}, verifying connection`);
+
+            // Verify the controller's interpreter path matches the expected venv path
+            // This handles cases where notebooks were used in VS Code and now opened in Cursor
+            const existingInterpreter = existingController.connection.interpreter;
+            if (existingInterpreter) {
+                const expectedInterpreter =
+                    process.platform === 'win32'
+                        ? Uri.joinPath(configuration.venvPath, 'Scripts', 'python.exe')
+                        : Uri.joinPath(configuration.venvPath, 'bin', 'python');
+
+                if (existingInterpreter.uri.fsPath !== expectedInterpreter.fsPath) {
+                    logger.warn(
+                        `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingInterpreter.uri.fsPath}. Recreating controller.`
+                    );
+                    // Dispose old controller and recreate it
+                    existingController.dispose();
+                    this.notebookControllers.delete(notebookKey);
+                } else {
+                    logger.info(`Controller verified, selecting it`);
+                    await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken);
+                    return;
+                }
+            }
         }
 
         // Ensure server is running (startServer is idempotent - returns early if already running)
@@ -582,6 +603,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector,
                 ? Uri.joinPath(configuration.venvPath, 'Scripts', 'python.exe')
                 : Uri.joinPath(configuration.venvPath, 'bin', 'python');
 
+        logger.info(`Using venv path: ${configuration.venvPath.fsPath}`);
+        logger.info(`Venv interpreter path: ${venvInterpreter.fsPath}`);
+
         // CRITICAL: Use unique notebook-based ID (includes query with notebook ID)
         // This ensures each notebook gets its own controller/kernel, even within the same project.
         // When switching environments, addOrUpdate will call updateConnection() on the existing
