Ability to track the user selected Environment (#25090)

DonJayamanne · web-flow · commit d7a5ab7848c7 · 2025-05-20T03:10:53.000+02:00
@karthiknadig /cc
diff --git a/src/client/chat/installPackagesTool.ts b/src/client/chat/installPackagesTool.ts
@@ -19,7 +19,6 @@ import { resolveFilePath } from './utils';
 import { IModuleInstaller } from '../common/installer/types';
 import { ModuleInstallerType } from '../pythonEnvironments/info';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
-import { trackEnvUsedByTool } from './lastUsedEnvs';
 
 export interface IInstallPackageArgs {
     resourcePath?: string;
@@ -67,7 +66,6 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             for (const packageName of options.input.packageList) {
                 await installer.installModule(packageName, resourcePath, token, undefined, { installAsProcess: true });
             }
-            trackEnvUsedByTool(resourcePath, environment);
             // format and return
             const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
             return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);
diff --git a/src/client/chat/lastUsedEnvs.ts b/src/client/chat/lastUsedEnvs.ts
diff --git a/src/client/chat/listPackagesTool.ts b/src/client/chat/listPackagesTool.ts
@@ -22,7 +22,6 @@ import { parsePipList } from './pipListUtils';
 import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
 import { traceError } from '../logging';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
-import { trackEnvUsedByTool } from './lastUsedEnvs';
 
 export interface IResourceReference {
     resourcePath?: string;
@@ -109,7 +108,6 @@ export async function getPythonPackagesResponse(
     if (!packages.length) {
         return 'No packages found';
     }
-    trackEnvUsedByTool(resourcePath, environment);
     // Installed Python packages, each in the format <name> or <name> (<version>). The version may be omitted if unknown. Returns an empty array if no packages are installed.
     const response = [
         'Below is a list of the Python packages, each in the format <name> or <name> (<version>). The version may be omitted if unknown: ',
diff --git a/src/client/chat/utils.ts b/src/client/chat/utils.ts
@@ -7,7 +7,6 @@ import { PythonExtension, ResolvedEnvironment } from '../api/types';
 import { ITerminalHelper, TerminalShellType } from '../common/terminal/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { Conda } from '../pythonEnvironments/common/environmentManagers/conda';
-import { trackEnvUsedByTool } from './lastUsedEnvs';
 
 export function resolveFilePath(filepath?: string): Uri | undefined {
     if (!filepath) {
@@ -71,7 +70,6 @@ export async function getEnvironmentDetails(
         getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper),
         token,
     );
-    trackEnvUsedByTool(resourcePath, environment);
     const message = [
         `Following is the information about the Python environment:`,
         `1. Environment Type: ${environment.environment?.type || 'unknown'}`,
diff --git a/src/client/common/utils/async.ts b/src/client/common/utils/async.ts
@@ -268,3 +268,26 @@ export async function waitForCondition(
         }, 10);
     });
 }
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function isPromiseLike<T>(v: any): v is PromiseLike<T> {
+    return typeof v?.then === 'function';
+}
+
+export function raceTimeout<T>(timeout: number, ...promises: Promise<T>[]): Promise<T | undefined>;
+export function raceTimeout<T>(timeout: number, defaultValue: T, ...promises: Promise<T>[]): Promise<T>;
+export function raceTimeout<T>(timeout: number, defaultValue: T, ...promises: Promise<T>[]): Promise<T> {
+    const resolveValue = isPromiseLike(defaultValue) ? undefined : defaultValue;
+    if (isPromiseLike(defaultValue)) {
+        promises.push((defaultValue as unknown) as Promise<T>);
+    }
+
+    let promiseResolve: ((value: T) => void) | undefined = undefined;
+
+    const timer = setTimeout(() => promiseResolve?.((resolveValue as unknown) as T), timeout);
+
+    return Promise.race([
+        Promise.race(promises).finally(() => clearTimeout(timer)),
+        new Promise<T>((resolve) => (promiseResolve = resolve)),
+    ]);
+}
diff --git a/src/client/extension.ts b/src/client/extension.ts
@@ -44,6 +44,7 @@ import { ProposedExtensionAPI } from './proposedApiTypes';
 import { buildProposedApi } from './proposedApi';
 import { GLOBAL_PERSISTENT_KEYS } from './common/persistentState';
 import { registerTools } from './chat';
+import { IRecommendedEnvironmentService } from './interpreter/configuration/types';
 
 durations.codeLoadingTime = stopWatch.elapsedTime;
 
@@ -164,6 +165,9 @@ async function activateUnsafe(
     );
     const proposedApi = buildProposedApi(components.pythonEnvs, ext.legacyIOC.serviceContainer);
     registerTools(context, components.pythonEnvs, api.environments, ext.legacyIOC.serviceContainer);
+    ext.legacyIOC.serviceContainer
+        .get<IRecommendedEnvironmentService>(IRecommendedEnvironmentService)
+        .registerEnvApi(api.environments);
     return [{ ...api, ...proposedApi }, activationPromise, ext.legacyIOC.serviceContainer];
 }
 
diff --git a/src/client/interpreter/configuration/interpreterSelector/commands/setInterpreter.ts b/src/client/interpreter/configuration/interpreterSelector/commands/setInterpreter.ts
@@ -565,8 +565,11 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem
         return Promise.resolve();
     }
 
