Detect if VSCode is launched from an activated environment and select it

src/client/common/persistentState.ts

@@ -6,7 +6,7 @@
 import { inject, injectable, named } from 'inversify';
 import { Memento } from 'vscode';
 import { IExtensionSingleActivationService } from '../activation/types';
-import { traceError } from '../logging';
+import { traceError, traceVerbose, traceWarn } from '../logging';
 import { ICommandManager } from './application/types';
 import { Commands } from './constants';
 import {
@@ -41,13 +41,24 @@ export class PersistentState<T> implements IPersistentState<T> {
         }
     }
 
-    public async updateValue(newValue: T): Promise<void> {
+    public async updateValue(newValue: T, retryOnce = true): Promise<void> {
         try {
             if (this.expiryDurationMs) {
                 await this.storage.update(this.key, { data: newValue, expiry: Date.now() + this.expiryDurationMs });
             } else {
                 await this.storage.update(this.key, newValue);
             }
+            if (retryOnce && JSON.stringify(this.value) != JSON.stringify(newValue)) {
+                // Due to a VSCode bug sometimes the changes are not reflected in the storage, atleast not immediately.
+                // It is noticed however that if we reset the storage first and then update it, it works.
+                // https://github.com/microsoft/vscode/issues/171827
+                traceVerbose('Storage update failed for key', this.key, ' retrying by resetting first');
+                await this.updateValue(undefined as any, false);
+                await this.updateValue(newValue, false);
+                if (JSON.stringify(this.value) != JSON.stringify(newValue)) {
+                    traceWarn('Retry failed, storage update failed for key', this.key);
+                }
+            }
         } catch (ex) {
             traceError('Error while updating storage for key:', this.key, ex);
         }

src/client/common/process/pythonExecutionFactory.ts

@@ -3,7 +3,7 @@
 import { inject, injectable } from 'inversify';
 
 import { IEnvironmentActivationService } from '../../interpreter/activation/types';
-import { IComponentAdapter } from '../../interpreter/contracts';
+import { IActivatedEnvironmentLaunch, IComponentAdapter } from '../../interpreter/contracts';
 import { IServiceContainer } from '../../ioc/types';
 import { sendTelemetryEvent } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
@@ -52,6 +52,10 @@ export class PythonExecutionFactory implements IPythonExecutionFactory {
     public async create(options: ExecutionFactoryCreationOptions): Promise<IPythonExecutionService> {
         let { pythonPath } = options;
         if (!pythonPath || pythonPath === 'python') {
+            const activatedEnvLaunch = this.serviceContainer.get<IActivatedEnvironmentLaunch>(
+                IActivatedEnvironmentLaunch,
+            );
+            await activatedEnvLaunch.selectIfLaunchedViaActivatedEnv();
             // If python path wasn't passed in, we need to auto select it and then read it
             // from the configuration.
             const interpreterPath = this.interpreterPathExpHelper.get(options.resource);

src/client/common/utils/localize.ts

@@ -191,6 +191,9 @@ export namespace Interpreters {
     export const condaInheritEnvMessage = l10n.t(
         'We noticed you\'re using a conda environment. If you are experiencing issues with this environment in the integrated terminal, we recommend that you let the Python extension change "terminal.integrated.inheritEnv" to false in your user settings.',
     );
+    export const activatedCondaEnvLaunch = l10n.t(
+        'We noticed VS Code was launched from an activated conda environment, would you like to select it?',
+    );
     export const environmentPromptMessage = l10n.t(
         'We noticed a new environment has been created. Do you want to select it for the workspace folder?',
     );

src/client/interpreter/contracts.ts

@@ -122,3 +122,8 @@ export type WorkspacePythonPath = {
     folderUri: Uri;
     configTarget: ConfigurationTarget.Workspace | ConfigurationTarget.WorkspaceFolder;
 };
+
+export const IActivatedEnvironmentLaunch = Symbol('IActivatedEnvironmentLaunch');
+export interface IActivatedEnvironmentLaunch {
+    selectIfLaunchedViaActivatedEnv(doNotBlockOnSelection?: boolean): Promise<string | undefined>;
+}

src/client/interpreter/interpreterService.ts

@@ -22,6 +22,7 @@ import {
 import { IServiceContainer } from '../ioc/types';
 import { PythonEnvironment } from '../pythonEnvironments/info';
 import {
+    IActivatedEnvironmentLaunch,
     IComponentAdapter,
     IInterpreterDisplay,
     IInterpreterService,
@@ -179,7 +180,13 @@ export class InterpreterService implements Disposable, IInterpreterService {
     }
 
     public async getActiveInterpreter(resource?: Uri): Promise<PythonEnvironment | undefined> {
-        let path = this.configService.getSettings(resource).pythonPath;
+        const activatedEnvLaunch = this.serviceContainer.get<IActivatedEnvironmentLaunch>(IActivatedEnvironmentLaunch);
+        let path = await activatedEnvLaunch.selectIfLaunchedViaActivatedEnv(true);
+        // This is being set as interpreter in background, after which it'll show up in `.pythonPath` config.
+        // However we need not wait on the update to take place, as we can use the value directly.
+        if (!path) {
+            path = this.configService.getSettings(resource).pythonPath;
+        }
         if (pathUtils.basename(path) === path) {
             // Value can be `python`, `python3`, `python3.9` etc.
             // Note the following triggers autoselection if no interpreter is explictly

src/client/interpreter/serviceRegistry.ts

@@ -25,11 +25,12 @@ import {
     IPythonPathUpdaterServiceFactory,
     IPythonPathUpdaterServiceManager,
 } from './configuration/types';
-import { IInterpreterDisplay, IInterpreterHelper, IInterpreterService } from './contracts';
+import { IActivatedEnvironmentLaunch, IInterpreterDisplay, IInterpreterHelper, IInterpreterService } from './contracts';
 import { InterpreterDisplay } from './display';
 import { InterpreterLocatorProgressStatubarHandler } from './display/progressDisplay';
 import { InterpreterHelper } from './helpers';
 import { InterpreterService } from './interpreterService';
+import { ActivatedEnvironmentLaunch } from './virtualEnvs/activatedEnvLaunch';
 import { CondaInheritEnvPrompt } from './virtualEnvs/condaInheritEnvPrompt';
 import { VirtualEnvironmentPrompt } from './virtualEnvs/virtualEnvPrompt';
 
@@ -90,6 +91,7 @@ export function registerInterpreterTypes(serviceManager: IServiceManager): void
     );
 
     serviceManager.addSingleton<IExtensionActivationService>(IExtensionActivationService, CondaInheritEnvPrompt);
+    serviceManager.addSingleton<IActivatedEnvironmentLaunch>(IActivatedEnvironmentLaunch, ActivatedEnvironmentLaunch);
 }
 
 export function registerTypes(serviceManager: IServiceManager): void {

src/client/interpreter/virtualEnvs/activatedEnvLaunch.ts

@@ -0,0 +1,160 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { inject, injectable, optional } from 'inversify';
+import { ConfigurationTarget } from 'vscode';
+import * as path from 'path';
+import { IApplicationShell, IWorkspaceService } from '../../common/application/types';
+import { IProcessServiceFactory } from '../../common/process/types';
+import { sleep } from '../../common/utils/async';
+import { cache } from '../../common/utils/decorators';
+import { Common, Interpreters } from '../../common/utils/localize';
+import { traceError, traceLog, traceWarn } from '../../logging';
+import { Conda } from '../../pythonEnvironments/common/environmentManagers/conda';
+import { sendTelemetryEvent } from '../../telemetry';
+import { EventName } from '../../telemetry/constants';
+import { IPythonPathUpdaterServiceManager } from '../configuration/types';
+import { IActivatedEnvironmentLaunch, IInterpreterService } from '../contracts';
+
+@injectable()
+export class ActivatedEnvironmentLaunch implements IActivatedEnvironmentLaunch {
+    public readonly supportedWorkspaceTypes = { untrustedWorkspace: false, virtualWorkspace: true };
+
+    private inMemorySelection: string | undefined;
+
+    constructor(
+        @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService,
+        @inject(IApplicationShell) private readonly appShell: IApplicationShell,
+        @inject(IPythonPathUpdaterServiceManager)
+        private readonly pythonPathUpdaterService: IPythonPathUpdaterServiceManager,
+        @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
+        @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory,
+        @optional() public wasSelected: boolean = false,
+    ) {}
+
+    @cache(-1, true)
+    public async _promptIfApplicable(): Promise<void> {
+        const baseCondaPrefix = getPrefixOfActivatedCondaEnv();
+        if (!baseCondaPrefix) {
+            return;
+        }
+        const info = await this.interpreterService.getInterpreterDetails(baseCondaPrefix);
+        if (info?.envName !== 'base') {
+            // Only show prompt for base conda environments, as we need to check config for such envs which can be slow.
+            return;
+        }
+        const conda = await Conda.getConda();
+        if (!conda) {
+            traceWarn('Conda not found even though activated environment vars are set');
+            return;
+        }
+        const service = await this.processServiceFactory.create();
+        const autoActivateBaseConfig = await service
+            .shellExec(`${conda.shellCommand} config --get auto_activate_base`)
+            .catch((ex) => {
+                traceError(ex);
+                return { stdout: '' };
+            });
+        if (autoActivateBaseConfig.stdout.trim().toLowerCase().endsWith('false')) {
+            await this.promptAndUpdate(baseCondaPrefix);
+        }
+    }
+
+    private async promptAndUpdate(prefix: string) {
+        this.wasSelected = true;
+        const prompts = [Common.bannerLabelYes, Common.bannerLabelNo];
+        const telemetrySelections: ['Yes', 'No'] = ['Yes', 'No'];
+        const selection = await this.appShell.showInformationMessage(Interpreters.activatedCondaEnvLaunch, ...prompts);
+        sendTelemetryEvent(EventName.ACTIVATED_CONDA_ENV_LAUNCH, undefined, {
+            selection: selection ? telemetrySelections[prompts.indexOf(selection)] : undefined,
+        });
+        if (!selection) {
+            return;
+        }
+        if (selection === prompts[0]) {
+            await this.setInterpeterInStorage(prefix);
+        }
+    }
+
+    public async selectIfLaunchedViaActivatedEnv(doNotBlockOnSelection = false): Promise<string | undefined> {
+        if (this.wasSelected) {
+            return this.inMemorySelection;
+        }
+        return this._selectIfLaunchedViaActivatedEnv(doNotBlockOnSelection);
+    }
+
+    @cache(-1, true)
+    private async _selectIfLaunchedViaActivatedEnv(doNotBlockOnSelection = false): Promise<string | undefined> {
+        if (this.workspaceService.workspaceFile) {
+            // Assuming multiroot workspaces cannot be directly launched via `code .` command.
+            return undefined;
+        }
+        const prefix = await this.getPrefixOfSelectedActivatedEnv();
+        if (!prefix) {
+            this._promptIfApplicable().ignoreErrors();
+            return undefined;
+        }
+        this.wasSelected = true;
+        this.inMemorySelection = prefix;
+        traceLog(
+            `VS Code was launched from an activated environment: '${path.basename(
+                prefix,
+            )}', selecting it as the interpreter for workspace.`,
+        );
+        if (doNotBlockOnSelection) {
+            this.setInterpeterInStorage(prefix).ignoreErrors();
+        } else {
+            await this.setInterpeterInStorage(prefix);
+            await sleep(1); // Yield control so config service can update itself.
+        }
+        this.inMemorySelection = undefined; // Once we have set the prefix in storage, clear the in memory selection.
+        return prefix;
+    }
+
+    private async setInterpeterInStorage(prefix: string) {
+        const { workspaceFolders } = this.workspaceService;
+        if (!workspaceFolders || workspaceFolders.length === 0) {
+            await this.pythonPathUpdaterService.updatePythonPath(prefix, ConfigurationTarget.Global, 'load');
+        } else {
+            await this.pythonPathUpdaterService.updatePythonPath(
+                prefix,
+                ConfigurationTarget.WorkspaceFolder,
+                'load',
+                workspaceFolders[0].uri,
+            );
+        }
+    }
+
+    private async getPrefixOfSelectedActivatedEnv(): Promise<string | undefined> {
+        const virtualEnvVar = process.env.VIRTUAL_ENV;
+        if (virtualEnvVar !== undefined && virtualEnvVar.length > 0) {
+            return virtualEnvVar;
+        }
+        const condaPrefixVar = getPrefixOfActivatedCondaEnv();
+        if (!condaPrefixVar) {
+            return undefined;
+        }
+        const info = await this.interpreterService.getInterpreterDetails(condaPrefixVar);
+        if (info?.envName !== 'base') {
+            return condaPrefixVar;
+        }
+        // Ignoring base conda environments, as they could be automatically set by conda.
+        if (process.env.CONDA_AUTO_ACTIVATE_BASE !== undefined) {
+            if (process.env.CONDA_AUTO_ACTIVATE_BASE.toLowerCase() === 'false') {
+                return condaPrefixVar;
+            }
+        }
+        return undefined;
+    }
+}
+
+function getPrefixOfActivatedCondaEnv() {
+    const condaPrefixVar = process.env.CONDA_PREFIX;
+    if (condaPrefixVar && condaPrefixVar.length > 0) {
+        const condaShlvl = process.env.CONDA_SHLVL;
+        if (condaShlvl !== undefined && condaShlvl.length > 0 && condaShlvl > '0') {
+            return condaPrefixVar;
+        }
+    }
+    return undefined;
+}

src/client/telemetry/constants.ts

@@ -30,6 +30,7 @@ export enum EventName {
     PYTHON_INTERPRETER_ACTIVATE_ENVIRONMENT_PROMPT = 'PYTHON_INTERPRETER_ACTIVATE_ENVIRONMENT_PROMPT',
     PYTHON_NOT_INSTALLED_PROMPT = 'PYTHON_NOT_INSTALLED_PROMPT',
     CONDA_INHERIT_ENV_PROMPT = 'CONDA_INHERIT_ENV_PROMPT',
+    ACTIVATED_CONDA_ENV_LAUNCH = 'ACTIVATED_CONDA_ENV_LAUNCH',
     ENVFILE_VARIABLE_SUBSTITUTION = 'ENVFILE_VARIABLE_SUBSTITUTION',
     ENVFILE_WORKSPACE = 'ENVFILE_WORKSPACE',
     EXECUTION_CODE = 'EXECUTION_CODE',

src/client/telemetry/index.ts

@@ -1299,8 +1299,9 @@ export interface IEventNamePropertyMapping {
         environmentsWithoutPython?: number;
     };
     /**
-     * Telemetry event sent with details when user clicks the prompt with the following message
-     * `Prompt message` :- 'We noticed you're using a conda environment. If you are experiencing issues with this environment in the integrated terminal, we suggest the "terminal.integrated.inheritEnv" setting to be changed to false. Would you like to update this setting?'
+     * Telemetry event sent with details when user clicks the prompt with the following message:
+     *
+     * 'We noticed you're using a conda environment. If you are experiencing issues with this environment in the integrated terminal, we suggest the "terminal.integrated.inheritEnv" setting to be changed to false. Would you like to update this setting?'
      */
     /* __GDPR__
        "conda_inherit_env_prompt" : {
@@ -1315,6 +1316,23 @@ export interface IEventNamePropertyMapping {
          */
         selection: 'Yes' | 'No' | 'More Info' | undefined;
     };
+    /**
+     * Telemetry event sent with details when user clicks the prompt with the following message:
+     *
+     * 'We noticed VS Code was launched from an activated conda environment, would you like to select it?'
+     */
+    /* __GDPR__
+       "activated_conda_env_launch" : {
+          "selection" : { "classification": "SystemMetaData", "purpose": "FeatureInsight", "owner": "karrtikr" }
+       }
+     */
+    [EventName.ACTIVATED_CONDA_ENV_LAUNCH]: {
+        /**
+         * `Yes` When 'Yes' option is selected
+         * `No` When 'No' option is selected
+         */
+        selection: 'Yes' | 'No' | undefined;
+    };
     /**
      * Telemetry event sent with details when user clicks a button in the virtual environment prompt.
      * `Prompt message` :- 'We noticed a new virtual environment has been created. Do you want to select it for the workspace folder?'

src/test/common/installer.test.ts

@@ -100,6 +100,14 @@ import { MockModuleInstaller } from '../mocks/moduleInstaller';
 import { MockProcessService } from '../mocks/proc';
 import { UnitTestIocContainer } from '../testing/serviceRegistry';
 import { closeActiveWindows, initializeTest, IS_MULTI_ROOT_TEST, TEST_TIMEOUT } from '../initialize';
+import { IActivatedEnvironmentLaunch } from '../../client/interpreter/contracts';
+import { ActivatedEnvironmentLaunch } from '../../client/interpreter/virtualEnvs/activatedEnvLaunch';
+import {
+    IPythonPathUpdaterServiceFactory,
+    IPythonPathUpdaterServiceManager,
+} from '../../client/interpreter/configuration/types';
+import { PythonPathUpdaterService } from '../../client/interpreter/configuration/pythonPathUpdaterService';
+import { PythonPathUpdaterServiceFactory } from '../../client/interpreter/configuration/pythonPathUpdaterServiceFactory';
 
 suite('Installer', () => {
     let ioc: UnitTestIocContainer;
@@ -169,6 +177,18 @@ suite('Installer', () => {
             TestFrameworkProductPathService,
             ProductType.TestFramework,
         );
+        ioc.serviceManager.addSingleton<IActivatedEnvironmentLaunch>(
+            IActivatedEnvironmentLaunch,
+            ActivatedEnvironmentLaunch,
+        );
+        ioc.serviceManager.addSingleton<IPythonPathUpdaterServiceManager>(
+            IPythonPathUpdaterServiceManager,
+            PythonPathUpdaterService,
+        );
+        ioc.serviceManager.addSingleton<IPythonPathUpdaterServiceFactory>(
+            IPythonPathUpdaterServiceFactory,
+            PythonPathUpdaterServiceFactory,
+        );
         ioc.serviceManager.addSingleton<IActiveResourceService>(IActiveResourceService, ActiveResourceService);
         ioc.serviceManager.addSingleton<IInterpreterPathService>(IInterpreterPathService, InterpreterPathService);
         ioc.serviceManager.addSingleton<IExtensions>(IExtensions, Extensions);

src/test/common/moduleInstaller.test.ts

@@ -1,7 +1,7 @@
 import { expect, should as chaiShould, use as chaiUse } from 'chai';
 import * as chaiAsPromised from 'chai-as-promised';
 import { SemVer } from 'semver';
-import { instance, mock } from 'ts-mockito';
+import { instance, mock, when } from 'ts-mockito';
 import * as TypeMoq from 'typemoq';
 import { ConfigurationTarget, Uri } from 'vscode';
 import { IExtensionSingleActivationService } from '../../client/activation/types';
@@ -90,7 +90,12 @@ import {
 import { IMultiStepInputFactory, MultiStepInputFactory } from '../../client/common/utils/multiStepInput';
 import { Architecture } from '../../client/common/utils/platform';
 import { Random } from '../../client/common/utils/random';
-import { ICondaService, IInterpreterService, IComponentAdapter } from '../../client/interpreter/contracts';
+import {
+    ICondaService,
+    IInterpreterService,
+    IComponentAdapter,
+    IActivatedEnvironmentLaunch,
+} from '../../client/interpreter/contracts';
 import { IServiceContainer } from '../../client/ioc/types';
 import { JupyterExtensionDependencyManager } from '../../client/jupyter/jupyterExtensionDependencyManager';
 import { EnvironmentType, PythonEnvironment } from '../../client/pythonEnvironments/info';
@@ -163,7 +168,12 @@ suite('Module Installer', () => {
                 ITerminalServiceFactory,
                 mockTerminalFactory.object,
             );
-
+            const activatedEnvironmentLaunch = mock<IActivatedEnvironmentLaunch>();
+            when(activatedEnvironmentLaunch.selectIfLaunchedViaActivatedEnv()).thenResolve(undefined);
+            ioc.serviceManager.addSingletonInstance<IActivatedEnvironmentLaunch>(
+                IActivatedEnvironmentLaunch,
+                instance(activatedEnvironmentLaunch),
+            );
             ioc.serviceManager.addSingleton<IModuleInstaller>(IModuleInstaller, PipInstaller);
             ioc.serviceManager.addSingleton<IModuleInstaller>(IModuleInstaller, CondaInstaller);
             ioc.serviceManager.addSingleton<IModuleInstaller>(IModuleInstaller, PipEnvInstaller);