Fix renderer losing connection on pty host restart

Fixes microsoft#186772
Fixes microsoft#186773

Author: Tyriar

src/vs/platform/terminal/electron-main/electronPtyHostStarter.ts

@@ -14,12 +14,12 @@ import { UtilityProcess } from 'vs/platform/utilityProcess/electron-main/utility
 import { Client as MessagePortClient } from 'vs/base/parts/ipc/electron-main/ipc.mp';
 import { IpcMainEvent } from 'electron';
 import { validatedIpcMain } from 'vs/base/parts/ipc/electron-main/ipcMain';
-import { DisposableStore, toDisposable } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore, toDisposable } from 'vs/base/common/lifecycle';
 import { Emitter } from 'vs/base/common/event';
 import { deepClone } from 'vs/base/common/objects';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 
-export class ElectronPtyHostStarter implements IPtyHostStarter {
+export class ElectronPtyHostStarter extends Disposable implements IPtyHostStarter {
 
 	private utilityProcess: UtilityProcess | undefined = undefined;
 
@@ -35,9 +35,14 @@ export class ElectronPtyHostStarter implements IPtyHostStarter {
 		@ILifecycleMainService private readonly _lifecycleMainService: ILifecycleMainService,
 		@ILogService private readonly _logService: ILogService
 	) {
+		super();
+
 		this._lifecycleMainService.onWillShutdown(() => this._onWillShutdown.fire());
 		// Listen for new windows to establish connection directly to pty host
 		validatedIpcMain.on('vscode:createPtyHostMessageChannel', (e, nonce) => this._onWindowConnection(e, nonce));
+		this._register(toDisposable(() => {
+			validatedIpcMain.removeHandler('vscode:createPtyHostMessageChannel');
+		}));
 	}
 
 	start(lastPtyId: number): IPtyHostConnection {
@@ -62,9 +67,9 @@ export class ElectronPtyHostStarter implements IPtyHostStarter {
 
 		const store = new DisposableStore();
 		store.add(client);
-		store.add(this.utilityProcess);
 		store.add(toDisposable(() => {
-			validatedIpcMain.removeHandler('vscode:createPtyHostMessageChannel');
+			this.utilityProcess?.kill();
+			this.utilityProcess?.dispose();
 			this.utilityProcess = undefined;
 		}));
 

src/vs/platform/terminal/node/nodePtyHostStarter.ts

@@ -3,19 +3,20 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { DisposableStore } from 'vs/base/common/lifecycle';
+import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { FileAccess } from 'vs/base/common/network';
 import { Client, IIPCOptions } from 'vs/base/parts/ipc/node/ipc.cp';
 import { IEnvironmentService, INativeEnvironmentService } from 'vs/platform/environment/common/environment';
 import { parsePtyHostDebugPort } from 'vs/platform/environment/node/environmentService';
 import { IReconnectConstants } from 'vs/platform/terminal/common/terminal';
 import { IPtyHostConnection, IPtyHostStarter } from 'vs/platform/terminal/node/ptyHost';
 
-export class NodePtyHostStarter implements IPtyHostStarter {
+export class NodePtyHostStarter extends Disposable implements IPtyHostStarter {
 	constructor(
 		private readonly _reconnectConstants: IReconnectConstants,
 		@IEnvironmentService private readonly _environmentService: INativeEnvironmentService
 	) {
+		super();
 	}
 
 	start(lastPtyId: number): IPtyHostConnection {

src/vs/platform/terminal/node/ptyHost.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Event } from 'vs/base/common/event';
-import { DisposableStore } from 'vs/base/common/lifecycle';
+import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
 import { IChannelClient } from 'vs/base/parts/ipc/common/ipc';
 
 export interface IPtyHostConnection {
@@ -13,7 +13,7 @@ export interface IPtyHostConnection {
 	readonly onDidProcessExit: Event<{ code: number; signal: string }>;
 }
 
-export interface IPtyHostStarter {
+export interface IPtyHostStarter extends IDisposable {
 	onBeforeWindowConnection?: Event<void>;
 	onWillShutdown?: Event<void>;
 

src/vs/platform/terminal/node/ptyHostService.ts

@@ -52,7 +52,7 @@ export class PtyHostService extends Disposable implements IPtyService {
 
 	private _ensurePtyHost() {
 		if (!this.__connection) {
-			[this.__connection, this.__proxy] = this._startPtyHost();
+			this._startPtyHost();
 		}
 	}
 
@@ -102,9 +102,9 @@ export class PtyHostService extends Disposable implements IPtyService {
 		// remote server).
 		registerTerminalPlatformConfiguration();
 
+		this._register(this._ptyHostStarter);
 		this._register(toDisposable(() => this._disposePtyHost()));
 
-
 		this._resolveVariablesRequestStore = this._register(new RequestStore(undefined, this._logService));
 		this._resolveVariablesRequestStore.onCreateRequest(this._onPtyHostRequestResolveVariables.fire, this._onPtyHostRequestResolveVariables);
 
@@ -145,8 +145,6 @@ export class PtyHostService extends Disposable implements IPtyService {
 		const connection = this._ptyHostStarter.start(lastPtyId);
 		const client = connection.client;
 
-		this._onPtyHostStart.fire();
-
 		// Setup heartbeat service and trigger a heartbeat immediately to reset the timeouts
 		const heartbeatService = ProxyChannel.toService<IHeartbeatService>(client.getChannel(TerminalIpcChannels.Heartbeat));
 		heartbeatService.onBeat(() => this._handleHeartbeat());
@@ -181,6 +179,8 @@ export class PtyHostService extends Disposable implements IPtyService {
 		this.__connection = connection;
 		this.__proxy = proxy;
 
+		this._onPtyHostStart.fire();
+
 		this._register(this._configurationService.onDidChangeConfiguration(async e => {
 			if (e.affectsConfiguration(TerminalSettingId.IgnoreProcessNames)) {
 				await this._refreshIgnoreProcessNames();

src/vs/workbench/contrib/terminal/browser/baseTerminalBackend.ts

@@ -22,6 +22,8 @@ export abstract class BaseTerminalBackend extends Disposable {
 
 	get isResponsive(): boolean { return !this._isPtyHostUnresponsive; }
 
+	protected readonly _onPtyHostConnected = this._register(new Emitter<void>());
+	readonly onPtyHostConnected = this._onPtyHostConnected.event;
 	protected readonly _onPtyHostRestart = this._register(new Emitter<void>());
 	readonly onPtyHostRestart = this._onPtyHostRestart.event;
 	protected readonly _onPtyHostUnresponsive = this._register(new Emitter<void>());
@@ -50,13 +52,14 @@ export abstract class BaseTerminalBackend extends Disposable {
 			}));
 		}
 		if (this._ptyHostController.onPtyHostStart) {
+			this.onPtyHostConnected(() => hasStarted = true);
 			this._register(this._ptyHostController.onPtyHostStart(() => {
 				this._logService.debug(`The terminal's pty host process is starting`);
-				// Only fire the event on the 2nd
+				// Only fire the _restart_ event after it has started
 				if (hasStarted) {
+					this._logService.trace('IPtyHostController#onPtyHostRestart');
 					this._onPtyHostRestart.fire();
 				}
-				hasStarted = true;
 				statusBarAccessor?.dispose();
 				this._isPtyHostUnresponsive = false;
 			}));

src/vs/workbench/contrib/terminal/electron-sandbox/localTerminalBackend.ts

@@ -107,6 +107,7 @@ class LocalTerminalBackend extends BaseTerminalBackend implements ITerminalBacke
 			// separate interface/service for this one.
 			const client = new MessagePortClient(port, `window:${this._environmentService.window.id}`);
 			clientEventually.complete(client);
+			this._onPtyHostConnected.fire();
 
 			// Attach process listeners
 			this._proxy.onProcessData(e => this._ptys.get(e.id)?.handleData(e.event));
@@ -190,7 +191,6 @@ class LocalTerminalBackend extends BaseTerminalBackend implements ITerminalBacke
 		shouldPersist: boolean
 	): Promise<ITerminalChildProcess> {
 		const executableEnv = await this._shellEnvironmentService.getShellEnv();
-		// TODO: Using _proxy here bypasses the lastPtyId tracking on the main process
 		const id = await this._proxy.createProcess(shellLaunchConfig, cwd, cols, rows, unicodeVersion, env, executableEnv, options, shouldPersist, this._getWorkspaceId(), this._getWorkspaceName());
 		const pty = new LocalPty(id, shouldPersist, this._proxy);
 		this._ptys.set(id, pty);