feat: add activation using shell startup

Author: karthiknadig

package.json

@@ -90,14 +90,14 @@
                 },
                 "python-envs.terminal.autoActivationType": {
                     "type": "string",
-                    "description": "%python-envs.terminal.autoActivationType.description%",
+                    "markdownDescription": "%python-envs.terminal.autoActivationType.description%",
                     "default": "command",
                     "enum": [
                         "command",
                         "shellStartup",
                         "off"
                     ],
-                    "enumDescriptions": [
+                    "markdownEnumDescriptions": [
                         "%python-envs.terminal.autoActivationType.command%",
                         "%python-envs.terminal.autoActivationType.shellStartup%",
                         "%python-envs.terminal.autoActivationType.off%"
@@ -247,6 +247,12 @@
                 "title": "%python-envs.copyProjectPath.title%",
                 "category": "Python Envs",
                 "icon": "$(copy)"
+            },
+            {
+                "command": "python-envs.terminal.revertStartupScriptChanges",
+                "title": "%python-envs.terminal.revertStartupScriptChanges.title%",
+                "category": "Python Envs",
+                "icon": "$(discard)"
             }
         ],
         "menus": {

package.nls.json

@@ -8,8 +8,9 @@
     "python-envs.terminal.showActivateButton.description": "Whether to show the 'Activate' button in the terminal menu",
     "python-envs.terminal.autoActivationType.description": "The type of activation to use when activating an environment in the terminal",
     "python-envs.terminal.autoActivationType.command": "Activation by executing a command in the terminal.",
-    "python-envs.terminal.autoActivationType.shellStartup": "Activation by modifying the terminal shell startup script.",
+    "python-envs.terminal.autoActivationType.shellStartup": "Activation by modifying the terminal shell startup script. To use this feature we will need to modify your shell startup scripts.",
     "python-envs.terminal.autoActivationType.off": "No automatic activation of environments.",
+    "python-envs.terminal.revertStartupScriptChanges.title": "Revert Shell Startup Script Changes",
     "python-envs.setEnvManager.title": "Set Environment Manager",
     "python-envs.setPkgManager.title": "Set Package Manager",
     "python-envs.addPythonProject.title": "Add Python Project",

src/common/localize.ts

@@ -144,3 +144,13 @@ export namespace EnvViewStrings {
     export const selectedGlobalTooltip = l10n.t('This environment is selected for non-workspace files');
     export const selectedWorkspaceTooltip = l10n.t('This environment is selected for workspace files');
 }
+
+export namespace ShellStartupActivationStrings {
+    export const shellStartupScriptEditPrompt = l10n.t(
+        'To support `shellStartup` we need to modify your shell profile. Do you want to proceed?',
+    );
+    export const updateScript = l10n.t('Update Shell Profile');
+    export const revertToCommandActivation = l10n.t(
+        'Auto Shell Activation type set to "command", due to removing shell startup from profile.',
+    );
+}

src/common/window.apis.ts

@@ -7,6 +7,8 @@ import {
     InputBox,
     InputBoxOptions,
     LogOutputChannel,
+    MessageItem,
+    MessageOptions,
     OpenDialogOptions,
     OutputChannel,
     Progress,
@@ -284,6 +286,21 @@ export async function showInputBoxWithButtons(
     }
 }
 
+export function showInformationMessage<T extends string>(message: string, ...items: T[]): Thenable<T | undefined>;
+export function showInformationMessage<T extends MessageItem>(message: string, ...items: T[]): Thenable<T | undefined>;
+export function showInformationMessage<T extends string>(
+    message: string,
+    options: MessageOptions,
+    ...items: T[]
+): Thenable<T | undefined>;
+export function showInformationMessage<T extends MessageItem>(
+    message: string,
+    options: MessageOptions,
+    ...items: T[]
+): Thenable<T | undefined> {
+    return window.showInformationMessage(message, options, ...items);
+}
+
 export function showWarningMessage(message: string, ...items: string[]): Thenable<string | undefined> {
     return window.showWarningMessage(message, ...items);
 }

src/extension.ts

@@ -57,6 +57,8 @@ import { registerTools } from './common/lm.apis';
 import { GetPackagesTool } from './features/copilotTools';
 import { TerminalActivationImpl } from './features/terminal/terminalActivationState';
 import { getEnvironmentForTerminal } from './features/terminal/utils';
+import { PowershellStartupProvider } from './features/terminal/startup/powershellStartup';
+import { ShellStartupActivationManagerImpl } from './features/terminal/startup/activateUsingShellStartup';
 
 export async function activate(context: ExtensionContext): Promise<PythonEnvironmentApi> {
     const start = new StopWatch();
@@ -83,8 +85,14 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
     context.subscriptions.push(envManagers);
 
     const terminalActivation = new TerminalActivationImpl();
-    const terminalManager: TerminalManager = new TerminalManagerImpl(terminalActivation);
-    context.subscriptions.push(terminalActivation, terminalManager);
+    const shellStartupProviders = [new PowershellStartupProvider()];
+    const shellStartupActivationManager = new ShellStartupActivationManagerImpl(
+        context.environmentVariableCollection,
+        shellStartupProviders,
+        envManagers,
+    );
+    const terminalManager: TerminalManager = new TerminalManagerImpl(terminalActivation, shellStartupProviders);
+    context.subscriptions.push(terminalActivation, terminalManager, shellStartupActivationManager);
 
     const projectCreators: ProjectCreators = new ProjectCreatorsImpl();
     context.subscriptions.push(
@@ -94,21 +102,23 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
     );
 
     setPythonApi(envManagers, projectManager, projectCreators, terminalManager, envVarManager);
+    const api = await getPythonApi();
 
     const managerView = new EnvManagerView(envManagers);
     context.subscriptions.push(managerView);
 
     const workspaceView = new ProjectView(envManagers, projectManager);
     context.subscriptions.push(workspaceView);
-
     workspaceView.initialize();
-    const api = await getPythonApi();
 
     const monitoredTerminals = new Map<Terminal, PythonEnvironment>();
 
     context.subscriptions.push(
         registerCompletionProvider(envManagers),
         registerTools('python_get_packages', new GetPackagesTool(api)),
+        commands.registerCommand('python-envs.terminal.revertStartupScriptChanges', async () => {
+            await shellStartupActivationManager.cleanupStartupScripts();
+        }),
         commands.registerCommand('python-envs.viewLogs', () => outputChannel.show()),
         commands.registerCommand('python-envs.refreshManager', async (item) => {
             await refreshManagerCommand(item);
@@ -262,6 +272,7 @@ export async function activate(context: ExtensionContext): Promise<PythonEnviron
         await Promise.all([
             registerSystemPythonFeatures(nativeFinder, context.subscriptions, outputChannel),
             registerCondaFeatures(nativeFinder, context.subscriptions, outputChannel),
+            shellStartupActivationManager.initialize(),
         ]);
         sendTelemetryEvent(EventNames.EXTENSION_MANAGER_REGISTRATION_DURATION, start.elapsedTime);
         await terminalManager.initialize(api);

src/features/common/activation.ts

@@ -15,7 +15,13 @@ export function getActivationCommand(
     environment: PythonEnvironment,
 ): PythonCommandRunConfiguration[] | undefined {
     const shell = identifyTerminalShell(terminal);
+    return getActivationCommandForShell(environment, shell);
+}
 
+export function getActivationCommandForShell(
+    environment: PythonEnvironment,
+    shell: TerminalShellType,
+): PythonCommandRunConfiguration[] | undefined {
     let activation: PythonCommandRunConfiguration[] | undefined;
     if (environment.execInfo?.shellActivation) {
         activation = environment.execInfo.shellActivation.get(shell);

src/features/terminal/startup/activateUsingShellStartup.ts

@@ -1,85 +1,158 @@
-import { Disposable, GlobalEnvironmentVariableCollection } from 'vscode';
-import { onDidChangeConfiguration } from '../../../common/workspace.apis';
-import { registerCommand } from '../../../common/command.api';
-import { getAutoActivationType } from '../utils';
+import { ConfigurationChangeEvent, Disposable, GlobalEnvironmentVariableCollection } from 'vscode';
+import { getWorkspaceFolder, getWorkspaceFolders, onDidChangeConfiguration } from '../../../common/workspace.apis';
+import { getAutoActivationType, setAutoActivationType } from '../utils';
+import { ShellStartupProvider } from './startupProvider';
+import { DidChangeEnvironmentEventArgs } from '../../../api';
 import { EnvironmentManagers } from '../../../internal.api';
+import { traceError, traceInfo } from '../../../common/logging';
+import { ShellStartupActivationStrings } from '../../../common/localize';
+import { showInformationMessage } from '../../../common/window.apis';
 
-export interface ActivateUsingShellStartup extends Disposable {}
+export interface ShellStartupActivationManager extends Disposable {
+    initialize(): Promise<void>;
+    updateStartupScripts(): Promise<void>;
+    cleanupStartupScripts(): Promise<void>;
+}
 
-class ActivateUsingShellStartupImpl implements ActivateUsingShellStartup {
+export class ShellStartupActivationManagerImpl implements ShellStartupActivationManager {
     private readonly disposables: Disposable[] = [];
     constructor(
         private readonly envCollection: GlobalEnvironmentVariableCollection,
+        private readonly shellStartupProviders: ShellStartupProvider[],
         private readonly em: EnvironmentManagers,
     ) {
         this.disposables.push(
-            onDidChangeConfiguration((e) => {
+            onDidChangeConfiguration((e: ConfigurationChangeEvent) => {
                 this.handleConfigurationChange(e);
             }),
+            this.em.onDidChangeEnvironmentFiltered((e: DidChangeEnvironmentEventArgs) => {
+                this.handleEnvironmentChange(e);
+            }),
         );
     }
 
-    private handleConfigurationChange(e) {
-        if (e.affectsConfiguration('python.terminal.autoActivationType')) {
+    private async handleConfigurationChange(e: ConfigurationChangeEvent) {
+        if (e.affectsConfiguration('python-envs.terminal.autoActivationType')) {
             const autoActType = getAutoActivationType();
             if (autoActType === 'shellStartup') {
-                s;
+                await this.initialize();
             } else {
+                // remove any contributed environment variables
+                const workspaces = getWorkspaceFolders() ?? [];
+                if (workspaces.length > 0) {
+                    // User has one or more workspaces open
+                    const promises: Promise<void>[] = [];
+                    workspaces.forEach((workspace) => {
+                        const collection = this.envCollection.getScoped({ workspaceFolder: workspace });
+                        promises.push(
+                            ...this.shellStartupProviders.map((provider) => provider.removeEnvVariables(collection)),
+                        );
+                    });
+                    await Promise.all(promises);
+                } else {
+                    // User has no workspaces open
+                    await Promise.all(
+                        this.shellStartupProviders.map((provider) => provider.removeEnvVariables(this.envCollection)),
+                    );
+                }
             }
         }
     }
 
-    private async addActivationVariables(): Promise<void> {}
-
-    private async removeActivationVariables(): Promise<void> {}
+    private async handleEnvironmentChange(e: DidChangeEnvironmentEventArgs) {
+        const autoActType = getAutoActivationType();
+        if (autoActType !== 'shellStartup') {
+            return;
+        }
 
-    dispose() {
-        this.disposables.forEach((disposable) => disposable.dispose());
+        if (e.uri) {
+            const wf = getWorkspaceFolder(e.uri);
+            if (wf) {
+                const envVars = this.envCollection.getScoped({ workspaceFolder: wf });
+                if (envVars) {
+                    await Promise.all(
+                        this.shellStartupProviders.map(async (provider) => {
+                            if (e.new) {
+                                await provider.updateEnvVariables(envVars, e.new);
+                            } else {
+                                await provider.removeEnvVariables(envVars);
+                            }
+                        }),
+                    );
+                }
+            }
+        }
     }
-}
-
-export async function checkAndUpdateStartupScripts(): Promise<void> {
-    // Implement the logic to check startup scripts
-    return Promise.resolve();
-}
 
-export async function removeAllStartupScripts(): Promise<void> {
-    // Implement the logic to remove all startup scripts
-    return Promise.resolve();
-}
+    private async isSetupRequired(): Promise<boolean> {
+        const results = await Promise.all(this.shellStartupProviders.map((provider) => provider.isSetup()));
+        return results.some((result) => !result);
+    }
 
-export function registerActivateUsingShellStartup(
-    disposables: Disposable[],
-    environmentVariableCollection: GlobalEnvironmentVariableCollection,
-    em: EnvironmentManagers,
-) {
-    let activateUsingShellStartup: ActivateUsingShellStartup | undefined;
+    public async initialize(): Promise<void> {
+        const autoActType = getAutoActivationType();
+        if (autoActType === 'shellStartup') {
+            if (await this.isSetupRequired()) {
+                const result = await showInformationMessage(
+                    ShellStartupActivationStrings.shellStartupScriptEditPrompt,
+                    { modal: true },
+                    ShellStartupActivationStrings.updateScript,
+                );
 
-    disposables.push(
-        onDidChangeConfiguration((e) => {
-            if (e.affectsConfiguration('python.terminal.autoActivationType')) {
-                const autoActType = getAutoActivationType();
-                if (autoActType === 'shellStartup') {
-                    if (!activateUsingShellStartup) {
-                        activateUsingShellStartup = new ActivateUsingShellStartupImpl(environmentVariableCollection);
-                    }
+                if (ShellStartupActivationStrings.updateScript === result) {
+                    await this.updateStartupScripts();
                 } else {
-                    activateUsingShellStartup?.dispose();
-                    activateUsingShellStartup = undefined;
+                    traceError('User declined to edit shell startup scripts. See <doc-link> for more information.');
+                    traceInfo('Setting `python-envs.terminal.autoActivationType` to `command`.');
+                    setAutoActivationType('command');
+                    return;
                 }
             }
-        }),
-        new Disposable(() => activateUsingShellStartup?.dispose()),
-        registerCommand('python-envs.removeStartupScripts', async () => {
-            await removeAllStartupScripts();
-        }),
-    );
 
-    const autoActType = getAutoActivationType();
-    if (autoActType === 'shellStartup') {
-        activateUsingShellStartup = new ActivateUsingShellStartupImpl(environmentVariableCollection);
-        setImmediate(async () => {
-            await checkAndUpdateStartupScripts();
-        });
+            const workspaces = getWorkspaceFolders() ?? [];
+
+            if (workspaces.length > 0) {
+                const promises: Promise<void>[] = [];
+                workspaces.forEach((workspace) => {
+                    const collection = this.envCollection.getScoped({ workspaceFolder: workspace });
+                    promises.push(
+                        ...this.shellStartupProviders.map(async (provider) => {
+                            const env = await this.em.getEnvironment(workspace.uri);
+                            if (env) {
+                                await provider.updateEnvVariables(collection, env);
+                            } else {
+                                await provider.removeEnvVariables(collection);
+                            }
+                        }),
+                    );
+                });
+                await Promise.all(promises);
+            } else {
+                await Promise.all(
+                    this.shellStartupProviders.map(async (provider) => {
+                        const env = await this.em.getEnvironment(undefined);
+                        if (env) {
+                            await provider.updateEnvVariables(this.envCollection, env);
+                        } else {
+                            await provider.removeEnvVariables(this.envCollection);
+                        }
+                    }),
+                );
+            }
+        }
+    }
+
+    public async updateStartupScripts(): Promise<void> {
+        await Promise.all(this.shellStartupProviders.map((provider) => provider.setupScripts()));
+    }
+
+    public async cleanupStartupScripts(): Promise<void> {
+        await Promise.all(this.shellStartupProviders.map((provider) => provider.teardownScripts()));
+        setAutoActivationType('command');
+        showInformationMessage(ShellStartupActivationStrings.revertToCommandActivation);
+    }
+
+    dispose() {
+        this.disposables.forEach((disposable) => disposable.dispose());
     }
 }