feat: enhance getVenvForWorkspace to prioritize persisted selections

src/managers/builtin/venvUtils.ts

@@ -61,23 +61,32 @@ export async function clearVenvCache(): Promise<void> {
 }
 
 export async function getVenvForWorkspace(fsPath: string): Promise<string | undefined> {
-    if (process.env.VIRTUAL_ENV) {
-        return process.env.VIRTUAL_ENV;
-    }
-
-    const state = await getWorkspacePersistentState();
-    const data: { [key: string]: string } | undefined = await state.get(VENV_WORKSPACE_KEY);
-    if (data) {
-        try {
+    // Check persisted user selection first — this should always take priority
+    // over process.env.VIRTUAL_ENV so that explicit selections survive reload.
+    try {
+        const state = await getWorkspacePersistentState();
+        const data: { [key: string]: string } | undefined = await state.get(VENV_WORKSPACE_KEY);
+        if (data) {
             const envPath = data[fsPath];
-            if (await fsapi.pathExists(envPath)) {
+            if (envPath && (await fsapi.pathExists(envPath))) {
                 return envPath;
             }
-            setVenvForWorkspace(fsPath, undefined);
-        } catch {
-            return undefined;
+            if (envPath) {
+                await setVenvForWorkspace(fsPath, undefined);
+            }
+        }
+    } catch {
+        // fall through to VIRTUAL_ENV check
+    }
+
+    // Fall back to VIRTUAL_ENV only if it points to a path inside this workspace.
+    if (process.env.VIRTUAL_ENV) {
+        const rel = path.relative(fsPath, process.env.VIRTUAL_ENV);
+        if (rel === '' || (!rel.startsWith('..') && !path.isAbsolute(rel))) {
+            return process.env.VIRTUAL_ENV;
         }
     }
+
     return undefined;
 }
 

src/test/managers/builtin/venvUtils.getVenvForWorkspace.unit.test.ts

@@ -0,0 +1,293 @@
+import * as assert from 'assert';
+import * as fse from 'fs-extra';
+import * as os from 'os';
+import * as path from 'path';
+import * as sinon from 'sinon';
+import * as persistentState from '../../../common/persistentState';
+import { getVenvForWorkspace, VENV_WORKSPACE_KEY } from '../../../managers/builtin/venvUtils';
+
+/**
+ * Helper that replicates the VIRTUAL_ENV subpath check from getVenvForWorkspace
+ * using a specific path module, allowing cross-platform verification.
+ */
+function isVenvInsideWorkspace(
+    fsPath: string,
+    virtualEnv: string,
+    pathModule: typeof path.posix | typeof path.win32,
+): boolean {
+    const rel = pathModule.relative(fsPath, virtualEnv);
+    return rel === '' || (!rel.startsWith('..') && !pathModule.isAbsolute(rel));
+}
+
+suite('VIRTUAL_ENV subpath check - cross-platform', () => {
+    suite('POSIX paths', () => {
+        const p = path.posix;
+
+        test('venv inside workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app/.venv', p), true);
+        });
+
+        test('venv deeply nested inside workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app/sub/dir/.venv', p), true);
+        });
+
+        test('venv equal to workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app', p), true);
+        });
+
+        test('sibling with shared prefix should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/app2/.venv', p), false);
+        });
+
+        test('venv in completely different location should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/other/place/.venv', p), false);
+        });
+
+        test('parent directory should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('/proj/app', '/proj/.venv', p), false);
+        });
+    });
+
+    suite('Windows paths', () => {
+        const p = path.win32;
+
+        test('venv inside workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app\\.venv', p), true);
+        });
+
+        test('venv deeply nested inside workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app\\sub\\dir\\.venv', p), true);
+        });
+
+        test('venv equal to workspace should match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app', p), true);
+        });
+
+        test('sibling with shared prefix should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\app2\\.venv', p), false);
+        });
+
+        test('venv on a different drive should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'D:\\proj\\app\\.venv', p), false);
+        });
+
+        test('parent directory should NOT match', () => {
+            assert.strictEqual(isVenvInsideWorkspace('C:\\proj\\app', 'C:\\proj\\.venv', p), false);
+        });
+
+        test('case-insensitive drive letters should match', () => {
+            // path.win32.relative handles case-insensitive drive letters
+            assert.strictEqual(isVenvInsideWorkspace('c:\\proj\\app', 'C:\\proj\\app\\.venv', p), true);
+        });
+
+        test('UNC path inside workspace should match', () => {
+            assert.strictEqual(
+                isVenvInsideWorkspace('\\\\server\\share\\proj', '\\\\server\\share\\proj\\.venv', p),
+                true,
+            );
+        });
+
+        test('UNC path sibling should NOT match', () => {
+            assert.strictEqual(
+                isVenvInsideWorkspace('\\\\server\\share\\proj', '\\\\server\\share\\proj2\\.venv', p),
+                false,
+            );
+        });
+    });
+});
+
+suite('getVenvForWorkspace', () => {
+    let mockState: {
+        get: sinon.SinonStub;
+        set: sinon.SinonStub;
+        clear: sinon.SinonStub;
+    };
+    let getWorkspacePersistentStateStub: sinon.SinonStub;
+    let originalVirtualEnv: string | undefined;
+    let tmpDir: string;
+
+    setup(async () => {
+        originalVirtualEnv = process.env.VIRTUAL_ENV;
+        delete process.env.VIRTUAL_ENV;
+
+        tmpDir = path.join(os.tmpdir(), `venv-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+        await fse.ensureDir(tmpDir);
+
+        mockState = {
+            get: sinon.stub(),
+            set: sinon.stub(),
+            clear: sinon.stub(),
+        };
+        getWorkspacePersistentStateStub = sinon.stub(persistentState, 'getWorkspacePersistentState');
+        getWorkspacePersistentStateStub.resolves(mockState);
+    });
+
+    teardown(async () => {
+        if (originalVirtualEnv !== undefined) {
+            process.env.VIRTUAL_ENV = originalVirtualEnv;
+        } else {
+            delete process.env.VIRTUAL_ENV;
+        }
+        sinon.restore();
+        await fse.remove(tmpDir);
+    });
+
+    test('should return persisted selection when available', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const venvPath = path.join(workspacePath, '.venv');
+        await fse.ensureDir(venvPath);
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: venvPath });
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, venvPath, 'Should return persisted venv path');
+    });
+
+    test('should return persisted selection even when VIRTUAL_ENV is set', async () => {
+        const projectA = path.join(tmpDir, 'projectA');
+        const projectB = path.join(tmpDir, 'projectB');
+        const persistedVenv = path.join(projectA, '.venv');
+        const otherVenv = path.join(projectB, '.venv');
+        await fse.ensureDir(persistedVenv);
+        await fse.ensureDir(otherVenv);
+        process.env.VIRTUAL_ENV = otherVenv;
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectA]: persistedVenv });
+
+        const result = await getVenvForWorkspace(projectA);
+        assert.strictEqual(result, persistedVenv, 'Persisted selection should take priority over VIRTUAL_ENV');
+    });
+
+    test('should fall back to VIRTUAL_ENV when no persisted selection and venv is inside workspace', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const virtualEnvPath = path.join(workspacePath, '.venv');
+        await fse.ensureDir(virtualEnvPath);
+        process.env.VIRTUAL_ENV = virtualEnvPath;
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, virtualEnvPath, 'Should use VIRTUAL_ENV when it is inside the workspace');
+    });
+
+    test('should NOT use VIRTUAL_ENV when it points to a different project', async () => {
+        const projectA = path.join(tmpDir, 'projectA');
+        const projectB = path.join(tmpDir, 'projectB');
+        const otherVenv = path.join(projectB, '.venv');
+        await fse.ensureDir(projectA);
+        await fse.ensureDir(otherVenv);
+        process.env.VIRTUAL_ENV = otherVenv;
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
+
+        const result = await getVenvForWorkspace(projectA);
+        assert.strictEqual(result, undefined, 'Should NOT use VIRTUAL_ENV from a different project');
+    });
+
+    test('should clear stale persisted path when venv no longer exists', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const staleVenv = path.join(workspacePath, '.venv-old');
+        await fse.ensureDir(workspacePath);
+        // Note: staleVenv directory does NOT exist on disk
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: staleVenv });
+
+        const result = await getVenvForWorkspace(workspacePath);
+
+        assert.strictEqual(result, undefined, 'Should return undefined for stale path');
+        assert.ok(mockState.set.called, 'Should clear the stale entry from persistent state');
+        const setArgs = mockState.set.firstCall.args;
+        assert.strictEqual(setArgs[0], VENV_WORKSPACE_KEY, 'Should update the correct key');
+        assert.strictEqual(setArgs[1][workspacePath], undefined, 'Should have removed the stale workspace entry');
+    });
+
+    test('should return undefined when no persisted selection and no VIRTUAL_ENV', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        await fse.ensureDir(workspacePath);
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, undefined, 'Should return undefined when nothing is available');
+    });
+
+    test('should return undefined when persisted data has no entry for this workspace', async () => {
+        const projectA = path.join(tmpDir, 'projectA');
+        const projectB = path.join(tmpDir, 'projectB');
+        await fse.ensureDir(projectA);
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectB]: '/some/path' });
+
+        const result = await getVenvForWorkspace(projectA);
+        assert.strictEqual(result, undefined, 'Should return undefined when no entry for this workspace');
+    });
+
+    test('should fall back to VIRTUAL_ENV when data access throws inside try block', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const virtualEnvPath = path.join(workspacePath, '.venv');
+        await fse.ensureDir(virtualEnvPath);
+        process.env.VIRTUAL_ENV = virtualEnvPath;
+
+        // Return data object with a getter that throws when accessing the workspace key
+        const badData: Record<string, string> = {};
+        Object.defineProperty(badData, workspacePath, {
+            get() {
+                throw new Error('corrupted data');
+            },
+            enumerable: true,
+            configurable: true,
+        });
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(badData);
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when try block throws');
+    });
+
+    test('should not clear state when no envPath exists for the workspace key', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        await fse.ensureDir(workspacePath);
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ other: '/some/path' });
+
+        await getVenvForWorkspace(workspacePath);
+
+        assert.ok(!mockState.set.called, 'Should not call set when there is no stale entry to clear');
+    });
+
+    test('should fall back to VIRTUAL_ENV when state.get rejects', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const virtualEnvPath = path.join(workspacePath, '.venv');
+        await fse.ensureDir(virtualEnvPath);
+        process.env.VIRTUAL_ENV = virtualEnvPath;
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).rejects(new Error('storage corrupted'));
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when state.get rejects');
+    });
+
+    test('should fall back to VIRTUAL_ENV when getWorkspacePersistentState rejects', async () => {
+        const workspacePath = path.join(tmpDir, 'projectA');
+        const virtualEnvPath = path.join(workspacePath, '.venv');
+        await fse.ensureDir(virtualEnvPath);
+        process.env.VIRTUAL_ENV = virtualEnvPath;
+
+        getWorkspacePersistentStateStub.rejects(new Error('persistent state unavailable'));
+
+        const result = await getVenvForWorkspace(workspacePath);
+        assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when persistent state fails');
+    });
+
+    test('should NOT use VIRTUAL_ENV for sibling path with shared prefix', async () => {
+        const projectA = path.join(tmpDir, 'app');
+        const siblingVenv = path.join(tmpDir, 'app2', '.venv');
+        await fse.ensureDir(projectA);
+        await fse.ensureDir(siblingVenv);
+        process.env.VIRTUAL_ENV = siblingVenv;
+
+        mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
+
+        const result = await getVenvForWorkspace(projectA);
+        assert.strictEqual(result, undefined, 'Should NOT match sibling path with shared prefix');
+    });
+});