Resolve shell env API performance issue for single env approach

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

@@ -157,7 +157,9 @@ export interface IShellEnvDetectionCapability {
 	setEnvironment(envs: { [key: string]: string | undefined } | undefined, isTrusted: boolean): void;
 	startEnvironmentSingleVar(isTrusted: boolean): void;
 	setEnvironmentSingleVar(key: string, value: string | undefined, isTrusted: boolean): void;
+	deleteEnvironmentSingleVar(key: string, value: string | undefined, isTrusted: boolean): void;
 	endEnvironmentSingleVar(isTrusted: boolean): void;
+	applyEnvironmentDiff(env: Map<string, string>, isTrusted: boolean): void;
 }
 
 export const enum CommandInvalidationReason {

src/vs/platform/terminal/common/capabilities/shellEnvDetectionCapability.ts

@@ -14,7 +14,6 @@ export class ShellEnvDetectionCapability extends Disposable implements IShellEnv
 	private _pendingEnv: Map<string, string> | undefined;
 	private _env: Map<string, string> = new Map();
 	get env(): Map<string, string> { return this._env; }
-
 	private readonly _onDidChangeEnv = this._register(new Emitter<Map<string, string>>());
 	readonly onDidChangeEnv = this._onDidChangeEnv.event;
 
@@ -59,8 +58,44 @@ export class ShellEnvDetectionCapability extends Disposable implements IShellEnv
 		if (!this._pendingEnv) {
 			return;
 		}
-		this._env = this._pendingEnv;
+		this.applyEnvironmentDiff(this._pendingEnv, isTrusted);
 		this._pendingEnv = undefined;
-		this._onDidChangeEnv.fire(this._env);
+	}
+
+	deleteEnvironmentSingleVar(key: string, value: string | undefined, isTrusted: boolean): void {
+		if (!isTrusted) {
+			return;
+		}
+		if (key !== undefined && value !== undefined) {
+			this._env.delete(key);
+			this._pendingEnv?.delete(key);
+		}
+		return;
+	}
+
+	// Make sure to update this.env to the latest, fire event if there is a diff
+	applyEnvironmentDiff(env: Map<string, string>, isTrusted: boolean): void {
+		if (!isTrusted) {
+			return;
+		}
+
+		let envDiffers: boolean = false;
+
+		for (const [key, value] of env.entries()) {
+			if (this._env.has(key) && this._env.get(key) === value) {
+				// Do nothing
+			} else if (this.env.has(key) && value !== this.env.get(key)) {
+				this._env.set(key, value);
+				envDiffers = true;
+			} else if (!this.env.has(key)) {
+				this._env.set(key, value);
+				envDiffers = true;
+			}
+		}
+
+		if (envDiffers) {
+			this._onDidChangeEnv.fire(this._env);
+			return;
+		}
 	}
 }

src/vs/platform/terminal/common/xterm/shellIntegrationAddon.ts

