Skip to content

Commit 5553a6e

Browse files
committed
feat: enhance getVenvForWorkspace to prioritize persisted selections and handle stale paths
1 parent e305eff commit 5553a6e

File tree

2 files changed

+179
-7
lines changed

2 files changed

+179
-7
lines changed

src/managers/builtin/venvUtils.ts

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -61,23 +61,29 @@ export async function clearVenvCache(): Promise<void> {
6161
}
6262

6363
export async function getVenvForWorkspace(fsPath: string): Promise<string | undefined> {
64-
if (process.env.VIRTUAL_ENV) {
65-
return process.env.VIRTUAL_ENV;
66-
}
67-
64+
// Check persisted user selection first — this should always take priority
65+
// over process.env.VIRTUAL_ENV so that explicit selections survive reload.
6866
const state = await getWorkspacePersistentState();
6967
const data: { [key: string]: string } | undefined = await state.get(VENV_WORKSPACE_KEY);
7068
if (data) {
7169
try {
7270
const envPath = data[fsPath];
73-
if (await fsapi.pathExists(envPath)) {
71+
if (envPath && (await fsapi.pathExists(envPath))) {
7472
return envPath;
7573
}
76-
setVenvForWorkspace(fsPath, undefined);
74+
if (envPath) {
75+
await setVenvForWorkspace(fsPath, undefined);
76+
}
7777
} catch {
78-
return undefined;
78+
// fall through to VIRTUAL_ENV check
7979
}
8080
}
81+
82+
// Fall back to VIRTUAL_ENV only if it points to a path inside this workspace.
83+
if (process.env.VIRTUAL_ENV && normalizePath(process.env.VIRTUAL_ENV).startsWith(normalizePath(fsPath))) {
84+
return process.env.VIRTUAL_ENV;
85+
}
86+
8187
return undefined;
8288
}
8389

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
import * as assert from 'assert';
2+
import * as fse from 'fs-extra';
3+
import * as os from 'os';
4+
import * as path from 'path';
5+
import * as sinon from 'sinon';
6+
import * as persistentState from '../../../common/persistentState';
7+
import { getVenvForWorkspace, VENV_WORKSPACE_KEY } from '../../../managers/builtin/venvUtils';
8+
9+
suite('getVenvForWorkspace', () => {
10+
let mockState: {
11+
get: sinon.SinonStub;
12+
set: sinon.SinonStub;
13+
clear: sinon.SinonStub;
14+
};
15+
let getWorkspacePersistentStateStub: sinon.SinonStub;
16+
let originalVirtualEnv: string | undefined;
17+
let tmpDir: string;
18+
19+
setup(async () => {
20+
originalVirtualEnv = process.env.VIRTUAL_ENV;
21+
delete process.env.VIRTUAL_ENV;
22+
23+
tmpDir = path.join(os.tmpdir(), `venv-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
24+
await fse.ensureDir(tmpDir);
25+
26+
mockState = {
27+
get: sinon.stub(),
28+
set: sinon.stub(),
29+
clear: sinon.stub(),
30+
};
31+
getWorkspacePersistentStateStub = sinon.stub(persistentState, 'getWorkspacePersistentState');
32+
getWorkspacePersistentStateStub.resolves(mockState);
33+
});
34+
35+
teardown(async () => {
36+
if (originalVirtualEnv !== undefined) {
37+
process.env.VIRTUAL_ENV = originalVirtualEnv;
38+
} else {
39+
delete process.env.VIRTUAL_ENV;
40+
}
41+
sinon.restore();
42+
await fse.remove(tmpDir);
43+
});
44+
45+
test('should return persisted selection when available', async () => {
46+
const workspacePath = path.join(tmpDir, 'projectA');
47+
const venvPath = path.join(workspacePath, '.venv');
48+
await fse.ensureDir(venvPath);
49+
50+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: venvPath });
51+
52+
const result = await getVenvForWorkspace(workspacePath);
53+
assert.strictEqual(result, venvPath, 'Should return persisted venv path');
54+
});
55+
56+
test('should return persisted selection even when VIRTUAL_ENV is set', async () => {
57+
const projectA = path.join(tmpDir, 'projectA');
58+
const projectB = path.join(tmpDir, 'projectB');
59+
const persistedVenv = path.join(projectA, '.venv');
60+
const otherVenv = path.join(projectB, '.venv');
61+
await fse.ensureDir(persistedVenv);
62+
await fse.ensureDir(otherVenv);
63+
process.env.VIRTUAL_ENV = otherVenv;
64+
65+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectA]: persistedVenv });
66+
67+
const result = await getVenvForWorkspace(projectA);
68+
assert.strictEqual(result, persistedVenv, 'Persisted selection should take priority over VIRTUAL_ENV');
69+
});
70+
71+
test('should fall back to VIRTUAL_ENV when no persisted selection and venv is inside workspace', async () => {
72+
const workspacePath = path.join(tmpDir, 'projectA');
73+
const virtualEnvPath = path.join(workspacePath, '.venv');
74+
await fse.ensureDir(virtualEnvPath);
75+
process.env.VIRTUAL_ENV = virtualEnvPath;
76+
77+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
78+
79+
const result = await getVenvForWorkspace(workspacePath);
80+
assert.strictEqual(result, virtualEnvPath, 'Should use VIRTUAL_ENV when it is inside the workspace');
81+
});
82+
83+
test('should NOT use VIRTUAL_ENV when it points to a different project', async () => {
84+
const projectA = path.join(tmpDir, 'projectA');
85+
const projectB = path.join(tmpDir, 'projectB');
86+
const otherVenv = path.join(projectB, '.venv');
87+
await fse.ensureDir(projectA);
88+
await fse.ensureDir(otherVenv);
89+
process.env.VIRTUAL_ENV = otherVenv;
90+
91+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
92+
93+
const result = await getVenvForWorkspace(projectA);
94+
assert.strictEqual(result, undefined, 'Should NOT use VIRTUAL_ENV from a different project');
95+
});
96+
97+
test('should clear stale persisted path when venv no longer exists', async () => {
98+
const workspacePath = path.join(tmpDir, 'projectA');
99+
const staleVenv = path.join(workspacePath, '.venv-old');
100+
await fse.ensureDir(workspacePath);
101+
// Note: staleVenv directory does NOT exist on disk
102+
103+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [workspacePath]: staleVenv });
104+
105+
const result = await getVenvForWorkspace(workspacePath);
106+
107+
assert.strictEqual(result, undefined, 'Should return undefined for stale path');
108+
assert.ok(mockState.set.called, 'Should clear the stale entry from persistent state');
109+
const setArgs = mockState.set.firstCall.args;
110+
assert.strictEqual(setArgs[0], VENV_WORKSPACE_KEY, 'Should update the correct key');
111+
assert.strictEqual(setArgs[1][workspacePath], undefined, 'Should have removed the stale workspace entry');
112+
});
113+
114+
test('should return undefined when no persisted selection and no VIRTUAL_ENV', async () => {
115+
const workspacePath = path.join(tmpDir, 'projectA');
116+
await fse.ensureDir(workspacePath);
117+
118+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(undefined);
119+
120+
const result = await getVenvForWorkspace(workspacePath);
121+
assert.strictEqual(result, undefined, 'Should return undefined when nothing is available');
122+
});
123+
124+
test('should return undefined when persisted data has no entry for this workspace', async () => {
125+
const projectA = path.join(tmpDir, 'projectA');
126+
const projectB = path.join(tmpDir, 'projectB');
127+
await fse.ensureDir(projectA);
128+
129+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ [projectB]: '/some/path' });
130+
131+
const result = await getVenvForWorkspace(projectA);
132+
assert.strictEqual(result, undefined, 'Should return undefined when no entry for this workspace');
133+
});
134+
135+
test('should fall back to VIRTUAL_ENV when data access throws inside try block', async () => {
136+
const workspacePath = path.join(tmpDir, 'projectA');
137+
const virtualEnvPath = path.join(workspacePath, '.venv');
138+
await fse.ensureDir(virtualEnvPath);
139+
process.env.VIRTUAL_ENV = virtualEnvPath;
140+
141+
// Return data object with a getter that throws when accessing the workspace key
142+
const badData: Record<string, string> = {};
143+
Object.defineProperty(badData, workspacePath, {
144+
get() {
145+
throw new Error('corrupted data');
146+
},
147+
enumerable: true,
148+
configurable: true,
149+
});
150+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves(badData);
151+
152+
const result = await getVenvForWorkspace(workspacePath);
153+
assert.strictEqual(result, virtualEnvPath, 'Should fall back to VIRTUAL_ENV when try block throws');
154+
});
155+
156+
test('should not clear state when no envPath exists for the workspace key', async () => {
157+
const workspacePath = path.join(tmpDir, 'projectA');
158+
await fse.ensureDir(workspacePath);
159+
160+
mockState.get.withArgs(VENV_WORKSPACE_KEY).resolves({ other: '/some/path' });
161+
162+
await getVenvForWorkspace(workspacePath);
163+
164+
assert.ok(!mockState.set.called, 'Should not call set when there is no stale entry to clear');
165+
});
166+
});

0 commit comments

Comments
 (0)