fix: ensure env managers are using workspace persistent state (#455)

Author: karthiknadig

src/common/persistentState.ts

@@ -1,6 +1,6 @@
 import { ExtensionContext, Memento } from 'vscode';
-import { createDeferred, Deferred } from './utils/deferred';
 import { traceError } from './logging';
+import { createDeferred, Deferred } from './utils/deferred';
 
 export interface PersistentState {
     get<T>(key: string, defaultValue?: T): Promise<T | undefined>;
@@ -9,7 +9,6 @@ export interface PersistentState {
 }
 
 class PersistentStateImpl implements PersistentState {
-    private keys: string[] = [];
     private clearing: Deferred<void>;
     constructor(private readonly momento: Memento) {
         this.clearing = createDeferred<void>();
@@ -24,10 +23,6 @@ class PersistentStateImpl implements PersistentState {
     }
     async set<T>(key: string, value: T): Promise<void> {
         await this.clearing.promise;
-
-        if (!this.keys.includes(key)) {
-            this.keys.push(key);
-        }
         await this.momento.update(key, value);
 
         const before = JSON.stringify(value);
@@ -40,9 +35,8 @@ class PersistentStateImpl implements PersistentState {
     async clear(keys?: string[]): Promise<void> {
         if (this.clearing.completed) {
             this.clearing = createDeferred<void>();
-            const _keys = keys ?? this.keys;
+            const _keys = keys ?? this.momento.keys();
             await Promise.all(_keys.map((key) => this.momento.update(key, undefined)));
-            this.keys = this.keys.filter((k) => _keys.includes(k));
             this.clearing.resolve();
         }
         return this.clearing.promise;
@@ -64,3 +58,9 @@ export function getWorkspacePersistentState(): Promise<PersistentState> {
 export function getGlobalPersistentState(): Promise<PersistentState> {
     return _global.promise;
 }
+
+export async function clearPersistentState(): Promise<void> {
+    const [workspace, global] = await Promise.all([_workspace.promise, _global.promise]);
+    await Promise.all([workspace.clear(), global.clear()]);
+    return undefined;
+}

src/extension.ts

@@ -2,7 +2,7 @@ import { commands, ExtensionContext, LogOutputChannel, Terminal, Uri, window } f
 import { PythonEnvironment, PythonEnvironmentApi, PythonProjectCreator } from './api';
 import { ensureCorrectVersion } from './common/extVersion';
 import { registerLogger, traceError, traceInfo } from './common/logging';
-import { setPersistentState } from './common/persistentState';
+import { clearPersistentState, setPersistentState } from './common/persistentState';
 import { newProjectSelection } from './common/pickers/managers';
 import { StopWatch } from './common/stopWatch';
 import { EventNames } from './common/telemetry/constants';
@@ -15,6 +15,10 @@ import {
     onDidChangeActiveTerminal,
     onDidChangeTerminalShellIntegration,
 } from './common/window.apis';
+import { CreateQuickVirtualEnvironmentTool } from './features/chat/createQuickVenvTool';
+import { GetEnvironmentInfoTool } from './features/chat/getEnvInfoTool';
+import { GetExecutableTool } from './features/chat/getExecutableTool';
+import { InstallPackageTool } from './features/chat/installPackagesTool';
 import { createManagerReady } from './features/common/managerReady';
 import { AutoFindProjects } from './features/creators/autoFindProjects';
 import { ExistingProjects } from './features/creators/existingProjects';
@@ -211,6 +215,7 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
             await removePythonProject(item, projectManager);
         }),
         commands.registerCommand('python-envs.clearCache', async () => {
+            await clearPersistentState();
             await envManagers.clearCache(undefined);
             await clearShellProfileCache(shellStartupProviders);
         }),

src/features/terminal/shells/pwsh/pwshStartup.ts

@@ -8,7 +8,7 @@ import { ShellScriptEditState, ShellSetupState, ShellStartupScriptProvider } fro
 import { runCommand } from '../utils';
 
 import assert from 'assert';
-import { getGlobalPersistentState } from '../../../../common/persistentState';
+import { getWorkspacePersistentState } from '../../../../common/persistentState';
 import { ShellConstants } from '../../../common/shellConstants';
 import { hasStartupCode, insertStartupCode, removeStartupCode } from '../common/editUtils';
 import { extractProfilePath, PROFILE_TAG_END, PROFILE_TAG_START } from '../common/shellUtils';
@@ -18,25 +18,19 @@ const PWSH_PROFILE_PATH_CACHE_KEY = 'PWSH_PROFILE_PATH_CACHE';
 const PS5_PROFILE_PATH_CACHE_KEY = 'PS5_PROFILE_PATH_CACHE';
 let pwshProfilePath: string | undefined;
 let ps5ProfilePath: string | undefined;
-async function clearPwshCache(shell: 'powershell' | 'pwsh'): Promise<void> {
-    const global = await getGlobalPersistentState();
-    if (shell === 'powershell') {
-        ps5ProfilePath = undefined;
-        await global.clear([PS5_PROFILE_PATH_CACHE_KEY]);
-    } else {
-        pwshProfilePath = undefined;
-        await global.clear([PWSH_PROFILE_PATH_CACHE_KEY]);
-    }
+function clearPwshCache() {
+    ps5ProfilePath = undefined;
+    pwshProfilePath = undefined;
 }
 
 async function setProfilePathCache(shell: 'powershell' | 'pwsh', profilePath: string): Promise<void> {
-    const global = await getGlobalPersistentState();
+    const state = await getWorkspacePersistentState();
     if (shell === 'powershell') {
         ps5ProfilePath = profilePath;
-        await global.set(PS5_PROFILE_PATH_CACHE_KEY, profilePath);
+        await state.set(PS5_PROFILE_PATH_CACHE_KEY, profilePath);
     } else {
         pwshProfilePath = profilePath;
-        await global.set(PWSH_PROFILE_PATH_CACHE_KEY, profilePath);
+        await state.set(PWSH_PROFILE_PATH_CACHE_KEY, profilePath);
     }
 }
 
@@ -308,10 +302,7 @@ export class PwshStartupProvider implements ShellStartupScriptProvider {
     }
 
     async clearCache(): Promise<void> {
-        for (const shell of this._supportedShells) {
-            await clearPwshCache(shell);
-        }
-
+        clearPwshCache();
         // Reset installation check cache as well
         this._isPwshInstalled = undefined;
         this._isPs5Installed = undefined;

src/managers/conda/condaUtils.ts

@@ -29,7 +29,7 @@ import { ENVS_EXTENSION_ID, EXTENSION_ROOT_DIR } from '../../common/constants';
 import { showErrorMessageWithLogs } from '../../common/errors/utils';
 import { Common, CondaStrings, PackageManagement, Pickers } from '../../common/localize';
 import { traceInfo } from '../../common/logging';
-import { getGlobalPersistentState, getWorkspacePersistentState } from '../../common/persistentState';
+import { getWorkspacePersistentState } from '../../common/persistentState';
 import { pickProject } from '../../common/pickers/projects';
 import { createDeferred } from '../../common/utils/deferred';
 import { untildify } from '../../common/utils/pathUtils';
@@ -62,10 +62,6 @@ export const CONDA_GLOBAL_KEY = `${ENVS_EXTENSION_ID}:conda:GLOBAL_SELECTED`;
 
 let condaPath: string | undefined;
 export async function clearCondaCache(): Promise<void> {
-    const state = await getWorkspacePersistentState();
-    await state.clear([CONDA_PATH_KEY, CONDA_WORKSPACE_KEY, CONDA_GLOBAL_KEY]);
-    const global = await getGlobalPersistentState();
-    await global.clear([CONDA_PREFIXES_KEY]);
     condaPath = undefined;
 }
 
@@ -245,7 +241,7 @@ async function getPrefixes(): Promise<string[]> {
         return prefixes;
     }
 
-    const state = await getGlobalPersistentState();
+    const state = await getWorkspacePersistentState();
     prefixes = await state.get<string[]>(CONDA_PREFIXES_KEY);
     if (prefixes) {
         return prefixes;

src/managers/poetry/main.ts

@@ -36,17 +36,17 @@ export async function registerPoetryFeatures(
                     return;
                 }
             }
-        }
 
-        const envManager = new PoetryManager(nativeFinder, api);
-        const pkgManager = new PoetryPackageManager(api, outputChannel, envManager);
+            const envManager = new PoetryManager(nativeFinder, api);
+            const pkgManager = new PoetryPackageManager(api, outputChannel, envManager);
 
-        disposables.push(
-            envManager,
-            pkgManager,
-            api.registerEnvironmentManager(envManager),
-            api.registerPackageManager(pkgManager),
-        );
+            disposables.push(
+                envManager,
+                pkgManager,
+                api.registerEnvironmentManager(envManager),
+                api.registerPackageManager(pkgManager),
+            );
+        }
     } catch (ex) {
         traceInfo('Poetry not found, turning off poetry features.', ex);
     }

src/managers/poetry/poetryUtils.ts

@@ -11,7 +11,7 @@ import {
 } from '../../api';
 import { ENVS_EXTENSION_ID } from '../../common/constants';
 import { traceError, traceInfo } from '../../common/logging';
-import { getGlobalPersistentState, getWorkspacePersistentState } from '../../common/persistentState';
+import { getWorkspacePersistentState } from '../../common/persistentState';
 import { getUserHomeDir, untildify } from '../../common/utils/pathUtils';
 import { isWindows } from '../../common/utils/platformUtils';
 import { ShellConstants } from '../../features/common/shellConstants';
@@ -43,36 +43,28 @@ let poetryPath: string | undefined;
 let poetryVirtualenvsPath: string | undefined;
 
 export async function clearPoetryCache(): Promise<void> {
-    // Clear workspace-specific settings 
-    const state = await getWorkspacePersistentState();
-    await state.clear([POETRY_WORKSPACE_KEY]);
-    
-    // Clear global settings
-    const global = await getGlobalPersistentState();
-    await global.clear([POETRY_PATH_KEY, POETRY_GLOBAL_KEY, POETRY_VIRTUALENVS_PATH_KEY]);
-    
     // Reset in-memory cache
     poetryPath = undefined;
     poetryVirtualenvsPath = undefined;
 }
 
 async function setPoetry(poetry: string): Promise<void> {
     poetryPath = poetry;
-    const global = await getGlobalPersistentState();
-    await global.set(POETRY_PATH_KEY, poetry);
+    const state = await getWorkspacePersistentState();
+    await state.set(POETRY_PATH_KEY, poetry);
 
     // Also get and cache the virtualenvs path
     await getPoetryVirtualenvsPath(poetry);
 }
 
 export async function getPoetryForGlobal(): Promise<string | undefined> {
-    const global = await getGlobalPersistentState();
-    return await global.get(POETRY_GLOBAL_KEY);
+    const state = await getWorkspacePersistentState();
+    return await state.get(POETRY_GLOBAL_KEY);
 }
 
 export async function setPoetryForGlobal(poetryPath: string | undefined): Promise<void> {
-    const global = await getGlobalPersistentState();
-    await global.set(POETRY_GLOBAL_KEY, poetryPath);
+    const state = await getWorkspacePersistentState();
+    await state.set(POETRY_GLOBAL_KEY, poetryPath);
 }
 
 export async function getPoetryForWorkspace(fsPath: string): Promise<string | undefined> {
@@ -117,8 +109,8 @@ export async function getPoetry(native?: NativePythonFinder): Promise<string | u
         return poetryPath;
     }
 
-    const global = await getGlobalPersistentState();
-    poetryPath = await global.get<string>(POETRY_PATH_KEY);
+    const state = await getWorkspacePersistentState();
+    poetryPath = await state.get<string>(POETRY_PATH_KEY);
     if (poetryPath) {
         traceInfo(`Using poetry from persistent state: ${poetryPath}`);
         // Also retrieve the virtualenvs path if we haven't already
@@ -174,8 +166,8 @@ export async function getPoetryVirtualenvsPath(poetryExe?: string): Promise<stri
     }
 
     // Check if we have it in persistent state
-    const global = await getGlobalPersistentState();
-    poetryVirtualenvsPath = await global.get<string>(POETRY_VIRTUALENVS_PATH_KEY);
+    const state = await getWorkspacePersistentState();
+    poetryVirtualenvsPath = await state.get<string>(POETRY_VIRTUALENVS_PATH_KEY);
     if (poetryVirtualenvsPath) {
         return untildify(poetryVirtualenvsPath);
     }
@@ -199,7 +191,7 @@ export async function getPoetryVirtualenvsPath(poetryExe?: string): Promise<stri
                 }
 
                 if (poetryVirtualenvsPath) {
-                    await global.set(POETRY_VIRTUALENVS_PATH_KEY, poetryVirtualenvsPath);
+                    await state.set(POETRY_VIRTUALENVS_PATH_KEY, poetryVirtualenvsPath);
                     return poetryVirtualenvsPath;
                 }
             }
@@ -212,7 +204,7 @@ export async function getPoetryVirtualenvsPath(poetryExe?: string): Promise<stri
     const home = getUserHomeDir();
     if (home) {
         poetryVirtualenvsPath = path.join(home, '.cache', 'pypoetry', 'virtualenvs');
-        await global.set(POETRY_VIRTUALENVS_PATH_KEY, poetryVirtualenvsPath);
+        await state.set(POETRY_VIRTUALENVS_PATH_KEY, poetryVirtualenvsPath);
         return poetryVirtualenvsPath;
     }
 

src/managers/pyenv/pyenvUtils.ts

@@ -11,7 +11,7 @@ import {
 } from '../../api';
 import { ENVS_EXTENSION_ID } from '../../common/constants';
 import { traceError, traceInfo } from '../../common/logging';
-import { getGlobalPersistentState, getWorkspacePersistentState } from '../../common/persistentState';
+import { getWorkspacePersistentState } from '../../common/persistentState';
 import { getUserHomeDir, normalizePath, untildify } from '../../common/utils/pathUtils';
 import { isWindows } from '../../common/utils/platformUtils';
 import {
@@ -40,10 +40,7 @@ export const PYENV_GLOBAL_KEY = `${ENVS_EXTENSION_ID}:pyenv:GLOBAL_SELECTED`;
 
 let pyenvPath: string | undefined;
 export async function clearPyenvCache(): Promise<void> {
-    const state = await getWorkspacePersistentState();
-    await state.clear([PYENV_WORKSPACE_KEY, PYENV_GLOBAL_KEY]);
-    const global = await getGlobalPersistentState();
-    await global.clear([PYENV_PATH_KEY]);
+    pyenvPath = undefined;
 }
 
 async function setPyenv(pyenv: string): Promise<void> {
@@ -225,7 +222,7 @@ export async function refreshPyenv(
             .filter((e) => !isNativeEnvInfo(e))
             .map((e) => e as NativeEnvManagerInfo)
             .filter((e) => e.tool.toLowerCase() === 'pyenv');
-        
+
         if (managers.length > 0) {
             pyenv = managers[0].executable;
             await setPyenv(pyenv);