fix: wire ext host restart through utility process api (microsoft#302633)

Author: deepak1556

src/vs/platform/extensions/common/extensionHostStarter.ts

@@ -9,6 +9,7 @@ import { createDecorator } from '../../instantiation/common/instantiation.js';
 export const IExtensionHostStarter = createDecorator<IExtensionHostStarter>('extensionHostStarter');
 
 export const ipcExtensionHostStarterChannelName = 'extensionHostStarter';
+export const extensionHostGraceTimeMs = 6000;
 
 export interface IExtensionHostProcessOptions {
 	responseWindowId: number;
@@ -31,6 +32,7 @@ export interface IExtensionHostStarter {
 	createExtensionHost(): Promise<{ id: string }>;
 	start(id: string, opts: IExtensionHostProcessOptions): Promise<{ pid: number | undefined }>;
 	enableInspectPort(id: string): Promise<boolean>;
+	waitForExit(id: string, maxWaitTimeMs: number): Promise<void>;
 	kill(id: string): Promise<void>;
 
 }

src/vs/platform/extensions/electron-main/extensionHostStarter.ts

@@ -7,7 +7,7 @@ import { Promises } from '../../../base/common/async.js';
 import { canceled } from '../../../base/common/errors.js';
 import { Event } from '../../../base/common/event.js';
 import { Disposable, IDisposable } from '../../../base/common/lifecycle.js';
-import { IExtensionHostProcessOptions, IExtensionHostStarter } from '../common/extensionHostStarter.js';
+import { extensionHostGraceTimeMs, IExtensionHostProcessOptions, IExtensionHostStarter } from '../common/extensionHostStarter.js';
 import { ILifecycleMainService } from '../../lifecycle/electron-main/lifecycleMainService.js';
 import { ILogService } from '../../log/common/log.js';
 import { ITelemetryService } from '../../telemetry/common/telemetry.js';
@@ -121,7 +121,7 @@ export class ExtensionHostStarter extends Disposable implements IDisposable, IEx
 			allowLoadingUnsignedLibraries: true,
 			respondToAuthRequestsFromMainProcess: true,
 			windowLifecycleBound: true,
-			windowLifecycleGraceTime: 6000,
+			windowLifecycleGraceTime: extensionHostGraceTimeMs,
 			correlationId: id
 		});
 		const pid = await Event.toPromise(extHost.onSpawn);
@@ -151,6 +151,17 @@ export class ExtensionHostStarter extends Disposable implements IDisposable, IEx
 		extHostProcess.kill();
 	}
 
