Skip to content

Commit 2181108

Browse files
TyriarTyriar
authored
Merge pull request microsoft#258811 from microsoft/tyriar/258721
Clean up terminals when chat session is disposed
2 parents 4baac87 + d1351ee commit 2181108

File tree

2 files changed

+119
-15
lines changed

2 files changed

+119
-15
lines changed

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

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ const telemetryIgnoredSequences = [
121121
export class RunInTerminalTool extends Disposable implements IToolImpl {
122122

123123
protected readonly _commandLineAutoApprover: CommandLineAutoApprover;
124-
private readonly _sessionTerminalAssociations: Map<string, IToolTerminal> = new Map();
124+
protected readonly _sessionTerminalAssociations: Map<string, IToolTerminal> = new Map();
125125

126126
// Immutable window state
127127
protected readonly _osBackend: Promise<OperatingSystem>;
@@ -166,6 +166,11 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
166166
}
167167
}
168168
}));
169+
170+
// Listen for chat session disposal to clean up associated terminals
171+
this._register(this._chatService.onDidDisposeSession(e => {
172+
this._cleanupSessionTerminals(e.sessionId);
173+
}));
169174
}
170175

171176
async prepareToolInvocation(context: IToolInvocationPreparationContext, token: CancellationToken): Promise<IPreparedToolInvocation | undefined> {
@@ -475,7 +480,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
475480
toolTerminal.instance.dispose();
476481
throw new CancellationError();
477482
}
478-
await this._setupTerminalAssociation(toolTerminal, chatSessionId, termId, true);
483+
await this._setupProcessIdAssociation(toolTerminal, chatSessionId, termId, true);
479484
return toolTerminal;
480485
}
481486

@@ -491,7 +496,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
491496
toolTerminal.instance.dispose();
492497
throw new CancellationError();
493498
}
494-
await this._setupTerminalAssociation(toolTerminal, chatSessionId, termId, false);
499+
await this._setupProcessIdAssociation(toolTerminal, chatSessionId, termId, false);
495500
return toolTerminal;
496501
}
497502

@@ -569,7 +574,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
569574

570575
// Listen for terminal disposal to clean up storage
571576
this._register(instance.onDisposed(() => {
572-
this._removeTerminalAssociation(instance.processId!);
577+
this._removeProcessIdAssociation(instance.processId!);
573578
}));
574579
}
575580
}
@@ -579,16 +584,16 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
579584
}
580585
}
581586

