dispose of all terminals for session on dispose or archive (#299816)

src/vs/workbench/contrib/chat/browser/agentSessions/agentSessionsService.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { Emitter, Event } from '../../../../../base/common/event.js';
 import { Disposable } from '../../../../../base/common/lifecycle.js';
 import { URI } from '../../../../../base/common/uri.js';
 import { createDecorator, IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
@@ -13,6 +14,7 @@ export interface IAgentSessionsService {
 	readonly _serviceBrand: undefined;
 
 	readonly model: IAgentSessionsModel;
+	readonly onDidChangeSessionArchivedState: Event<IAgentSession>;
 
 	getSession(resource: URI): IAgentSession | undefined;
 }
@@ -20,11 +22,14 @@ export interface IAgentSessionsService {
 export class AgentSessionsService extends Disposable implements IAgentSessionsService {
 
 	declare readonly _serviceBrand: undefined;
+	private readonly _onDidChangeSessionArchivedState = this._register(new Emitter<IAgentSession>());
+	readonly onDidChangeSessionArchivedState = this._onDidChangeSessionArchivedState.event;
 
 	private _model: IAgentSessionsModel | undefined;
 	get model(): IAgentSessionsModel {
 		if (!this._model) {
 			this._model = this._register(this.instantiationService.createInstance(AgentSessionsModel));
+			this._register(this._model.onDidChangeSessionArchivedState(session => this._onDidChangeSessionArchivedState.fire(session)));
 			this._model.resolve(undefined /* all providers */);
 		}
 

src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts

@@ -10,7 +10,7 @@ import { Codicon } from '../../../../../../base/common/codicons.js';
 import { CancellationError } from '../../../../../../base/common/errors.js';
 import { Event } from '../../../../../../base/common/event.js';
 import { MarkdownString, type IMarkdownString } from '../../../../../../base/common/htmlContent.js';
-import { Disposable, DisposableStore } from '../../../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, MutableDisposable } from '../../../../../../base/common/lifecycle.js';
 import { ResourceMap } from '../../../../../../base/common/map.js';
 import { basename, posix, win32 } from '../../../../../../base/common/path.js';
 import { OperatingSystem, OS } from '../../../../../../base/common/platform.js';
@@ -69,6 +69,7 @@ import { TerminalChatCommandId } from '../../../chat/browser/terminalChat.js';
 import { clamp } from '../../../../../../base/common/numbers.js';
 import { IOutputAnalyzer } from './outputAnalyzer.js';
 import { SandboxOutputAnalyzer } from './sandboxOutputAnalyzer.js';
+import { IAgentSessionsService } from '../../../../chat/browser/agentSessions/agentSessionsService.js';
 
 // #region Tool data
 
@@ -322,8 +323,11 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 	private readonly _commandLineAnalyzers: ICommandLineAnalyzer[];
 	private readonly _commandLinePresenters: ICommandLinePresenter[];
 	private readonly _outputAnalyzers: IOutputAnalyzer[];
+	private readonly _archivedSessionListener = this._register(new MutableDisposable());
 
 	protected readonly _sessionTerminalAssociations = new ResourceMap<IToolTerminal>();
+	protected readonly _sessionTerminalInstances = new ResourceMap<Set<ITerminalInstance>>();
+	private readonly _terminalsBeingDisposedBySessionCleanup = new Set<ITerminalInstance>();
 
 	// Immutable window state
 	protected readonly _osBackend: Promise<OperatingSystem>;
@@ -373,6 +377,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 		@ITerminalService private readonly _terminalService: ITerminalService,
 		@IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService,
 		@IChatWidgetService private readonly _chatWidgetService: IChatWidgetService,
+		@IAgentSessionsService private readonly _agentSessionsService: IAgentSessionsService,
 	) {
 		super();
 
@@ -417,11 +422,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 		// Restore terminal associations from storage
 		this._restoreTerminalAssociations();
 		this._register(this._terminalService.onDidDisposeInstance(e => {
-			for (const [sessionResource, toolTerminal] of this._sessionTerminalAssociations.entries()) {
-				if (e === toolTerminal.instance) {
-					this._sessionTerminalAssociations.delete(sessionResource);
-				}
-			}
+			this._removeTerminalAssociations(e);
 		}));
 
 		// Listen for chat session disposal to clean up associated terminals
@@ -430,6 +431,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 				this._cleanupSessionTerminals(resource);
 			}
 		}));
+
 	}
 
 	async prepareToolInvocation(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
@@ -1108,7 +1110,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 		this._terminalChatService.registerTerminalInstanceWithToolSession(terminalToolSessionId, toolTerminal.instance);
 		this._terminalChatService.registerTerminalInstanceWithChatSession(chatSessionResource, toolTerminal.instance);
 		this._registerInputListener(toolTerminal);
-		this._sessionTerminalAssociations.set(chatSessionResource, toolTerminal);
+		this._addSessionTerminalAssociation(chatSessionResource, toolTerminal);
 		if (token.isCancellationRequested) {
 			toolTerminal.instance.dispose();
 			throw new CancellationError();
@@ -1149,7 +1151,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 							shellIntegrationQuality: association.shellIntegrationQuality,
 							isBackground: association.isBackground
 						};
-						this._sessionTerminalAssociations.set(chatSessionResource, toolTerminal);
+						this._addSessionTerminalAssociation(chatSessionResource, toolTerminal);
 						this._terminalChatService.registerTerminalInstanceWithChatSession(chatSessionResource, instance);
 
 						// Listen for terminal disposal to clean up storage
@@ -1220,23 +1222,83 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
 	}
 
 	private _cleanupSessionTerminals(chatSessionResource: URI): void {
+		const sessionTerminals = this._sessionTerminalInstances.get(chatSessionResource);
 		const toolTerminal = this._sessionTerminalAssociations.get(chatSessionResource);
-		if (toolTerminal) {
-			this._logService.debug(`RunInTerminalTool: Cleaning up terminal for disposed chat session ${chatSessionResource}`);
+		const terminalsToDispose = sessionTerminals ?? (toolTerminal ? new Set([toolTerminal.instance]) : undefined);
+		if (!terminalsToDispose || terminalsToDispose.size === 0) {
+			return;
+		}
 
-			this._sessionTerminalAssociations.delete(chatSessionResource);
-			toolTerminal.instance.dispose();
+		this._logService.debug(`RunInTerminalTool: Cleaning up ${terminalsToDispose.size} terminal(s) for ended chat session ${chatSessionResource}`);
 
-			// Clean up any active executions associated with this session
-			const terminalToRemove: string[] = [];
-			for (const [termId, execution] of RunInTerminalTool._activeExecutions.entries()) {
-				if (execution.instance === toolTerminal.instance) {
-					execution.dispose();
-					terminalToRemove.push(termId);
-				}
+		this._sessionTerminalAssociations.delete(chatSessionResource);
+		this._sessionTerminalInstances.delete(chatSessionResource);
+
+		for (const terminal of terminalsToDispose) {
+			// Skip redundant map walks in onDidDispose since this session has already been removed.
+			this._terminalsBeingDisposedBySessionCleanup.add(terminal);
+			terminal.dispose();
+		}
+
+		// Clean up any active executions associated with this session
+		const terminalToRemove: string[] = [];
+		for (const [termId, execution] of RunInTerminalTool._activeExecutions.entries()) {
+			if (terminalsToDispose.has(execution.instance)) {
+				execution.dispose();
+				terminalToRemove.push(termId);
 			}
-			for (const termId of terminalToRemove) {
-				RunInTerminalTool._activeExecutions.delete(termId);
+		}
+		for (const termId of terminalToRemove) {
+			RunInTerminalTool._activeExecutions.delete(termId);
+		}
+	}
+
+	private _addSessionTerminalAssociation(chatSessionResource: URI, toolTerminal: IToolTerminal): void {
+		this._ensureArchivedSessionListener();
+
+		let sessionTerminals = this._sessionTerminalInstances.get(chatSessionResource);
+		if (!sessionTerminals) {
+			sessionTerminals = new Set<ITerminalInstance>();
+			this._sessionTerminalInstances.set(chatSessionResource, sessionTerminals);
+		}
+		sessionTerminals.add(toolTerminal.instance);
+
+		if (!toolTerminal.isBackground) {
+			this._sessionTerminalAssociations.set(chatSessionResource, toolTerminal);
+		}
+	}
+
+	private _ensureArchivedSessionListener(): void {
+		if (this._archivedSessionListener.value) {
+			return;
+		}
+
+		// Archiving a session does not fire onDidDisposeSession, but we still need to dispose
+		// any terminals associated with the archived session to avoid process accumulation.
+		this._archivedSessionListener.value = this._agentSessionsService.onDidChangeSessionArchivedState(session => {
+			if (session.isArchived()) {
+				this._cleanupSessionTerminals(session.resource);
+			}
+		});
+	}
+
+	private _removeTerminalAssociations(terminal: ITerminalInstance): void {
+		if (this._terminalsBeingDisposedBySessionCleanup.delete(terminal)) {
+			return;
+		}
+
+		for (const [sessionResource, toolTerminal] of this._sessionTerminalAssociations.entries()) {
+			if (terminal === toolTerminal.instance) {
+				this._sessionTerminalAssociations.delete(sessionResource);
+			}
+		}
+
+		for (const [sessionResource, sessionTerminals] of this._sessionTerminalInstances.entries()) {
+			if (!sessionTerminals.delete(terminal)) {
+				continue;
+			}
+			if (sessionTerminals.size === 0) {
+				this._sessionTerminalInstances.delete(sessionResource);
 			}
 		}
 	}

src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/runInTerminalTool.test.ts

@@ -41,11 +41,14 @@ import { ShellIntegrationQuality } from '../../browser/toolTerminalCreator.js';
 import { terminalChatAgentToolsConfiguration, TerminalChatAgentToolsSettingId } from '../../common/terminalChatAgentToolsConfiguration.js';
 import { TerminalChatService } from '../../../chat/browser/terminalChatService.js';
 import type { IMarkdownString } from '../../../../../../base/common/htmlContent.js';
+import { IAgentSessionsService } from '../../../../chat/browser/agentSessions/agentSessionsService.js';
+import { IAgentSession } from '../../../../chat/browser/agentSessions/agentSessionsModel.js';
 
 class TestRunInTerminalTool extends RunInTerminalTool {
 	protected override _osBackend: Promise<OperatingSystem> = Promise.resolve(OperatingSystem.Windows);
 
 	get sessionTerminalAssociations() { return this._sessionTerminalAssociations; }
+	get sessionTerminalInstances() { return this._sessionTerminalInstances; }
 	get profileFetcher() { return this._profileFetcher; }
 
 	setBackendOs(os: OperatingSystem) {
@@ -63,6 +66,7 @@ suite('RunInTerminalTool', () => {
 	let workspaceContextService: TestContextService;
 	let terminalServiceDisposeEmitter: Emitter<ITerminalInstance>;
 	let chatServiceDisposeEmitter: Emitter<{ sessionResource: URI[]; reason: 'cleared' }>;
+	let chatSessionArchivedEmitter: Emitter<IAgentSession>;
 
 	let runInTerminalTool: TestRunInTerminalTool;
 
@@ -79,6 +83,7 @@ suite('RunInTerminalTool', () => {
 		setConfig(TerminalChatAgentToolsSettingId.BlockDetectedFileWrites, 'outsideWorkspace');
 		terminalServiceDisposeEmitter = new Emitter<ITerminalInstance>();
 		chatServiceDisposeEmitter = new Emitter<{ sessionResource: URI[]; reason: 'cleared' }>();
+		chatSessionArchivedEmitter = new Emitter<IAgentSession>();
 
 		instantiationService = workbenchInstantiationService({
 			configurationService: () => configurationService,
@@ -89,6 +94,12 @@ suite('RunInTerminalTool', () => {
 			onDidDisposeSession: chatServiceDisposeEmitter.event,
 			getSession: () => undefined,
 		});
+		instantiationService.stub(IAgentSessionsService, {
+			onDidChangeSessionArchivedState: chatSessionArchivedEmitter.event,
+			model: {
+				onDidChangeSessionArchivedState: chatSessionArchivedEmitter.event,
+			} as IAgentSessionsService['model']
+		});
 		instantiationService.stub(ITerminalChatService, store.add(instantiationService.createInstance(TerminalChatService)));
 		instantiationService.stub(IWorkspaceContextService, workspaceContextService);
 		instantiationService.stub(IHistoryService, {
@@ -505,8 +516,7 @@ suite('RunInTerminalTool', () => {
 
 			// Verify that auto-approve information is included
 			ok(result?.toolSpecificData, 'Expected toolSpecificData to be defined');
-			// eslint-disable-next-line local/code-no-any-casts
-			const terminalData = result!.toolSpecificData as any;
+			const terminalData = result!.toolSpecificData as IChatTerminalToolInvocationData;
 			ok(terminalData.autoApproveInfo, 'Expected autoApproveInfo to be defined for auto-approved background command');
 			ok(terminalData.autoApproveInfo.value, 'Expected autoApproveInfo to have a value');
 			ok(terminalData.autoApproveInfo.value.includes('npm'), 'Expected autoApproveInfo to mention the approved rule');
@@ -1140,13 +1150,149 @@ suite('RunInTerminalTool', () => {
 	});
 
 	suite('chat session disposal cleanup', () => {
+		const createMockTerminal = (processId: number): ITerminalInstance => ({
+			dispose: () => { /* Mock dispose */ },
+			processId
+		} as unknown as ITerminalInstance);
+
+		test('should restore all terminals into the session terminal map and dispose them when archived', () => {
+			const sessionId = 'test-session-restored-archive';
+			const sessionResource = LocalChatSessionUri.forSession(sessionId);
+
+			let terminal1Disposed = false;
+			let terminal2Disposed = false;
+			const terminal1DisposedEmitter = new Emitter<void>();
+			const terminal2DisposedEmitter = new Emitter<void>();
+			const mockTerminal1 = {
+				dispose: () => {
+					terminal1Disposed = true;
+					terminal1DisposedEmitter.fire();
+				},
+				onDisposed: terminal1DisposedEmitter.event,
+				processId: 55555,
+			} as unknown as ITerminalInstance;
+			const mockTerminal2 = {
+				dispose: () => {
+					terminal2Disposed = true;
+					terminal2DisposedEmitter.fire();
+				},
+				onDisposed: terminal2DisposedEmitter.event,
+				processId: 66666,
+			} as unknown as ITerminalInstance;
+
+			storageService.store('chat.terminalSessions', JSON.stringify({
+				[mockTerminal1.processId!]: {
+					sessionId,
+					id: 'restored-1',
+					shellIntegrationQuality: ShellIntegrationQuality.None,
+					isBackground: true,
+				},
+				[mockTerminal2.processId!]: {
+					sessionId,
+					id: 'restored-2',
+					shellIntegrationQuality: ShellIntegrationQuality.None,
+					isBackground: false,
+				}
+			}), StorageScope.WORKSPACE, StorageTarget.USER);
+
+			instantiationService.stub(ITerminalService, {
+				onDidDisposeInstance: terminalServiceDisposeEmitter.event,
+				instances: [mockTerminal1, mockTerminal2],
+				setNextCommandId: async () => { }
+			});
+
+			const restoredRunInTerminalTool = store.add(instantiationService.createInstance(TestRunInTerminalTool));
+			const restoredSessionTerminals = restoredRunInTerminalTool.sessionTerminalInstances.get(sessionResource);
+			strictEqual(restoredSessionTerminals?.size, 2, 'Both restored terminals should be tracked for the session');
+
+			chatSessionArchivedEmitter.fire({
+				resource: sessionResource,
+				isArchived: () => true,
+			} as unknown as IAgentSession);
+
+			strictEqual(terminal1Disposed, true, 'Restored background terminal should have been disposed');
+			strictEqual(terminal2Disposed, true, 'Restored foreground terminal should have been disposed');
+			ok(!restoredRunInTerminalTool.sessionTerminalAssociations.has(sessionResource), 'Foreground terminal association should be removed after archive');
+			ok(!restoredRunInTerminalTool.sessionTerminalInstances.has(sessionResource), 'All restored terminals for the session should be removed after archive');
+		});
+
+		test('should dispose all terminals associated with a single chat session when archived', () => {
+			const sessionId = 'test-session-archive';
+			const sessionResource = LocalChatSessionUri.forSession(sessionId);
+			const mockTerminal1 = { dispose: () => { /* Mock dispose */ }, processId: 33333 } as unknown as ITerminalInstance;
+			const mockTerminal2 = { dispose: () => { /* Mock dispose */ }, processId: 44444 } as unknown as ITerminalInstance;
+
+			let terminal1Disposed = false;
+			let terminal2Disposed = false;
+			mockTerminal1.dispose = () => { terminal1Disposed = true; };
+			mockTerminal2.dispose = () => { terminal2Disposed = true; };
+
+			runInTerminalTool.sessionTerminalAssociations.set(sessionResource, {
+				instance: mockTerminal2,
+				shellIntegrationQuality: ShellIntegrationQuality.None
+			});
+			runInTerminalTool.sessionTerminalInstances.set(sessionResource, new Set([mockTerminal1, mockTerminal2]));
+
+			// Initialize lazy archive listener before firing the archive event.
+			const ensureArchivedSessionListener = (runInTerminalTool as unknown as Record<string, () => void>)['_ensureArchivedSessionListener'];
+			ensureArchivedSessionListener.call(runInTerminalTool);
+
+			chatSessionArchivedEmitter.fire({
+				resource: sessionResource,
+				isArchived: () => true,
+			} as unknown as IAgentSession);
+
+			strictEqual(terminal1Disposed, true, 'Terminal 1 should have been disposed');
+			strictEqual(terminal2Disposed, true, 'Terminal 2 should have been disposed');
+			ok(!runInTerminalTool.sessionTerminalAssociations.has(sessionResource), 'Terminal association should be removed after archive');
+			ok(!runInTerminalTool.sessionTerminalInstances.has(sessionResource), 'All tracked terminals for the session should be removed after archive');
+		});
+
+		test('should not access agent sessions model when initializing archive listener', () => {
+			let modelAccessed = false;
+			instantiationService.stub(IAgentSessionsService, {
+				onDidChangeSessionArchivedState: chatSessionArchivedEmitter.event,
+				get model() {
+					modelAccessed = true;
+					throw new Error('model should not be accessed when wiring archive listener');
+				},
+			} as unknown as IAgentSessionsService);
+
+			const noModelAccessRunInTerminalTool = store.add(instantiationService.createInstance(TestRunInTerminalTool));
+			const ensureArchivedSessionListener = (noModelAccessRunInTerminalTool as unknown as Record<string, () => void>)['_ensureArchivedSessionListener'];
+			ensureArchivedSessionListener.call(noModelAccessRunInTerminalTool);
+
+			strictEqual(modelAccessed, false, 'Agent sessions model should not be accessed when initializing archive listener');
+		});
+
+		test('should dispose all terminals associated with a single chat session', () => {
+			const sessionId = 'test-session-multiple-terminals';
+			const mockTerminal1 = createMockTerminal(11111);
+			const mockTerminal2 = createMockTerminal(22222);
+
+			let terminal1Disposed = false;
+			let terminal2Disposed = false;
+			mockTerminal1.dispose = () => { terminal1Disposed = true; };
+			mockTerminal2.dispose = () => { terminal2Disposed = true; };
+
+			const sessionResource = LocalChatSessionUri.forSession(sessionId);
+			runInTerminalTool.sessionTerminalAssociations.set(sessionResource, {
+				instance: mockTerminal2,
+				shellIntegrationQuality: ShellIntegrationQuality.None
+			});
+			runInTerminalTool.sessionTerminalInstances.set(sessionResource, new Set([mockTerminal1, mockTerminal2]));
+
+			chatServiceDisposeEmitter.fire({ sessionResource: [sessionResource], reason: 'cleared' });
+
+			strictEqual(terminal1Disposed, true, 'Terminal 1 should have been disposed');
+			strictEqual(terminal2Disposed, true, 'Terminal 2 should have been disposed');
+			ok(!runInTerminalTool.sessionTerminalAssociations.has(sessionResource), 'Terminal association should be removed after disposal');
+			ok(!runInTerminalTool.sessionTerminalInstances.has(sessionResource), 'All tracked terminals for the session should be removed after disposal');
+		});
+
 		test('should dispose associated terminals when chat session is disposed', () => {
 			const sessionId = 'test-session-123';
-			// eslint-disable-next-line local/code-no-any-casts
-			const mockTerminal: ITerminalInstance = {
-				dispose: () => { /* Mock dispose */ },
-				processId: 12345
-			} as any;
+			const mockTerminal = createMockTerminal(12345);
 			let terminalDisposed = false;
 			mockTerminal.dispose = () => { terminalDisposed = true; };
 
@@ -1167,16 +1313,8 @@ suite('RunInTerminalTool', () => {
 		test('should not affect other sessions when one session is disposed', () => {
 			const sessionId1 = 'test-session-1';
 			const sessionId2 = 'test-session-2';
-			// eslint-disable-next-line local/code-no-any-casts
-			const mockTerminal1: ITerminalInstance = {
-				dispose: () => { /* Mock dispose */ },
-				processId: 12345
-			} as any;
-			// eslint-disable-next-line local/code-no-any-casts
-			const mockTerminal2: ITerminalInstance = {
-				dispose: () => { /* Mock dispose */ },
-				processId: 67890
-			} as any;
+			const mockTerminal1 = createMockTerminal(12345);
+			const mockTerminal2 = createMockTerminal(67890);
 
 			let terminal1Disposed = false;
 			let terminal2Disposed = false;