dispose extension pseudoterminals when terminals close

Author: SimonSiefke

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

@@ -309,7 +309,7 @@ export class ExtHostTerminal extends Disposable {
 	}
 }
 
-class ExtHostPseudoterminal implements ITerminalChildProcess {
+class ExtHostPseudoterminal extends Disposable implements ITerminalChildProcess {
 	readonly id = 0;
 	readonly shouldPersist = false;
 
@@ -322,7 +322,9 @@ class ExtHostPseudoterminal implements ITerminalChildProcess {
 	private readonly _onProcessExit = new Emitter<number | undefined>();
 	public readonly onProcessExit: Event<number | undefined> = this._onProcessExit.event;
 
-	constructor(private readonly _pty: vscode.Pseudoterminal) { }
+	constructor(private readonly _pty: vscode.Pseudoterminal) {
+		super();
+	}
 
 	refreshProperty<T extends ProcessPropertyType>(property: ProcessPropertyType): Promise<IProcessPropertyMap[T]> {
 		throw new Error(`refreshProperty is not suppported in extension owned terminals. property: ${property}`);
@@ -380,18 +382,18 @@ class ExtHostPseudoterminal implements ITerminalChildProcess {
 
 	startSendingEvents(initialDimensions: ITerminalDimensionsDto | undefined): void {
 		// Attach the listeners
-		this._pty.onDidWrite(e => this._onProcessData.fire(e));
-		this._pty.onDidClose?.((e: number | void = undefined) => {
+		this._register(this._pty.onDidWrite(e => this._onProcessData.fire(e)));
+		this._register(this._pty.onDidClose?.((e: number | void = undefined) => {
 			this._onProcessExit.fire(e === void 0 ? undefined : e);
-		});
-		this._pty.onDidOverrideDimensions?.(e => {
+		}) ?? Disposable.None);
+		this._register(this._pty.onDidOverrideDimensions?.(e => {
 			if (e) {
 				this._onDidChangeProperty.fire({ type: ProcessPropertyType.OverrideDimensions, value: { cols: e.columns, rows: e.rows } });
 			}
-		});
-		this._pty.onDidChangeName?.(title => {
+		}) ?? Disposable.None);
+		this._register(this._pty.onDidChangeName?.(title => {
 			this._onDidChangeProperty.fire({ type: ProcessPropertyType.Title, value: title });
-		});
+		}) ?? Disposable.None);
 
 		this._pty.open(initialDimensions ? initialDimensions : undefined);
 
@@ -520,8 +522,10 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 		const terminal = new ExtHostTerminal(this._proxy, generateUuid(), options, options.name);
 		const p = new ExtHostPseudoterminal(options.pty);
 		terminal.createExtensionTerminal(options.location, internalOptions, this._serializeParentTerminal(options, internalOptions).resolvedExtHostIdentifier, asTerminalIcon(options.iconPath), asTerminalColor(options.color), options.shellIntegrationNonce, options.titleTemplate).then(id => {
-			const disposable = this._setupExtHostProcessListeners(id, p);
-			this._terminalProcessDisposables[id] = disposable;
+			if (!this.getTerminalById(id)) {
+				return;
+			}
+			this._setupExtHostProcessListeners(id, p);
 		});
 		this._terminals.push(terminal);
 		return terminal.value;
@@ -551,8 +555,7 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 			throw new Error(`Cannot resolve terminal with id ${id} for virtual process`);
 		}
 		const p = new ExtHostPseudoterminal(pty);
-		const disposable = this._setupExtHostProcessListeners(id, p);
-		this._terminalProcessDisposables[id] = disposable;
+		this._setupExtHostProcessListeners(id, p);
 	}
 
 	public async $acceptActiveTerminalChanged(id: number | null): Promise<void> {
@@ -613,6 +616,8 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 	}
 
 	public async $acceptTerminalClosed(id: number, exitCode: number | undefined, exitReason: TerminalExitReason): Promise<void> {
+		this._cleanUpTerminalProcess(id);
+
 		// Release any cached terminal links and cancel in-flight link providers for this terminal
 		this._terminalLinkCache.delete(id);
 		const cancellationSource = this._terminalLinkCancellationSource.get(id);
@@ -694,23 +699,26 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 		return undefined;
 	}
 
-	protected _setupExtHostProcessListeners(id: number, p: ITerminalChildProcess): IDisposable {
+	protected _setupExtHostProcessListeners(id: number, p: ITerminalChildProcess): void {
 		const disposables = new DisposableStore();
+		if (p instanceof ExtHostPseudoterminal) {
+			disposables.add(p);
+		}
 		disposables.add(p.onProcessReady(e => this._proxy.$sendProcessReady(id, e.pid, e.cwd, e.windowsPty)));
 		disposables.add(p.onDidChangeProperty(property => this._proxy.$sendProcessProperty(id, property)));
 
 		// Buffer data events to reduce the amount of messages going to the renderer
-		this._bufferer.startBuffering(id, p.onProcessData);
+		disposables.add(this._bufferer.startBuffering(id, p.onProcessData));
 		disposables.add(p.onProcessExit(exitCode => this._onProcessExit(id, exitCode)));
 		this._terminalProcesses.set(id, p);
+		this._terminalProcessDisposables[id] = disposables;
 
 		const awaitingStart = this._extensionTerminalAwaitingStart[id];
 		if (awaitingStart && p instanceof ExtHostPseudoterminal) {
 			p.startSendingEvents(awaitingStart.initialDimensions);
 			delete this._extensionTerminalAwaitingStart[id];
 		}
 
-		return disposables;
 	}
 
 	public $acceptProcessAckDataEvent(id: number, charCount: number): void {
@@ -980,6 +988,13 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 	}
 
 	private _onProcessExit(id: number, exitCode: number | undefined): void {
+		this._cleanUpTerminalProcess(id);
+
+		// Send exit event to main side
+		this._proxy.$sendProcessExit(id, exitCode);
+	}
+
+	private _cleanUpTerminalProcess(id: number): void {
 		this._bufferer.stopBuffering(id);
 
 		// Remove process reference
@@ -992,8 +1007,6 @@ export abstract class BaseExtHostTerminalService extends Disposable implements I
 			processDiposable.dispose();
 			delete this._terminalProcessDisposables[id];
 		}
-		// Send exit event to main side
-		this._proxy.$sendProcessExit(id, exitCode);
 	}
 
 	public getTerminalById(id: number): ExtHostTerminal | null {

src/vs/workbench/api/test/common/extHostTerminalService.test.ts

@@ -5,6 +5,7 @@
 
 import * as assert from 'assert';
 import type * as vscode from 'vscode';
+import { Emitter } from '../../../../base/common/event.js';
 import { mock } from '../../../../base/test/common/mock.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js';
 import { IShellLaunchConfigDto } from '../../../../platform/terminal/common/terminal.js';
@@ -78,4 +79,62 @@ suite('ExtHostTerminalService', () => {
 			{ providerTokenCancelled: true, firstLinks: [], linksAfterClose: [], handledAfterClose: false }
 		);
 	});
+
+	test('extension terminal processes are released when terminals close', async () => {
+		const rpcProtocol = new TestRPCProtocol();
+		let processExitCalls = 0;
+		rpcProtocol.set(MainContext.MainThreadTerminalService, new class extends mock<MainThreadTerminalServiceShape>() {
+			override async $registerProcessSupport(): Promise<void> { }
+			override async $sendProcessExit(): Promise<void> { processExitCalls++; }
+			override async $sendProcessReady(): Promise<void> { }
+		});
+
+		const commands = new class extends mock<ExtHostCommands>() {
+			override registerArgumentProcessor(_processor: ArgumentProcessor): void { }
+		};
+		const initData = new class extends mock<IExtHostInitDataService>() {
+			override readonly remote = { authority: 'test+remote', isRemote: true, connectionData: null };
+		};
+		const service = store.add(new WorkerExtHostTerminalService(commands, rpcProtocol, initData));
+
+		const terminalId = 42;
+		service.$acceptTerminalOpened(terminalId, undefined, 'test', {} as IShellLaunchConfigDto);
+		const terminal = service.getTerminalById(terminalId)!;
+		store.add(terminal);
+
+		const writeEmitter = store.add(new Emitter<string>());
+		const closeEmitter = store.add(new Emitter<number | void>());
+		const inputs: string[] = [];
+		let closeOnOpen = false;
+		const pty: vscode.Pseudoterminal = {
+			onDidWrite: writeEmitter.event,
+			onDidClose: closeEmitter.event,
+			open(): void { if (closeOnOpen) { closeEmitter.fire(); } },
+			close(): void { },
+			handleInput(data: string): void { inputs.push(data); }
+		};
+		service.attachPtyToTerminal(terminalId, pty);
+		await service.$startExtensionTerminal(terminalId, undefined);
+		await rpcProtocol.sync();
+
+		service.$acceptProcessInput(terminalId, 'before');
+		await service.$acceptTerminalClosed(terminalId, undefined, TerminalExitReason.Unknown);
+		service.$acceptProcessInput(terminalId, 'after');
+		closeEmitter.fire();
+		await rpcProtocol.sync();
+
+		assert.deepStrictEqual(inputs, ['before']);
+		assert.strictEqual(processExitCalls, 0);
+
+		const synchronousTerminalId = 43;
+		service.$acceptTerminalOpened(synchronousTerminalId, undefined, 'test', {} as IShellLaunchConfigDto);
+		store.add(service.getTerminalById(synchronousTerminalId)!);
+		closeOnOpen = true;
+		service.attachPtyToTerminal(synchronousTerminalId, pty);
+		await service.$startExtensionTerminal(synchronousTerminalId, undefined);
+		await rpcProtocol.sync();
+
+		assert.strictEqual(closeEmitter.hasListeners(), false);
+		assert.strictEqual(processExitCalls, 1);
+	});
 });