582-
private async _setupTerminalAssociation(toolTerminal: IToolTerminal, chatSessionId: string, termId: string, isBackground: boolean) {
583-
await this._associateTerminalWithSession(toolTerminal.instance, chatSessionId, termId, toolTerminal.shellIntegrationQuality, isBackground);
587+
private async _setupProcessIdAssociation(toolTerminal: IToolTerminal, chatSessionId: string, termId: string, isBackground: boolean) {
588+
await this._associateProcessIdWithSession(toolTerminal.instance, chatSessionId, termId, toolTerminal.shellIntegrationQuality, isBackground);
584589
this._register(toolTerminal.instance.onDisposed(() => {
585590
if (toolTerminal!.instance.processId) {
586-
this._removeTerminalAssociation(toolTerminal!.instance.processId);
591+
this._removeProcessIdAssociation(toolTerminal!.instance.processId);
587592
}
588593
}));
589594
}
590595

591-
private async _associateTerminalWithSession(terminal: ITerminalInstance, sessionId: string, id: string, shellIntegrationQuality: ShellIntegrationQuality, isBackground?: boolean): Promise<void> {
596+
private async _associateProcessIdWithSession(terminal: ITerminalInstance, sessionId: string, id: string, shellIntegrationQuality: ShellIntegrationQuality, isBackground?: boolean): Promise<void> {
592597
try {
593598
// Wait for process ID with timeout
594599
const pid = await Promise.race([
@@ -617,7 +622,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
617622
}
618623
}
619624

620-
private async _removeTerminalAssociation(pid: number): Promise<void> {
625+
private async _removeProcessIdAssociation(pid: number): Promise<void> {
621626
try {
622627
const storedAssociations = this._storageService.get(TERMINAL_SESSION_STORAGE_KEY, StorageScope.WORKSPACE, '{}');
623628
const associations: Record<number, IStoredTerminalAssociation> = JSON.parse(storedAssociations);
@@ -632,6 +637,28 @@ export class RunInTerminalTool extends Disposable implements IToolImpl {
632637
}
633638
}
634639

640+
private _cleanupSessionTerminals(sessionId: string): void {
641+
const toolTerminal = this._sessionTerminalAssociations.get(sessionId);
642+
if (toolTerminal) {
643+
this._logService.debug(`RunInTerminalTool: Cleaning up terminal for disposed chat session ${sessionId}`);
644+
645+
this._sessionTerminalAssociations.delete(sessionId);
646+
toolTerminal.instance.dispose();
647+
648+
// Clean up any background executions associated with this session
649+
const terminalToRemove: string[] = [];
650+
for (const [termId, execution] of RunInTerminalTool._backgroundExecutions.entries()) {
651+
if (execution.instance === toolTerminal.instance) {
652+
execution.dispose();
653+
terminalToRemove.push(termId);
654+
}
655+
}
656+
for (const termId of terminalToRemove) {
657+
RunInTerminalTool._backgroundExecutions.delete(termId);
658+
}
659+
}
660+
}
661+
635662
private _sendTelemetry(instance: ITerminalInstance, state: {
636663
didUserEditCommand: boolean;
637664
didToolEditCommand: boolean;
@@ -711,16 +738,16 @@ class BackgroundTerminalExecution extends Disposable {
711738
private _startMarker?: IXtermMarker;
712739

713740
constructor(
714-
private readonly _instance: ITerminalInstance,
741+
readonly instance: ITerminalInstance,
715742
private readonly _xterm: XtermTerminal,
716743
private readonly _commandLine: string
717744
) {
718745
super();
719746

720747
this._startMarker = this._register(this._xterm.raw.registerMarker());
721-
this._instance.runCommand(this._commandLine, true);
748+
this.instance.runCommand(this._commandLine, true);
722749
}
723750
getOutput(): string {
724-
return getOutput(this._instance, this._startMarker);
751+
return getOutput(this.instance, this._startMarker);
725752
}
726753
}

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

Lines changed: 80 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import { ConfigurationTarget } from '../../../../../../platform/configuration/co
1111
import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js';
1212
import { workbenchInstantiationService } from '../../../../../test/browser/workbenchTestServices.js';
1313
import { IToolInvocationPreparationContext, IPreparedToolInvocation, ILanguageModelToolsService } from '../../../../chat/common/languageModelToolsService.js';
14-
import { CommandLineAutoApprover } from '../../browser/commandLineAutoApprover.js';
1514
import { RunInTerminalTool, type IRunInTerminalInputParams } from '../../browser/runInTerminalTool.js';
1615
import { TerminalChatAgentToolsSettingId } from '../../common/terminalChatAgentToolsConfiguration.js';
1716
import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js';
@@ -20,11 +19,14 @@ import type { TestInstantiationService } from '../../../../../../platform/instan
2019
import { ITerminalService, type ITerminalInstance } from '../../../../terminal/browser/terminal.js';
2120
import { OperatingSystem } from '../../../../../../base/common/platform.js';
2221
import { Emitter } from '../../../../../../base/common/event.js';
22+
import { IChatService } from '../../../../chat/common/chatService.js';
23+
import { ShellIntegrationQuality } from '../../browser/toolTerminalCreator.js';
2324

2425
class TestRunInTerminalTool extends RunInTerminalTool {
2526
protected override _osBackend: Promise<OperatingSystem> = Promise.resolve(OperatingSystem.Windows);
2627

27-
get commandLineAutoApprover(): CommandLineAutoApprover { return this._commandLineAutoApprover; }
28+
get commandLineAutoApprover() { return this._commandLineAutoApprover; }
29+
get sessionTerminalAssociations() { return this._sessionTerminalAssociations; }
2830

2931
async rewriteCommandIfNeeded(args: IRunInTerminalInputParams, instance: Pick<ITerminalInstance, 'getCwdResource'> | undefined, shell: string): Promise<string> {
3032
return this._rewriteCommandIfNeeded(args, instance, shell);
@@ -41,11 +43,16 @@ suite('RunInTerminalTool', () => {
4143
let instantiationService: TestInstantiationService;
4244
let configurationService: TestConfigurationService;
4345
let workspaceService: TestContextService;
46+
let terminalServiceDisposeEmitter: Emitter<ITerminalInstance>;
47+
let chatServiceDisposeEmitter: Emitter<{ sessionId: string; reason: 'cleared' }>;
4448

4549
let runInTerminalTool: TestRunInTerminalTool;
4650

4751
setup(() => {
4852
configurationService = new TestConfigurationService();
53+
terminalServiceDisposeEmitter = new Emitter<ITerminalInstance>();
54+
chatServiceDisposeEmitter = new Emitter<{ sessionId: string; reason: 'cleared' }>();
55+
4956
instantiationService = workbenchInstantiationService({
5057
configurationService: () => configurationService,
5158
}, store);
@@ -55,7 +62,10 @@ suite('RunInTerminalTool', () => {
5562
},
5663
});
5764
instantiationService.stub(ITerminalService, {
58-
onDidDisposeInstance: new Emitter<ITerminalInstance>().event
65+
onDidDisposeInstance: terminalServiceDisposeEmitter.event
66+
});
67+
instantiationService.stub(IChatService, {
68+
onDidDisposeSession: chatServiceDisposeEmitter.event
5969
});
6070
workspaceService = instantiationService.invokeFunction(accessor => accessor.get(IWorkspaceContextService)) as TestContextService;
6171

@@ -611,4 +621,71 @@ suite('RunInTerminalTool', () => {
611621
});
612622
});
613623
});
624+
625+
suite('chat session disposal cleanup', () => {
626+
test('should dispose associated terminals when chat session is disposed', () => {
627+
const sessionId = 'test-session-123';
628+
const mockTerminal: ITerminalInstance = {
629+
dispose: () => { /* Mock dispose */ },
630+
processId: 12345
631+
} as any;
632+
let terminalDisposed = false;
633+
mockTerminal.dispose = () => { terminalDisposed = true; };
634+
635+
runInTerminalTool.sessionTerminalAssociations.set(sessionId, {
636+
instance: mockTerminal,
637+
shellIntegrationQuality: ShellIntegrationQuality.None
638+
});
639+
640+
ok(runInTerminalTool.sessionTerminalAssociations.has(sessionId), 'Terminal association should exist before disposal');
641+
642+
chatServiceDisposeEmitter.fire({ sessionId, reason: 'cleared' });
643+
644+
strictEqual(terminalDisposed, true, 'Terminal should have been disposed');
645+
ok(!runInTerminalTool.sessionTerminalAssociations.has(sessionId), 'Terminal association should be removed after disposal');
646+
});
647+
648+
test('should not affect other sessions when one session is disposed', () => {
649+
const sessionId1 = 'test-session-1';
650+
const sessionId2 = 'test-session-2';
651+
const mockTerminal1: ITerminalInstance = {
652+
dispose: () => { /* Mock dispose */ },
653+
processId: 12345
654+
} as any;
655+
const mockTerminal2: ITerminalInstance = {
656+
dispose: () => { /* Mock dispose */ },
657+
processId: 67890
658+
} as any;
659+
660+
let terminal1Disposed = false;
661+
let terminal2Disposed = false;
662+
mockTerminal1.dispose = () => { terminal1Disposed = true; };
663+
mockTerminal2.dispose = () => { terminal2Disposed = true; };
664+
665+
runInTerminalTool.sessionTerminalAssociations.set(sessionId1, {
666+
instance: mockTerminal1,
667+
shellIntegrationQuality: ShellIntegrationQuality.None
668+
});
669+
runInTerminalTool.sessionTerminalAssociations.set(sessionId2, {
670+
instance: mockTerminal2,
671+
shellIntegrationQuality: ShellIntegrationQuality.None
672+
});
673+
674+
ok(runInTerminalTool.sessionTerminalAssociations.has(sessionId1), 'Session 1 terminal association should exist');
675+
ok(runInTerminalTool.sessionTerminalAssociations.has(sessionId2), 'Session 2 terminal association should exist');
676+
677+
chatServiceDisposeEmitter.fire({ sessionId: sessionId1, reason: 'cleared' });
678+
679+
strictEqual(terminal1Disposed, true, 'Terminal 1 should have been disposed');
680+
strictEqual(terminal2Disposed, false, 'Terminal 2 should NOT have been disposed');
681+
ok(!runInTerminalTool.sessionTerminalAssociations.has(sessionId1), 'Session 1 terminal association should be removed');
682+
ok(runInTerminalTool.sessionTerminalAssociations.has(sessionId2), 'Session 2 terminal association should remain');
683+
});
684+
685+
test('should handle disposal of non-existent session gracefully', () => {
686+
strictEqual(runInTerminalTool.sessionTerminalAssociations.size, 0, 'No associations should exist initially');
687+
chatServiceDisposeEmitter.fire({ sessionId: 'non-existent-session', reason: 'cleared' });
688+
strictEqual(runInTerminalTool.sessionTerminalAssociations.size, 0, 'No associations should exist after handling non-existent session');
689+
});
690+
});
614691
});

0 commit comments

Comments
 (0)