+    /**
+     * @returns true when an interpreter was set, undefined if the user cancelled the quickpick.
+     */
     @captureTelemetry(EventName.SELECT_INTERPRETER)
-    public async setInterpreter(): Promise<void> {
+    public async setInterpreter(): Promise<true | undefined> {
         const targetConfig = await this.getConfigTargets();
         if (!targetConfig) {
             return;
@@ -588,6 +591,7 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem
             if (useEnvExtension()) {
                 await setInterpreterLegacy(interpreterState.path, wkspace);
             }
+            return true;
         }
     }
 
diff --git a/src/client/interpreter/configuration/pythonPathUpdaterService.ts b/src/client/interpreter/configuration/pythonPathUpdaterService.ts
@@ -7,14 +7,19 @@ import { sendTelemetryEvent } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
 import { PythonInterpreterTelemetry } from '../../telemetry/types';
 import { IComponentAdapter } from '../contracts';
-import { IPythonPathUpdaterServiceFactory, IPythonPathUpdaterServiceManager } from './types';
+import {
+    IRecommendedEnvironmentService,
+    IPythonPathUpdaterServiceFactory,
+    IPythonPathUpdaterServiceManager,
+} from './types';
 
 @injectable()
 export class PythonPathUpdaterService implements IPythonPathUpdaterServiceManager {
     constructor(
         @inject(IPythonPathUpdaterServiceFactory)
         private readonly pythonPathSettingsUpdaterFactory: IPythonPathUpdaterServiceFactory,
         @inject(IComponentAdapter) private readonly pyenvs: IComponentAdapter,
+        @inject(IRecommendedEnvironmentService) private readonly preferredEnvService: IRecommendedEnvironmentService,
     ) {}
 
     public async updatePythonPath(
@@ -28,6 +33,9 @@ export class PythonPathUpdaterService implements IPythonPathUpdaterServiceManage
         let failed = false;
         try {
             await pythonPathUpdater.updatePythonPath(pythonPath);
+            if (trigger === 'ui') {
+                this.preferredEnvService.trackUserSelectedEnvironment(pythonPath, wkspace);
+            }
         } catch (err) {
             failed = true;
             const reason = err as Error;
diff --git a/src/client/interpreter/configuration/recommededEnvironmentService.ts b/src/client/interpreter/configuration/recommededEnvironmentService.ts
@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { inject, injectable } from 'inversify';
+import { IRecommendedEnvironmentService } from './types';
+import { PythonExtension } from '../../api/types';
+import { IExtensionContext, Resource } from '../../common/types';
+import { Uri, workspace } from 'vscode';
+import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../../common/persistentState';
+import { traceError } from '../../logging';
+
+const MEMENTO_KEY = 'userSelectedEnvPath';
+
+@injectable()
+export class RecommendedEnvironmentService implements IRecommendedEnvironmentService {
+    private api?: PythonExtension['environments'];
+    constructor(@inject(IExtensionContext) private readonly extensionContext: IExtensionContext) {}
+
+    registerEnvApi(api: PythonExtension['environments']) {
+        this.api = api;
+    }
+
+    trackUserSelectedEnvironment(environmentPath: string | undefined, uri: Uri | undefined) {
+        if (workspace.workspaceFolders?.length) {
+            try {
+                void updateWorkspaceStateValue(MEMENTO_KEY, getDataToStore(environmentPath, uri));
+            } catch (ex) {
+                traceError('Failed to update workspace state for preferred environment', ex);
+            }
+        } else {
+            void this.extensionContext.globalState.update(MEMENTO_KEY, environmentPath);
+        }
+    }
+
+    getRecommededEnvironment(
+        resource: Resource,
+    ):
+        | { environmentPath: string; reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended' }
+        | undefined {
+        let workspaceState: string | undefined = undefined;
+        try {
+            workspaceState = getWorkspaceStateValue<string>(MEMENTO_KEY);
+        } catch (ex) {
+            traceError('Failed to get workspace state for preferred environment', ex);
+        }
+
+        if (workspace.workspaceFolders?.length && workspaceState) {
+            const workspaceUri = (
+                (resource ? workspace.getWorkspaceFolder(resource)?.uri : undefined) ||
+                workspace.workspaceFolders[0].uri
+            ).toString();
+
+            try {
+                const existingJson: Record<string, string> = JSON.parse(workspaceState);
+                const selectedEnvPath = existingJson[workspaceUri];
+                if (selectedEnvPath) {
+                    return { environmentPath: selectedEnvPath, reason: 'workspaceUserSelected' };
+                }
+            } catch (ex) {
+                traceError('Failed to parse existing workspace state value for preferred environment', ex);
+            }
+        }
+
+        const globalSelectedEnvPath = this.extensionContext.globalState.get<string | undefined>(MEMENTO_KEY);
+        if (globalSelectedEnvPath) {
+            return { environmentPath: globalSelectedEnvPath, reason: 'globalUserSelected' };
+        }
+        return this.api && workspace.isTrusted
+            ? {
+                  environmentPath: this.api.getActiveEnvironmentPath(resource).path,
+                  reason: 'defaultRecommended',
+              }
+            : undefined;
+    }
+}
+
+function getDataToStore(environmentPath: string | undefined, uri: Uri | undefined): string | undefined {
+    if (!workspace.workspaceFolders?.length) {
+        return environmentPath;
+    }
+    const workspaceUri = (
+        (uri ? workspace.getWorkspaceFolder(uri)?.uri : undefined) || workspace.workspaceFolders[0].uri
+    ).toString();
+    const existingData = getWorkspaceStateValue<string>(MEMENTO_KEY);
+    if (!existingData) {
+        return JSON.stringify(environmentPath ? { [workspaceUri]: environmentPath } : {});
+    }
+    try {
+        const existingJson: Record<string, string> = JSON.parse(existingData);
+        if (environmentPath) {
+            existingJson[workspaceUri] = environmentPath;
+        } else {
+            delete existingJson[workspaceUri];
+        }
+        return JSON.stringify(existingJson);
+    } catch (ex) {
+        traceError('Failed to parse existing workspace state value for preferred environment', ex);
+        return JSON.stringify({
+            [workspaceUri]: environmentPath,
+        });
+    }
+}
diff --git a/src/client/interpreter/configuration/types.ts b/src/client/interpreter/configuration/types.ts
@@ -1,6 +1,7 @@
 import { ConfigurationTarget, Disposable, QuickPickItem, Uri } from 'vscode';
 import { Resource } from '../../common/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
+import { PythonExtension } from '../../api/types';
 
 export interface IPythonPathUpdaterService {
     updatePythonPath(pythonPath: string | undefined): Promise<void>;
@@ -96,3 +97,14 @@ export interface IInterpreterQuickPick {
         params?: InterpreterQuickPickParams,
     ): Promise<string | undefined>;
 }
+
+export const IRecommendedEnvironmentService = Symbol('IRecommendedEnvironmentService');
+export interface IRecommendedEnvironmentService {
+    registerEnvApi(api: PythonExtension['environments']): void;
+    trackUserSelectedEnvironment(environmentPath: string | undefined, uri: Uri | undefined): void;
+    getRecommededEnvironment(
+        resource: Resource,
+    ):
+        | { environmentPath: string; reason: 'globalUserSelected' | 'workspaceUserSelected' | 'defaultRecommended' }
+        | undefined;
+}
diff --git a/src/client/interpreter/serviceRegistry.ts b/src/client/interpreter/serviceRegistry.ts
@@ -16,12 +16,14 @@ import { InstallPythonViaTerminal } from './configuration/interpreterSelector/co
 import { ResetInterpreterCommand } from './configuration/interpreterSelector/commands/resetInterpreter';
 import { SetInterpreterCommand } from './configuration/interpreterSelector/commands/setInterpreter';
 import { InterpreterSelector } from './configuration/interpreterSelector/interpreterSelector';
+import { RecommendedEnvironmentService } from './configuration/recommededEnvironmentService';
 import { PythonPathUpdaterService } from './configuration/pythonPathUpdaterService';
 import { PythonPathUpdaterServiceFactory } from './configuration/pythonPathUpdaterServiceFactory';
 import {
     IInterpreterComparer,
     IInterpreterQuickPick,
     IInterpreterSelector,
+    IRecommendedEnvironmentService,
     IPythonPathUpdaterServiceFactory,
     IPythonPathUpdaterServiceManager,
 } from './configuration/types';
@@ -59,6 +61,10 @@ export function registerInterpreterTypes(serviceManager: IServiceManager): void
         IExtensionSingleActivationService,
         ResetInterpreterCommand,
     );
+    serviceManager.addSingleton<IRecommendedEnvironmentService>(
+        IRecommendedEnvironmentService,
+        RecommendedEnvironmentService,
+    );
     serviceManager.addSingleton(IInterpreterQuickPick, SetInterpreterCommand);
 
     serviceManager.addSingleton<IExtensionActivationService>(IExtensionActivationService, VirtualEnvironmentPrompt);
diff --git a/src/client/jupyter/jupyterIntegration.ts b/src/client/jupyter/jupyterIntegration.ts
diff --git a/src/test/debugger/envVars.test.ts b/src/test/debugger/envVars.test.ts
diff --git a/src/test/interpreters/serviceRegistry.unit.test.ts b/src/test/interpreters/serviceRegistry.unit.test.ts

Original file line number	Diff line number	Diff line change
`@@ -565,8 +565,11 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem`
`565`	`565`	`return Promise.resolve();`
`566`	`566`	`}`
`567`	`567`
	`568`	`+ /**`
	`569`	`+ * @returns true when an interpreter was set, undefined if the user cancelled the quickpick.`
	`570`	`+ */`
`568`	`571`	`@captureTelemetry(EventName.SELECT_INTERPRETER)`
`569`		`- public async setInterpreter(): Promise<void> {`
	`572`	`+ public async setInterpreter(): Promise<true \| undefined> {`
`570`	`573`	`const targetConfig = await this.getConfigTargets();`
`571`	`574`	`if (!targetConfig) {`
`572`	`575`	`return;`
`@@ -588,6 +591,7 @@ export class SetInterpreterCommand extends BaseInterpreterSelectorCommand implem`
`588`	`591`	`if (useEnvExtension()) {`
`589`	`592`	`await setInterpreterLegacy(interpreterState.path, wkspace);`
`590`	`593`	`}`
	`594`	`+ return true;`
`591`	`595`	`}`
`592`	`596`	`}`
`593`	`597`