@@ -239,6 +239,16 @@ const enum VSCodeOscPt {
 	 */
 	EnvJson = 'EnvJson',
 
+	/**
+	 * Delete a single environment variable from cached environment.
+	 *
+	 * Format: `OSC 633 ; EnvSingleDelete ; <EnvironmentKey> ; <EnvironmentValue> ; <Nonce>`
+	 *
+	 * WARNING: This sequence is unfinalized, DO NOT use this in your shell integration script.
+	 */
+
+	EnvSingleDelete = 'EnvSingleDelete',
+
 	/**
 	 * The start of the collecting user's environment variables individually.
 	 * Clears any environment residuals in previous sessions.
@@ -467,7 +477,7 @@ export class ShellIntegrationAddon extends Disposable implements IShellIntegrati
 				if (arg0 !== undefined) {
 					try {
 						const env = JSON.parse(deserializeMessage(arg0));
-						this._createOrGetShellEnvDetection().setEnvironment(env, arg1 === this._nonce);
+						this._createOrGetShellEnvDetection().applyEnvironmentDiff(env, arg1 === this._nonce);
 					} catch (e) {
 						this._logService.warn('Failed to parse environment from shell integration sequence', arg0);
 					}
@@ -478,6 +488,16 @@ export class ShellIntegrationAddon extends Disposable implements IShellIntegrati
 				this._createOrGetShellEnvDetection().startEnvironmentSingleVar(args[0] === this._nonce);
 				return true;
 			}
+			case VSCodeOscPt.EnvSingleDelete: {
+				const arg0 = args[0];
+				const arg1 = args[1];
+				const arg2 = args[2];
+				if (arg0 !== undefined && arg1 !== undefined) {
+					const env = deserializeMessage(arg1);
+					this._createOrGetShellEnvDetection().deleteEnvironmentSingleVar(arg0, env, arg2 === this._nonce);
+				}
+				return true;
+			}
 			case VSCodeOscPt.EnvSingleEntry: {
 				const arg0 = args[0];
 				const arg1 = args[1];

src/vs/workbench/api/common/extHostTerminalShellIntegration.ts

@@ -55,30 +55,30 @@ export class ExtHostTerminalShellIntegration extends Disposable implements IExtH
 		}));
 
 		// Convenient test code:
-		// this.onDidChangeTerminalShellIntegration(e => {
-		// 	console.log('*** onDidChangeTerminalShellIntegration', e);
-		// });
-		// this.onDidStartTerminalShellExecution(async e => {
-		// 	console.log('*** onDidStartTerminalShellExecution', e);
-		// 	// new Promise<void>(r => {
-		// 	// 	(async () => {
-		// 	// 		for await (const d of e.execution.read()) {
-		// 	// 			console.log('data2', d);
-		// 	// 		}
-		// 	// 	})();
-		// 	// });
-		// 	for await (const d of e.execution.read()) {
-		// 		console.log('data', d);
-		// 	}
-		// });
-		// this.onDidEndTerminalShellExecution(e => {
-		// 	console.log('*** onDidEndTerminalShellExecution', e);
-		// });
-		// setTimeout(() => {
-		// 	console.log('before executeCommand(\"echo hello\")');
-		// 	Array.from(this._activeShellIntegrations.values())[0].value.executeCommand('echo hello');
-		// 	console.log('after executeCommand(\"echo hello\")');
-		// }, 4000);
+		this.onDidChangeTerminalShellIntegration(e => {
+			console.log('*** onDidChangeTerminalShellIntegration', e.shellIntegration.env);
+		});
+		this.onDidStartTerminalShellExecution(async e => {
+			console.log('*** onDidStartTerminalShellExecution', e);
+			// new Promise<void>(r => {
+			// 	(async () => {
+			// 		for await (const d of e.execution.read()) {
+			// 			console.log('data2', d);
+			// 		}
+			// 	})();
+			// });
+			for await (const d of e.execution.read()) {
+				console.log('data', d);
+			}
+		});
+		this.onDidEndTerminalShellExecution(e => {
+			console.log('*** onDidEndTerminalShellExecution', e);
+		});
+		setTimeout(() => {
+			console.log('before executeCommand(\"echo hello\")');
+			Array.from(this._activeShellIntegrations.values())[0].value.executeCommand('echo hello');
+			console.log('after executeCommand(\"echo hello\")');
+		}, 4000);
 	}
 
 	public $shellIntegrationChange(instanceId: number): void {

src/vs/workbench/contrib/terminal/common/scripts/shellIntegration-bash.sh

@@ -10,6 +10,9 @@ fi
 
 VSCODE_SHELL_INTEGRATION=1
 
+vsc_env_keys=()
+vsc_env_values=()
+
 # Run relevant rc/profile only if shell integration has been injected, not when run manually
 if [ "$VSCODE_INJECTION" == "1" ]; then
 	if [ -z "$VSCODE_SHELL_LOGIN" ]; then
@@ -214,16 +217,93 @@ __vsc_update_cwd() {
 	builtin printf '\e]633;P;Cwd=%s\a' "$(__vsc_escape_value "$__vsc_cwd")"
 }
 
-# __vsc_update_env() {
-# 	builtin printf '\e]633;EnvSingleStart;%s;\a' $__vsc_nonce
-# 	for var in $(compgen -v); do
-# 		if printenv "$var" >/dev/null 2>&1; then
-# 			value=$(builtin printf '%s' "${!var}")
-# 			builtin printf '\e]633;EnvSingleEntry;%s;%s;%s\a' "$var" "$(__vsc_escape_value "$value")" $__vsc_nonce
-# 		fi
-# 	done
-# 	builtin printf '\e]633;EnvSingleEnd;%s;\a' $__vsc_nonce
-# }
+
+updateEnvCache() {
+	local key="$1"
+	local value="$2"
+
+	# builtin printf "Start %s\n" $(date +%s%3N)
+	for i in "${!vsc_env_keys[@]}"; do
+		# builtin printf "i $i $(date +%s%3N)\n"
+		if [[ "${vsc_env_keys[$i]}" == "$key" ]]; then
+			if [[ "${vsc_env_values[$i]}" != "$value" ]]; then
+				vsc_env_values[$i]="$value"
+				# printf "Updated value for %s\n" "$key"
+				builtin printf '\e]633;EnvSingleEntry;%s;%s;%s\a' "$key" "$(__vsc_escape_value "$value")" "$__vsc_nonce"
+			fi
+			return
+		fi
+	done
+
+	vsc_env_keys+=("$key")
+	vsc_env_values+=("$value")
+	# printf "Added new variable %s\n" "$key"
+	builtin printf '\e]633;EnvSingleEntry;%s;%s;%s\a' "$key" "$(__vsc_escape_value "$value")" "$__vsc_nonce"
+}
+
+trackMissingEnvVars() {
+	# Capture current environment variables in an array
+	local current_env_keys=()
+
+	# TODO: I think this will break for values that contain `=`, see `cut` command in `VSCODE_ENV_REPLACE` code
+	while IFS='=' read -r key value; do
+		current_env_keys+=("$key")
+	done < <(env)
+
+	# Compare vsc_env_keys with current_env_keys
+	for key in "${vsc_env_keys[@]}"; do
+		local found=false
+		for env_key in "${current_env_keys[@]}"; do
+			if [[ "$key" == "$env_key" ]]; then
+				# TODO: Use 1 and 0 over true
+				found=true
+				break
+			fi
+		done
+		if [[ "$found" == false ]]; then
+			builtin printf '\e]633;EnvSingleDelete;%s;%s;%s\a' "${vsc_env_keys[i]}" "$(__vsc_escape_value "${vsc_env_values[i]}")" "$__vsc_nonce"
+			# TODO: Use `builtin unset`
+			unset 'vsc_env_keys[i]'
+			unset 'vsc_env_values[i]'
+		fi
+	done
+
+	# Remove gaps from unset
+	vsc_env_keys=("${vsc_env_keys[@]}")
+	vsc_env_values=("${vsc_env_values[@]}")
+}
+
+__vsc_update_env() {
+
+	# start_time=$(date +%s%3N)
+	# builtin printf "Start time: %s\n" "$start_time"
+	# TODO: Send everything the first time, don't bother diffing (array is empty)
+
+	builtin printf '\e]633;EnvSingleStart;%s;\a' $__vsc_nonce
+
+	# This surprisingly doesnt break even when value contains = or multiple =
+	while IFS='=' read -r key value; do
+		updateEnvCache "$key" "$value"
+	done < <(env)
+
+	# IMPORTANT: Trying to use cut like this with echo really slow things down.
+	# while IFS= read -r line; do
+	# 	key=$(echo "$line" | cut -d '=' -f 1)
+	# 	value=$(echo "$line" | cut -d '=' -f 2-)
+	# 	updateEnvCache "$key" "$value"
+	# done < <(env)
+
+	# mid_time=$(date +%s%3N)
+	trackMissingEnvVars
+
+	builtin printf '\e]633;EnvSingleEnd;%s;\a' $__vsc_nonce
+
+	# end_time=$(date +%s%3N)
+	# elapsed_time=$(($end_time - $start_time))
+	# builtin printf "updateEnvCache: $(($mid_time - $start_time)) ms\n"
+	# builtin printf "trackMissingEnvVars: $(($end_time - $mid_time)) ms\n"
+	# builtin printf "Total: $(($end_time - $start_time)) ms\n"
+}
 
 __vsc_command_output_start() {
 	if [[ -z "${__vsc_first_prompt-}" ]]; then
@@ -251,10 +331,10 @@ __vsc_command_complete() {
 		builtin printf '\e]633;D;%s\a' "$__vsc_status"
 	fi
 	__vsc_update_cwd
-
-	# if [ "$__vsc_stable" = "0" ]; then
-	# 	__vsc_update_env
-	# fi
+	# IMPORTANT: We may have to bring back __vsc_update_env here, as now it takes one additional prompt execution for us to see unset env in effect
+	if [ "$__vsc_stable" = "0" ]; then
+		__vsc_update_env
+	fi
 }
 __vsc_update_prompt() {
 	# in command execution
@@ -284,10 +364,9 @@ __vsc_precmd() {
 	fi
 	__vsc_first_prompt=1
 	__vsc_update_prompt
-
-	# if [ "$__vsc_stable" = "0" ]; then
-	# 	__vsc_update_env
-	# fi
+	if [ "$__vsc_stable" = "0" ]; then
+		__vsc_update_env
+	fi
 }
 
 __vsc_preexec() {