Merge pull request #179472 from microsoft/tyriar/apply_si

Support VSCODE_ENV_* variables to adjust environment in shell integration scripts

src/vs/base/common/processes.ts

@@ -108,7 +108,7 @@ export function sanitizeProcessEnvironment(env: IProcessEnvironment, ...preserve
 	}, {} as Record<string, boolean>);
 	const keysToRemove = [
 		/^ELECTRON_.+$/,
-		/^VSCODE_(?!(PORTABLE|SHELL_LOGIN)).+$/,
+		/^VSCODE_(?!(PORTABLE|SHELL_LOGIN|ENV_REPLACE|ENV_APPEND|ENV_PREPEND)).+$/,
 		/^SNAP(|_.*)$/,
 		/^GDK_PIXBUF_.+$/,
 	];

src/vs/platform/terminal/common/environmentVariable.ts

@@ -10,9 +10,15 @@ export enum EnvironmentVariableMutatorType {
 	Append = 2,
 	Prepend = 3
 }
+// export enum EnvironmentVariableMutatorTiming {
+// 	AtSpawn = 1,
+// 	AfterShellIntegration = 2
+// 	// TODO: Do we need a both?
+// }
 export interface IEnvironmentVariableMutator {
 	readonly value: string;
 	readonly type: EnvironmentVariableMutatorType;
+	// readonly timing?: EnvironmentVariableMutatorTiming;
 }
 
 export interface IEnvironmentVariableCollection {

src/vs/platform/terminal/common/environmentVariableCollection.ts

@@ -8,6 +8,12 @@ import { EnvironmentVariableMutatorType, IEnvironmentVariableCollection, IExtens
 
 type VariableResolver = (str: string) => Promise<string>;
 
+// const mutatorTypeToLabelMap: Map<EnvironmentVariableMutatorType, string> = new Map([
+// 	[EnvironmentVariableMutatorType.Append, 'APPEND'],
+// 	[EnvironmentVariableMutatorType.Prepend, 'PREPEND'],
+// 	[EnvironmentVariableMutatorType.Replace, 'REPLACE']
+// ]);
+
 export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVariableCollection {
 	readonly map: Map<string, IExtensionOwnedEnvironmentVariableMutator[]> = new Map();
 
@@ -55,6 +61,11 @@ export class MergedEnvironmentVariableCollection implements IMergedEnvironmentVa
 			const actualVariable = isWindows ? lowerToActualVariableNames![variable.toLowerCase()] || variable : variable;
 			for (const mutator of mutators) {
 				const value = variableResolver ? await variableResolver(mutator.value) : mutator.value;
+				// if (mutator.timing === EnvironmentVariableMutatorTiming.AfterShellIntegration) {
+				// 	const key = `VSCODE_ENV_${mutatorTypeToLabelMap.get(mutator.type)!}`;
+				// 	env[key] = (env[key] ? env[key] + ':' : '') + variable + '=' + value;
+				// 	continue;
+				// }
 				switch (mutator.type) {
 					case EnvironmentVariableMutatorType.Append:
 						env[actualVariable] = (env[actualVariable] || '') + value;

src/vs/workbench/contrib/terminal/browser/media/shellIntegration-bash.sh

@@ -45,6 +45,35 @@ if [ -z "$VSCODE_SHELL_INTEGRATION" ]; then
 	builtin return
 fi
 
+# Apply EnvironmentVariableCollections if needed
+if [ -n "$VSCODE_ENV_REPLACE" ]; then
+	IFS=':' read -ra ADDR <<< "$VSCODE_ENV_REPLACE"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo $ITEM | cut -d "=" -f 1)"
+		VALUE="$(echo $ITEM | cut -d "=" -f 2)"
+		export $VARNAME="$VALUE"
+	done
+	builtin unset VSCODE_ENV_REPLACE
+fi
+if [ -n "$VSCODE_ENV_PREPEND" ]; then
+	IFS=':' read -ra ADDR <<< "$VSCODE_ENV_PREPEND"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo $ITEM | cut -d "=" -f 1)"
+		VALUE="$(echo $ITEM | cut -d "=" -f 2)"
+		export $VARNAME="$VALUE${!VARNAME}"
+	done
+	builtin unset VSCODE_ENV_REPLACE
+fi
+if [ -n "$VSCODE_ENV_APPEND" ]; then
+	IFS=':' read -ra ADDR <<< "$VSCODE_ENV_APPEND"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo $ITEM | cut -d "=" -f 1)"
+		VALUE="$(echo $ITEM | cut -d "=" -f 2)"
+		export $VARNAME="${!VARNAME}$VALUE"
+	done
+	builtin unset VSCODE_ENV_REPLACE
+fi
+
 __vsc_get_trap() {
 	# 'trap -p DEBUG' outputs a shell command like `trap -- '…shellcode…' DEBUG`.
 	# The terms are quoted literals, but are not guaranteed to be on a single line.

src/vs/workbench/contrib/terminal/browser/media/shellIntegration-rc.zsh

@@ -38,6 +38,35 @@ if [ -z "$VSCODE_SHELL_INTEGRATION" ]; then
 	builtin return
 fi
 
+# Apply EnvironmentVariableCollections if needed
+if [ -n "$VSCODE_ENV_REPLACE" ]; then
+	echo "VSCODE_ENV_REPLACE: $VSCODE_ENV_REPLACE"
+	IFS=':' read -rA ADDR <<< "$VSCODE_ENV_REPLACE"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo ${ITEM%%=*})"
+		export $VARNAME="${ITEM#*=}"
+	done
+	unset VSCODE_ENV_REPLACE
+fi
+if [ -n "$VSCODE_ENV_PREPEND" ]; then
+	echo "VSCODE_ENV_PREPEND: $VSCODE_ENV_PREPEND"
+	IFS=':' read -rA ADDR <<< "$VSCODE_ENV_PREPEND"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo ${ITEM%%=*})"
+		export $VARNAME="${ITEM#*=}${(P)VARNAME}"
+	done
+	unset VSCODE_ENV_PREPEND
+fi
+if [ -n "$VSCODE_ENV_APPEND" ]; then
+	echo "VSCODE_ENV_APPEND: $VSCODE_ENV_APPEND"
+	IFS=':' read -rA ADDR <<< "$VSCODE_ENV_APPEND"
+	for ITEM in "${ADDR[@]}"; do
+		VARNAME="$(echo ${ITEM%%=*})"
+		export $VARNAME="${(P)VARNAME}${ITEM#*=}"
+	done
+	unset VSCODE_ENV_APPEND
+fi
+
 # The property (P) and command (E) codes embed values which require escaping.
 # Backslashes are doubled. Non-alphanumeric characters are converted to escaped hex.
 __vsc_escape_value() {

src/vs/workbench/contrib/terminal/browser/media/shellIntegration.fish

@@ -28,6 +28,32 @@ if status --is-login; and set -q VSCODE_PATH_PREFIX
 end
 set -e VSCODE_PATH_PREFIX
 
+# Apply EnvironmentVariableCollections if needed
+if test -n "$VSCODE_ENV_REPLACE"
+	set ITEMS $(string split : $VSCODE_ENV_REPLACE)
+	for B in $ITEMS
+		set split $(string split = $B)
+		set -gx "$split[1]" "$split[2]"
+	end
+	set -e VSCODE_ENV_REPLACE
+end
+if test -n "$VSCODE_ENV_PREPEND"
+	set ITEMS $(string split : $VSCODE_ENV_PREPEND)
+	for B in $ITEMS
+		set split $(string split = $B)
+		set -gx "$split[1]" "$split[2]$$split[1]" # avoid -p as it adds a space
+	end
+	set -e VSCODE_ENV_PREPEND
+end
+if test -n "$VSCODE_ENV_APPEND"
+	set ITEMS $(string split : $VSCODE_ENV_APPEND)
+	for B in $ITEMS
+		set split $(string split = $B)
+		set -gx "$split[1]" "$$split[1]$split[2]" # avoid -a as it adds a space
+	end
+	set -e VSCODE_ENV_APPEND
+end
+
 # Helper function
 function __vsc_esc -d "Emit escape sequences for VS Code shell integration"
 	builtin printf "\e]633;%s\a" (string join ";" $argv)

src/vs/workbench/contrib/terminal/browser/media/shellIntegration.ps1

@@ -17,6 +17,31 @@ $Global:__VSCodeOriginalPrompt = $function:Prompt
 
 $Global:__LastHistoryId = -1
 
+if ($env:VSCODE_ENV_REPLACE) {
+	$Split = $env:VSCODE_ENV_REPLACE.Split(":")
+	foreach ($Item in $Split) {
+		$Inner = $Item.Split('=')
+		[Environment]::SetEnvironmentVariable($Inner[0], $Inner[1])
+	}
+	$env:VSCODE_ENV_REPLACE = $null
+}
+if ($env:VSCODE_ENV_PREPEND) {
+	$Split = $env:VSCODE_ENV_PREPEND.Split(":")
+	foreach ($Item in $Split) {
+		$Inner = $Item.Split('=')
+		[Environment]::SetEnvironmentVariable($Inner[0], $Inner[1] + [Environment]::GetEnvironmentVariable($Inner[0]))
+	}
+	$env:VSCODE_ENV_PREPEND = $null
+}
+if ($env:VSCODE_ENV_APPEND) {
+	$Split = $env:VSCODE_ENV_APPEND.Split(":")
+	foreach ($Item in $Split) {
+		$Inner = $Item.Split('=')
+		[Environment]::SetEnvironmentVariable($Inner[0], [Environment]::GetEnvironmentVariable($Inner[0]) + $Inner[1])
+	}
+	$env:VSCODE_ENV_APPEND = $null
+}
+
 function Global:__VSCode-Escape-Value([string]$value) {
 	# NOTE: In PowerShell v6.1+, this can be written `$value -replace '…', { … }` instead of `[regex]::Replace`.
 	# Replace any non-alphanumeric characters.