+	async waitForExit(id: string, maxWaitTimeMs: number): Promise<void> {
+		if (this._shutdown) {
+			throw canceled();
+		}
+		const extHostProcess = this._extHosts.get(id);
+		if (!extHostProcess) {
+			return;
+		}
+		await extHostProcess.waitForExit(maxWaitTimeMs);
+	}
+
 	async _killAllNow(): Promise<void> {
 		for (const [, extHost] of this._extHosts) {
 			extHost.kill();

src/vs/workbench/services/extensions/electron-browser/localProcessExtensionHost.ts

@@ -19,7 +19,7 @@ import { BufferedEmitter } from '../../../../base/parts/ipc/common/ipc.net.js';
 import { acquirePort } from '../../../../base/parts/ipc/electron-browser/ipc.mp.js';
 import * as nls from '../../../../nls.js';
 import { IExtensionHostDebugService } from '../../../../platform/debug/common/extensionHostDebug.js';
-import { IExtensionHostProcessOptions, IExtensionHostStarter } from '../../../../platform/extensions/common/extensionHostStarter.js';
+import { extensionHostGraceTimeMs, IExtensionHostProcessOptions, IExtensionHostStarter } from '../../../../platform/extensions/common/extensionHostStarter.js';
 import { ILabelService } from '../../../../platform/label/common/label.js';
 import { ILogService, ILoggerService } from '../../../../platform/log/common/log.js';
 import { INativeHostService } from '../../../../platform/native/common/native.js';
@@ -32,7 +32,7 @@ import { IWorkspaceContextService, WorkbenchState, isUntitledWorkspace } from '.
 import { INativeWorkbenchEnvironmentService } from '../../environment/electron-browser/environmentService.js';
 import { IShellEnvironmentService } from '../../environment/electron-browser/shellEnvironmentService.js';
 import { MessagePortExtHostConnection, writeExtHostConnection } from '../common/extensionHostEnv.js';
-import { IExtensionHostInitData, MessageType, NativeLogMarkers, UIKind, isMessageOfType } from '../common/extensionHostProtocol.js';
+import { createMessageOfType, IExtensionHostInitData, MessageType, NativeLogMarkers, UIKind, isMessageOfType } from '../common/extensionHostProtocol.js';
 import { LocalProcessRunningLocation } from '../common/extensionRunningLocation.js';
 import { ExtensionHostExtensions, ExtensionHostStartup, IExtensionHost, IExtensionInspectInfo } from '../common/extensions.js';
 import { IHostService } from '../../host/browser/host.js';
@@ -83,6 +83,10 @@ export class ExtensionHostProcess {
 		return this._extensionHostStarter.enableInspectPort(this._id);
 	}
 
+	public waitForExit(maxWaitTimeMs: number): Promise<void> {
+		return this._extensionHostStarter.waitForExit(this._id, maxWaitTimeMs);
+	}
+
 	public kill(): Promise<void> {
 		return this._extensionHostStarter.kill(this._id);
 	}
@@ -161,14 +165,39 @@ export class NativeLocalProcessExtensionHost extends Disposable implements IExte
 	}
 
 	public override dispose(): void {
-		if (this._terminating) {
-			return;
+		if (!this._terminating) {
+			this._terminating = true;
 		}
-		this._terminating = true;
 		super.dispose();
 		this._messageProtocol = null;
 	}
 
+	public async disconnect(): Promise<void> {
+		this._terminating = true;
+
+		if (this._messageProtocol) {
+			try {
+				const protocol = await Promise.race([
+					this._messageProtocol.then(protocol => protocol, () => undefined),
+					timeout(1000).then(() => undefined)
+				]);
+				protocol?.send(createMessageOfType(MessageType.Terminate));
+			} catch {
+				// ignore - extension host may have already exited
+			}
+		}
+
+		if (this._extensionHostProcess) {
+			try {
+				await this._extensionHostProcess.waitForExit(extensionHostGraceTimeMs);
+			} catch {
+				// best-effort: waitForExit may reject with canceled() if the main side is already shutting down
+			}
+		}
+
+		this._messageProtocol = null;
+	}
+
 	public start(): Promise<IMessagePassingProtocol> {
 		if (this._terminating) {
 			// .terminate() was called

test/smoke/extensions/vscode-smoketest-ext-host/extension.js

@@ -16,19 +16,26 @@ let deactivateMarkerFile;
  * @param {vscode.ExtensionContext} context
  */
 function activate(context) {
+	// Record extension host pid on every activation so smoke tests can validate
+	// that a new extension host process was started after a restart action.
+	try {
+		const pid = String(process.pid);
+		const activationPidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid-on-activate.txt');
+		fs.writeFileSync(activationPidFile, pid, 'utf-8');
+	} catch {
+		// Ignore errors in smoke helper setup.
+	}
+
 	// This is used to verify that the extension host process is properly killed
 	// when window reloads even if the extension host is blocked
 	// Refs: https://github.com/microsoft/vscode/issues/291346
 	context.subscriptions.push(
 		vscode.commands.registerCommand('smoketest.getExtensionHostPidAndBlock', (delayMs = 100, durationMs = 60000) => {
 			const pid = process.pid;
 
-			// Write PID file to workspace folder if available, otherwise temp dir
+			// Write PID file to temp dir to avoid polluting workspace search results
 			// Note: filename must match name in extension-host-restart.test.ts
-			const workspaceFolder = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath;
-			const pidFile = workspaceFolder
-				? path.join(workspaceFolder, 'vscode-ext-host-pid.txt')
-				: path.join(os.tmpdir(), 'vscode-ext-host-pid.txt');
+			const pidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid.txt');
 			setTimeout(() => {
 				fs.writeFileSync(pidFile, String(pid), 'utf-8');
 
@@ -57,13 +64,8 @@ function activate(context) {
 	context.subscriptions.push(
 		vscode.commands.registerCommand('smoketest.setupGracefulDeactivation', () => {
 			const pid = process.pid;
-			const workspaceFolder = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath;
-			const pidFile = workspaceFolder
-				? path.join(workspaceFolder, 'vscode-ext-host-pid-graceful.txt')
-				: path.join(os.tmpdir(), 'vscode-ext-host-pid-graceful.txt');
-			deactivateMarkerFile = workspaceFolder
-				? path.join(workspaceFolder, 'vscode-ext-host-deactivated.txt')
-				: path.join(os.tmpdir(), 'vscode-ext-host-deactivated.txt');
+			const pidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid-graceful.txt');
+			deactivateMarkerFile = path.join(os.tmpdir(), 'vscode-ext-host-deactivated.txt');
 
 			// Write PID file immediately so test knows the extension is ready
 			fs.writeFileSync(pidFile, String(pid), 'utf-8');

test/smoke/src/areas/extensions/extension-host-restart.test.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as fs from 'fs';
+import * as os from 'os';
 import * as path from 'path';
 import { Application, Logger } from '../../../../automation';
 import { installAllHandlers, timeout } from '../../utils';
@@ -30,7 +31,7 @@ export function setup(logger: Logger) {
 			this.timeout(60_000);
 
 			const app = this.app as Application;
-			const pidFile = path.join(app.workspacePathOrFolder, 'vscode-ext-host-pid.txt');
+			const pidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid.txt');
 
 			if (fs.existsSync(pidFile)) {
 				fs.unlinkSync(pidFile);
@@ -78,8 +79,8 @@ export function setup(logger: Logger) {
 			this.timeout(60_000);
 
 			const app = this.app as Application;
-			const pidFile = path.join(app.workspacePathOrFolder, 'vscode-ext-host-pid-graceful.txt');
-			const markerFile = path.join(app.workspacePathOrFolder, 'vscode-ext-host-deactivated.txt');
+			const pidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid-graceful.txt');
+			const markerFile = path.join(os.tmpdir(), 'vscode-ext-host-deactivated.txt');
 
 			// Clean up any existing files
 			if (fs.existsSync(pidFile)) {
@@ -139,5 +140,79 @@ export function setup(logger: Logger) {
 
 			logger.log('Extension host was properly terminated after graceful deactivation');
 		});
+
+		it('kills blocked extension host on restart extension host (issue #296681)', async function () {
+			this.timeout(90_000);
+
+			const app = this.app as Application;
+			const pidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid.txt');
+			const activationPidFile = path.join(os.tmpdir(), 'vscode-ext-host-pid-on-activate.txt');
+
+			if (fs.existsSync(pidFile)) {
+				fs.unlinkSync(pidFile);
+			}
+
+			await app.workbench.quickaccess.runCommand('smoketest.getExtensionHostPidAndBlock');
+
+			let retries = 0;
+			while (!fs.existsSync(pidFile) && retries < 20) {
+				await timeout(500);
+				retries++;
+			}
+
+			if (!fs.existsSync(pidFile)) {
+				throw new Error('PID file was not created - extension may not have activated');
+			}
+
+			const oldPid = parseInt(fs.readFileSync(pidFile, 'utf-8'), 10);
+			logger.log(`Old extension host PID: ${oldPid}`);
+
+			if (fs.existsSync(activationPidFile)) {
+				fs.unlinkSync(activationPidFile);
+			}
+
+			await app.workbench.quickaccess.runCommand('Developer: Restart Extension Host', { keepOpen: true });
+
+			const maxWaitMs = 10_000;
+			const pollIntervalMs = 500;
+			let waitedMs = 0;
+
+			let newPid: number | undefined;
+			while (waitedMs < maxWaitMs) {
+				if (fs.existsSync(activationPidFile)) {
+					const pidText = fs.readFileSync(activationPidFile, 'utf-8').trim();
+					const parsedPid = parseInt(pidText, 10);
+					if (!Number.isNaN(parsedPid) && parsedPid !== oldPid) {
+						newPid = parsedPid;
+						break;
+					}
+				}
+
+				await timeout(pollIntervalMs);
+				waitedMs += pollIntervalMs;
+			}
+
+			if (!newPid) {
+				throw new Error(`New extension host PID was not observed after restart (waited ${maxWaitMs}ms)`);
+			}
+
+			if (newPid === oldPid) {
+				throw new Error(`Extension host PID did not change after restart (pid: ${oldPid})`);
+			}
+
+			logger.log(`New extension host PID observed: ${newPid}`);
+
+			waitedMs = 0;
+			while (processExists(oldPid) && waitedMs < maxWaitMs) {
+				await timeout(pollIntervalMs);
+				waitedMs += pollIntervalMs;
+			}
+
+			if (processExists(oldPid)) {
+				throw new Error(`Old extension host ${oldPid} still running after restart (waited ${maxWaitMs}ms)`);
+			}
+
+			logger.log(`Extension host restarted successfully (old: ${oldPid}, new: ${newPid})`);
+		});
 	});
 }