diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentSessionsService.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentSessionsService.ts index efdfc49fa93..3bce41c7d02 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentSessionsService.ts +++ b/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; 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()); + 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 */); } diff --git a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts index fd8fcd1f37a..2947937e8e5 100644 --- a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts +++ b/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(); + protected readonly _sessionTerminalInstances = new ResourceMap>(); + private readonly _terminalsBeingDisposedBySessionCleanup = new Set(); // Immutable window state protected readonly _osBackend: Promise; @@ -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 { @@ -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(); + 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); } } } diff --git a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/runInTerminalTool.test.ts b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/runInTerminalTool.test.ts index 48ff9c22554..43c58ddd915 100644 --- a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/test/electron-browser/runInTerminalTool.test.ts +++ b/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 = 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; let chatServiceDisposeEmitter: Emitter<{ sessionResource: URI[]; reason: 'cleared' }>; + let chatSessionArchivedEmitter: Emitter; let runInTerminalTool: TestRunInTerminalTool; @@ -79,6 +83,7 @@ suite('RunInTerminalTool', () => { setConfig(TerminalChatAgentToolsSettingId.BlockDetectedFileWrites, 'outsideWorkspace'); terminalServiceDisposeEmitter = new Emitter(); chatServiceDisposeEmitter = new Emitter<{ sessionResource: URI[]; reason: 'cleared' }>(); + chatSessionArchivedEmitter = new Emitter(); 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(); + const terminal2DisposedEmitter = new Emitter(); + 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 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 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;