From 30cfcc9385eb6862964a079f6ec07a87fde56b2e Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Sun, 23 Nov 2025 21:21:12 -0800 Subject: [PATCH] Fix ChatService.onDidDispose (#278917) * Fix ChatService.onDidDispose * Fix? * Wait for disposal * Mop up remaining disposable leaks * Dispose IChatService * And this * Simplify * Dispose Emitter --- .../contrib/chat/browser/chatWidget.ts | 11 ++++++----- .../contrib/chat/common/chatModelStore.ts | 10 +++++++++- .../contrib/chat/common/chatService.ts | 1 - .../contrib/chat/common/chatServiceImpl.ts | 18 +++--------------- .../test/browser/chatEditingService.test.ts | 4 +--- .../chat/test/common/chatService.test.ts | 19 ++++++++++++++++++- .../chat/test/common/mockChatService.ts | 3 --- .../browser/inlineChatSessionServiceImpl.ts | 7 ------- .../test/browser/inlineChatController.test.ts | 3 ++- .../test/browser/inlineChatSession.test.ts | 4 +++- 10 files changed, 42 insertions(+), 38 deletions(-) diff --git a/src/vs/workbench/contrib/chat/browser/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/chatWidget.ts index 581570b9c3f..e1c1bc4124e 100644 --- a/src/vs/workbench/contrib/chat/browser/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/chatWidget.ts @@ -1632,9 +1632,9 @@ export class ChatWidget extends Disposable implements IChatWidget { } }, 0); - dom.scheduleAtNextAnimationFrame(dom.getWindow(this.listContainer), () => { + this._register(dom.scheduleAtNextAnimationFrame(dom.getWindow(this.listContainer), () => { this._onDidShow.fire(); - }); + })); } } else if (wasVisible) { this._onDidHide.fire(); @@ -1991,11 +1991,11 @@ export class ChatWidget extends Disposable implements IChatWidget { // Consider the tree to be scrolled all the way down if it is within 2px of the bottom. const lastElementWasVisible = this.tree.scrollTop + this.tree.renderHeight >= this.previousTreeScrollHeight - 2; if (lastElementWasVisible) { - dom.scheduleAtNextAnimationFrame(dom.getWindow(this.listContainer), () => { + this._register(dom.scheduleAtNextAnimationFrame(dom.getWindow(this.listContainer), () => { // Can't set scrollTop during this event listener, the list might overwrite the change this.scrollToEnd(); - }, 0); + }, 0)); } } } @@ -2205,7 +2205,8 @@ export class ChatWidget extends Disposable implements IChatWidget { const renderImmediately = this.configurationService.getValue('chat.experimental.renderMarkdownImmediately'); const delay = renderImmediately ? MicrotaskDelay : 0; this.viewModelDisposables.add(Event.runAndSubscribe(Event.accumulate(this.viewModel.onDidChange, delay), (events => { - if (!this.viewModel) { + if (!this.viewModel || this._store.isDisposed) { + // See https://github.com/microsoft/vscode/issues/278969 return; } diff --git a/src/vs/workbench/contrib/chat/common/chatModelStore.ts b/src/vs/workbench/contrib/chat/common/chatModelStore.ts index e83dee43860..fe6a9b169f3 100644 --- a/src/vs/workbench/contrib/chat/common/chatModelStore.ts +++ b/src/vs/workbench/contrib/chat/common/chatModelStore.ts @@ -4,7 +4,8 @@ *--------------------------------------------------------------------------------------------*/ import { CancellationToken } from '../../../../base/common/cancellation.js'; -import { IDisposable, IReference, ReferenceCollection } from '../../../../base/common/lifecycle.js'; +import { Emitter } from '../../../../base/common/event.js'; +import { DisposableStore, IDisposable, IReference, ReferenceCollection } from '../../../../base/common/lifecycle.js'; import { ObservableMap } from '../../../../base/common/observable.js'; import { URI } from '../../../../base/common/uri.js'; import { ILogService } from '../../../../platform/log/common/log.js'; @@ -28,10 +29,15 @@ export interface ChatModelStoreDelegate { } export class ChatModelStore extends ReferenceCollection implements IDisposable { + private readonly _store = new DisposableStore(); + private readonly _models = new ObservableMap(); private readonly _modelsToDispose = new Set(); private readonly _pendingDisposals = new Set>(); + private readonly _onDidDisposeModel = this._store.add(new Emitter()); + public readonly onDidDisposeModel = this._onDidDisposeModel.event; + constructor( private readonly delegate: ChatModelStoreDelegate, @ILogService private readonly logService: ILogService, @@ -108,6 +114,7 @@ export class ChatModelStore extends ReferenceCollection implements ID if (this._modelsToDispose.has(key)) { this.logService.trace(`Disposing chat session ${key}`); this._models.delete(key); + this._onDidDisposeModel.fire(object); object.dispose(); } this._modelsToDispose.delete(key); @@ -126,6 +133,7 @@ export class ChatModelStore extends ReferenceCollection implements ID } dispose(): void { + this._store.dispose(); this._models.forEach(model => model.dispose()); } } diff --git a/src/vs/workbench/contrib/chat/common/chatService.ts b/src/vs/workbench/contrib/chat/common/chatService.ts index 3a870cc688d..790dc303214 100644 --- a/src/vs/workbench/contrib/chat/common/chatService.ts +++ b/src/vs/workbench/contrib/chat/common/chatService.ts @@ -979,7 +979,6 @@ export interface IChatService { adoptRequest(sessionResource: URI, request: IChatRequestModel): Promise; removeRequest(sessionResource: URI, requestId: string): Promise; cancelCurrentRequestForSession(sessionResource: URI): void; - forceClearSession(sessionResource: URI): Promise; addCompleteRequest(sessionResource: URI, message: IParsedChatRequest | string, variableData: IChatRequestVariableData | undefined, attempt: number | undefined, response: IChatCompleteResponse): void; setChatSessionTitle(sessionResource: URI, title: string): void; getLocalSessionHistory(): Promise; diff --git a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts index 88563efb462..f30dad6467f 100644 --- a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts @@ -175,6 +175,9 @@ export class ChatService extends Disposable implements IChatService { } } })); + this._register(this._sessionModels.onDidDisposeModel(model => { + this._onDidDisposeSession.fire({ sessionResource: model.sessionResource, reason: 'cleared' }); + })); this._chatServiceTelemetry = this.instantiationService.createInstance(ChatServiceTelemetry); @@ -1252,21 +1255,6 @@ export class ChatService extends Disposable implements IChatService { this._pendingRequests.deleteAndDispose(sessionResource); } - // TODO should not exist - async forceClearSession(sessionResource: URI): Promise { - this.trace('clearSession', `session: ${sessionResource}`); - const model = this._sessionModels.get(sessionResource); - if (!model) { - throw new Error(`Unknown session: ${sessionResource}`); - } - - // this._sessionModels.delete(sessionResource); - model.dispose(); - this._pendingRequests.get(sessionResource)?.cancel(); - this._pendingRequests.deleteAndDispose(sessionResource); - this._onDidDisposeSession.fire({ sessionResource, reason: 'cleared' }); - } - public hasSessions(): boolean { return this._chatSessionStore.hasSessions(); } diff --git a/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts b/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts index fbd19e46702..f76c1a9db03 100644 --- a/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts @@ -161,7 +161,7 @@ suite('ChatEditingService', function () { const uri = URI.from({ scheme: 'test', path: 'HelloWorld' }); - const modelRef = chatService.startSession(ChatAgentLocation.Chat, CancellationToken.None); + const modelRef = store.add(chatService.startSession(ChatAgentLocation.Chat, CancellationToken.None)); const model = modelRef.object as ChatModel; const session = model.editingSession; if (!session) { @@ -188,8 +188,6 @@ suite('ChatEditingService', function () { await unset; await entry.reject(); - - modelRef.dispose(); }); async function idleAfterEdit(session: IChatEditingSession, model: ChatModel, uri: URI, edits: TextEdit[]) { diff --git a/src/vs/workbench/contrib/chat/test/common/chatService.test.ts b/src/vs/workbench/contrib/chat/test/common/chatService.test.ts index dcd14ae96bc..acf46701097 100644 --- a/src/vs/workbench/contrib/chat/test/common/chatService.test.ts +++ b/src/vs/workbench/contrib/chat/test/common/chatService.test.ts @@ -7,6 +7,7 @@ import assert from 'assert'; import { CancellationToken } from '../../../../../base/common/cancellation.js'; import { Event } from '../../../../../base/common/event.js'; import { MarkdownString } from '../../../../../base/common/htmlContent.js'; +import { DisposableStore } from '../../../../../base/common/lifecycle.js'; import { observableValue } from '../../../../../base/common/observable.js'; import { URI } from '../../../../../base/common/uri.js'; import { assertSnapshot } from '../../../../../base/test/common/snapshot.js'; @@ -43,7 +44,6 @@ import { IChatVariablesService } from '../../common/chatVariables.js'; import { ChatAgentLocation, ChatModeKind } from '../../common/constants.js'; import { MockChatService } from './mockChatService.js'; import { MockChatVariablesService } from './mockChatVariables.js'; -import { DisposableStore } from '../../../../../base/common/lifecycle.js'; const chatAgentWithUsedContextId = 'ChatProviderWithUsedContext'; const chatAgentWithUsedContext: IChatAgent = { @@ -397,6 +397,23 @@ suite('ChatService', () => { await assertSnapshot(toSnapshotExportData(chatModel2)); }); + + test('onDidDisposeSession', async () => { + const testService = createChatService(); + const modelRef = testService.startSession(ChatAgentLocation.Chat, CancellationToken.None); + const model = modelRef.object; + + let disposed = false; + testDisposables.add(testService.onDidDisposeSession(e => { + if (e.sessionResource.toString() === model.sessionResource.toString()) { + disposed = true; + } + })); + + modelRef.dispose(); + await testService.waitForModelDisposals(); + assert.strictEqual(disposed, true); + }); }); diff --git a/src/vs/workbench/contrib/chat/test/common/mockChatService.ts b/src/vs/workbench/contrib/chat/test/common/mockChatService.ts index 0d26c7eab66..15086524079 100644 --- a/src/vs/workbench/contrib/chat/test/common/mockChatService.ts +++ b/src/vs/workbench/contrib/chat/test/common/mockChatService.ts @@ -81,9 +81,6 @@ export class MockChatService implements IChatService { cancelCurrentRequestForSession(sessionResource: URI): void { throw new Error('Method not implemented.'); } - forceClearSession(sessionResource: URI): Promise { - throw new Error('Method not implemented.'); - } addCompleteRequest(sessionResource: URI, message: IParsedChatRequest | string, variableData: IChatRequestVariableData | undefined, attempt: number | undefined, response: IChatCompleteResponse): void { throw new Error('Method not implemented.'); } diff --git a/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts b/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts index 41bae92b154..832edc0aada 100644 --- a/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts +++ b/src/vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl.ts @@ -134,13 +134,6 @@ export class InlineChatSessionServiceImpl implements IInlineChatSessionService { store.add(chatModelRef); } - store.add(toDisposable(() => { - const doesOtherSessionUseChatModel = [...this._sessions.values()].some(data => data.session !== session && data.session.chatModel === chatModel); - - if (!doesOtherSessionUseChatModel) { - this._chatService.forceClearSession(chatModel.sessionResource); - } - })); const lastResponseListener = store.add(new MutableDisposable()); store.add(chatModel.onDidChange(e => { if (e.kind !== 'addRequest' || !e.request.response) { diff --git a/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts b/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts index b75de3fbf19..13522f2c673 100644 --- a/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts +++ b/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts @@ -275,9 +275,10 @@ suite('InlineChatController', function () { }); - teardown(function () { + teardown(async function () { store.clear(); ctrl?.dispose(); + await chatService.waitForModelDisposals(); }); // TODO@jrieken re-enable, looks like List/ChatWidget is leaking diff --git a/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatSession.test.ts b/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatSession.test.ts index d4c489ddeb8..425b3525bfa 100644 --- a/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatSession.test.ts +++ b/src/vs/workbench/contrib/inlineChat/test/browser/inlineChatSession.test.ts @@ -144,6 +144,7 @@ suite('InlineChatSession', function () { instaService = store.add(workbenchInstantiationService(undefined, store).createChild(serviceCollection)); inlineChatSessionService = store.add(instaService.get(IInlineChatSessionService)); store.add(instaService.get(IChatSessionsService) as ChatSessionsService); // Needs to be disposed in between test runs to clear extensionPoint contribution + store.add(instaService.get(IChatService) as ChatService); instaService.get(IChatAgentService).registerDynamicAgent({ extensionId: nullExtensionDescription.identifier, @@ -171,8 +172,9 @@ suite('InlineChatSession', function () { editor = store.add(instantiateTestCodeEditor(instaService, model)); }); - teardown(function () { + teardown(async function () { store.clear(); + await instaService.get(IChatService).waitForModelDisposals(); }); ensureNoDisposablesAreLeakedInTestSuite();