From f6da4b4e1a07823b5183b4e25416cceced47d99b Mon Sep 17 00:00:00 2001
From: deepak1556 <hop2deep@gmail.com>
Date: Wed, 4 Mar 2026 13:08:01 +0900
Subject: [PATCH] fix: add graceful shutdown path when heapprofile is enabled

---
 .../src/tsServer/serverProcess.electron.ts    | 79 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 3 deletions(-)

diff --git a/extensions/typescript-language-features/src/tsServer/serverProcess.electron.ts b/extensions/typescript-language-features/src/tsServer/serverProcess.electron.ts
index 992cae925df..a356fd7817b 100644
--- a/extensions/typescript-language-features/src/tsServer/serverProcess.electron.ts
+++ b/extensions/typescript-language-features/src/tsServer/serverProcess.electron.ts
@@ -24,6 +24,12 @@ const contentLengthSize: number = Buffer.byteLength(contentLength, 'utf8');
 const blank: number = Buffer.from(' ', 'utf8')[0];
 const backslashR: number = Buffer.from('\r', 'utf8')[0];
 const backslashN: number = Buffer.from('\n', 'utf8')[0];
+const gracefulExitTimeout = 5000;
+const tsServerExitRequest: Proto.Request = {
+	seq: 0,
+	type: 'request',
+	command: 'exit',
+};
 
 class ProtocolBuffer {
 
@@ -207,10 +213,15 @@ function getTssDebugBrk(): string | undefined {
 }
 
 class IpcChildServerProcess extends Disposable implements TsServerProcess {
+	private _killTimeout: NodeJS.Timeout | undefined;
+	private _isShuttingDown = false;
+
 	constructor(
 		private readonly _process: child_process.ChildProcess,
+		private readonly _useGracefulShutdown: boolean,
 	) {
 		super();
+		this._process.once('exit', () => this.clearKillTimeout());
 	}
 
 	write(serverRequest: Proto.Request): void {
@@ -230,18 +241,47 @@ class IpcChildServerProcess extends Disposable implements TsServerProcess {
 	}
 
 	kill(): void {
-		this._process.kill();
+		if (!this._useGracefulShutdown) {
+			this._process.kill();
+			return;
+		}
+
+		if (this._isShuttingDown) {
+			return;
+		}
+		this._isShuttingDown = true;
+
+		try {
+			this._process.send(tsServerExitRequest);
+		} catch {
+			this._process.kill();
+			return;
+		}
+
+		this._killTimeout = setTimeout(() => this._process.kill(), gracefulExitTimeout);
+		this._killTimeout.unref?.();
+	}
+
+	private clearKillTimeout(): void {
+		if (this._killTimeout) {
+			clearTimeout(this._killTimeout);
+			this._killTimeout = undefined;
+		}
 	}
 }
 
 class StdioChildServerProcess extends Disposable implements TsServerProcess {
 	private readonly _reader: Reader<Proto.Response>;
+	private _killTimeout: NodeJS.Timeout | undefined;
+	private _isShuttingDown = false;
 
 	constructor(
 		private readonly _process: child_process.ChildProcess,
+		private readonly _useGracefulShutdown: boolean,
 	) {
 		super();
 		this._reader = this._register(new Reader<Proto.Response>(this._process.stdout!));
+		this._process.once('exit', () => this.clearKillTimeout());
 	}
 
 	write(serverRequest: Proto.Request): void {
@@ -262,7 +302,39 @@ class StdioChildServerProcess extends Disposable implements TsServerProcess {
 	}
 
 	kill(): void {
-		this._process.kill();
+		if (!this._useGracefulShutdown) {
+			this._process.kill();
+			this._reader.dispose();
+			return;
+		}
+
+		if (this._isShuttingDown) {
+			return;
+		}
+		this._isShuttingDown = true;
+
+		try {
+			this._process.stdin?.write(JSON.stringify(tsServerExitRequest) + '\r\n', 'utf8');
+			this._process.stdin?.end();
+		} catch {
+			this._process.kill();
+			this._reader.dispose();
+			return;
+		}
+
+		this._killTimeout = setTimeout(() => {
+			this._process.kill();
+			this._reader.dispose();
+		}, gracefulExitTimeout);
+		this._killTimeout.unref?.();
+	}
+
+	private clearKillTimeout(): void {
+		if (this._killTimeout) {
+			clearTimeout(this._killTimeout);
+			this._killTimeout = undefined;
+		}
+
 		this._reader.dispose();
 	}
 }
@@ -290,6 +362,7 @@ export class ElectronServiceProcessFactory implements TsServerProcessFactory {
 		const env = generatePatchedEnv(process.env, tsServerPath, !!execPath);
 		const runtimeArgs = [...args];
 		const execArgv = getExecArgv(kind, configuration);
+		const useGracefulShutdown = configuration.heapProfile.enabled;
 		const useIpc = !execPath && version.apiVersion?.gte(API.v460);
 		if (useIpc) {
 			runtimeArgs.push('--useNodeIpc');
@@ -309,6 +382,6 @@ export class ElectronServiceProcessFactory implements TsServerProcessFactory {
 				stdio: useIpc ? ['pipe', 'pipe', 'pipe', 'ipc'] : undefined,
 			});
 
-		return useIpc ? new IpcChildServerProcess(childProcess) : new StdioChildServerProcess(childProcess);
+		return useIpc ? new IpcChildServerProcess(childProcess, useGracefulShutdown) : new StdioChildServerProcess(childProcess, useGracefulShutdown);
 	}
 }