Merge pull request #77824 from microsoft/roblou/remoteEnvReconnection

Fix remote EH environment after reconnection

src/vs/workbench/api/common/extHost.protocol.ts

@@ -876,6 +876,7 @@ export interface ExtHostExtensionServiceShape {
 	$startExtensionHost(enabledExtensionIds: ExtensionIdentifier[]): Promise<void>;
 	$activateByEvent(activationEvent: string): Promise<void>;
 	$activate(extensionId: ExtensionIdentifier, activationEvent: string): Promise<boolean>;
+	$setRemoteEnvironment(env: { [key: string]: string | null }): Promise<void>;
 
 	$deltaExtensions(toAdd: IExtensionDescription[], toRemove: ExtensionIdentifier[]): Promise<void>;
 

src/vs/workbench/api/node/extHostExtensionService.ts

@@ -749,6 +749,20 @@ export class ExtHostExtensionService implements ExtHostExtensionServiceShape {
 		return buff;
 	}
 
+	public async $setRemoteEnvironment(env: { [key: string]: string | null }): Promise<void> {
+		if (!this._initData.remote.isRemote) {
+			return;
+		}
+
+		for (const key in env) {
+			const value = env[key];
+			if (value === null) {
+				delete process.env[key];
+			} else {
+				process.env[key] = value;
+			}
+		}
+	}
 }
 
 function loadCommonJSModule<T>(logService: ILogService, modulePath: string, activationTimesBuilder: ExtensionActivationTimesBuilder): Promise<T> {

src/vs/workbench/contrib/remote/electron-browser/remote.contribution.ts

@@ -284,6 +284,26 @@ class ProgressReporter {
 	}
 }
 
+class RemoteExtensionHostEnvironmentUpdater implements IWorkbenchContribution {
+	constructor(
+		@IRemoteAgentService remoteAgentService: IRemoteAgentService,
+		@IRemoteAuthorityResolverService remoteResolverService: IRemoteAuthorityResolverService,
+		@IExtensionService extensionService: IExtensionService
+	) {
+		const connection = remoteAgentService.getConnection();
+		if (connection) {
+			connection.onDidStateChange(async e => {
+				if (e.type === PersistentConnectionEventType.ConnectionGain) {
+					const resolveResult = await remoteResolverService.resolveAuthority(connection.remoteAuthority);
+					if (resolveResult.options && resolveResult.options.extensionHostEnv) {
+						await extensionService.setRemoteEnvironment(resolveResult.options.extensionHostEnv);
+					}
+				}
+			});
+		}
+	}
+}
+
 class RemoteAgentConnectionStatusListener implements IWorkbenchContribution {
 	constructor(
 		@IRemoteAgentService remoteAgentService: IRemoteAgentService,
@@ -440,6 +460,7 @@ const workbenchContributionsRegistry = Registry.as<IWorkbenchContributionsRegist
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteChannelsContribution, LifecyclePhase.Starting);
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteAgentDiagnosticListener, LifecyclePhase.Eventually);
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteAgentConnectionStatusListener, LifecyclePhase.Eventually);
+workbenchContributionsRegistry.registerWorkbenchContribution(RemoteExtensionHostEnvironmentUpdater, LifecyclePhase.Eventually);
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteWindowActiveIndicator, LifecyclePhase.Starting);
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteTelemetryEnablementUpdater, LifecyclePhase.Ready);
 workbenchContributionsRegistry.registerWorkbenchContribution(RemoteEmptyWorkbenchPresentation, LifecyclePhase.Starting);

src/vs/workbench/services/extensions/common/abstractExtensionService.ts

@@ -263,6 +263,11 @@ export abstract class AbstractExtensionService extends Disposable implements IEx
 		return 0;
 	}
 
+	public async setRemoteEnvironment(env: { [key: string]: string | null }): Promise<void> {
+		await this._extensionHostProcessManagers
+			.map(manager => manager.setRemoteEnvironment(env));
+	}
+
 	//#endregion
 
 	// --- impl

src/vs/workbench/services/extensions/common/extensionHostProcessManager.ts

@@ -288,6 +288,15 @@ export class ExtensionHostProcessManager extends Disposable {
 		}
 		return proxy.$deltaExtensions(toAdd, toRemove);
 	}
+
+	public async setRemoteEnvironment(env: { [key: string]: string | null }): Promise<void> {
+		const proxy = await this._getExtensionHostProcessProxy();
+		if (!proxy) {
+			return;
+		}
+
+		return proxy.$setRemoteEnvironment(env);
+	}
 }
 
 const colorTables = [

src/vs/workbench/services/extensions/common/extensions.ts

@@ -229,6 +229,12 @@ export interface IExtensionService {
 	 */
 	stopExtensionHost(): void;
 
+	/**
+	 * Modify the environment of the remote extension host
+	 * @param env New properties for the remote extension host
+	 */
+	setRemoteEnvironment(env: { [key: string]: string | null }): Promise<void>;
+
 	_logOrShowMessage(severity: Severity, msg: string): void;
 	_activateById(extensionId: ExtensionIdentifier, activationEvent: string): Promise<void>;
 	_onWillActivateExtension(extensionId: ExtensionIdentifier): void;
@@ -278,6 +284,7 @@ export class NullExtensionService implements IExtensionService {
 	restartExtensionHost(): void { }
 	startExtensionHost(): void { }
 	stopExtensionHost(): void { }
+	async setRemoteEnvironment(_env: { [key: string]: string | null }): Promise<void> { }
 	canAddExtension(): boolean { return false; }
 	canRemoveExtension(): boolean { return false; }
 	_logOrShowMessage(_severity: Severity, _msg: string): void { }