Fix editing session lifecycle issues (microsoft#229807)

* Redo the way an edit session gets disposed and add more assertions to catch any incorrect usages

* Fix lifecycle issues, also don't capture in the click closure the initial contents of the editing session

Author: alexdima

src/vs/workbench/contrib/chat/browser/chatEditingService.ts

@@ -7,8 +7,7 @@ import { Sequencer } from '../../../../base/common/async.js';
 import { Codicon } from '../../../../base/common/codicons.js';
 import { BugIndicatingError } from '../../../../base/common/errors.js';
 import { Emitter } from '../../../../base/common/event.js';
-import { Disposable, IReference } from '../../../../base/common/lifecycle.js';
-import { ResourceSet } from '../../../../base/common/map.js';
+import { Disposable, IReference, MutableDisposable } from '../../../../base/common/lifecycle.js';
 import { Schemas } from '../../../../base/common/network.js';
 import { derived, IObservable, ITransaction, observableValue, ValueWithChangeEventFromObservable } from '../../../../base/common/observable.js';
 import { URI } from '../../../../base/common/uri.js';
@@ -52,6 +51,7 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 	_serviceBrand: undefined;
 
 	private readonly _currentSessionObs = observableValue<ChatEditingSession | null>(this, null);
+	private readonly _currentSessionDisposeListener = this._register(new MutableDisposable());
 
 	get currentEditingSession(): IChatEditingSession | null {
 		return this._currentSessionObs.get();
@@ -62,11 +62,6 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		return this._onDidCreateEditingSession.event;
 	}
 
-	private readonly _onDidDisposeEditingSession = new Emitter<IChatEditingSession>();
-	get onDidDisposeEditingSession() {
-		return this._onDidDisposeEditingSession.event;
-	}
-
 	constructor(
 		@IEditorGroupsService private readonly _editorGroupsService: IEditorGroupsService,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
@@ -90,11 +85,16 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		}));
 		this._register(this._chatService.onDidDisposeSession((e) => {
 			if (e.reason === 'cleared' && this._currentSessionObs.get()?.chatSessionId === e.sessionId) {
-				this._killCurrentEditingSession();
+				void this._currentSessionObs.get()?.stop();
 			}
 		}));
 	}
 
+	override dispose(): void {
+		this._currentSessionObs.get()?.dispose();
+		super.dispose();
+	}
+
 	async startOrContinueEditingSession(chatSessionId: string, builder?: (stream: IChatEditingSessionStream) => Promise<void>, options?: { silent: boolean }): Promise<void> {
 		const session = this._currentSessionObs.get();
 		if (session) {
@@ -121,12 +121,10 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		const editorPane = options?.silent ? undefined : await this._editorGroupsService.activeGroup.openEditor(input, { pinned: true, activation: EditorActivation.ACTIVATE }) as MultiDiffEditor | undefined;
 
 		const session = this._instantiationService.createInstance(ChatEditingSession, chatSessionId, editorPane);
-		const disposeListener = session.onDidDispose(() => {
+		this._currentSessionDisposeListener.value = session.onDidDispose(() => {
+			this._currentSessionDisposeListener.clear();
 			this._currentSessionObs.set(null, undefined);
-			this._onDidDisposeEditingSession.fire(session);
-			disposeListener.dispose();
 		});
-		this._register(disposeListener);
 
 		this._currentSessionObs.set(session, undefined);
 		this._onDidCreateEditingSession.fire(session);
@@ -179,15 +177,6 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		}
 	}
 
-	private _killCurrentEditingSession() {
-		const currentSession = this._currentSessionObs.get();
-		if (currentSession) {
-			this._onDidDisposeEditingSession.fire(currentSession);
-			currentSession.dispose();
-			this._currentSessionObs.set(null, undefined);
-		}
-	}
-
 	private _findGroupedEditors() {
 		const editors: [IEditorGroup, EditorInput][] = [];
 		for (const group of this._editorGroupsService.groups) {
@@ -495,7 +484,7 @@ export class ChatEditingStopSessionAction extends Action2 {
 
 	async run(accessor: ServicesAccessor, ...args: any[]): Promise<void> {
 		const chatEditingService = accessor.get(IChatEditingService);
-		chatEditingService.currentEditingSession?.dispose();
+		await chatEditingService.currentEditingSession?.stop();
 	}
 }
 registerAction2(ChatEditingStopSessionAction);
@@ -572,28 +561,32 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 	private readonly _state = observableValue<ChatEditingSessionState>(this, ChatEditingSessionState.Initial);
 	private readonly _entriesObs = observableValue<readonly ModifiedFileEntry[]>(this, []);
 	public get entries(): IObservable<readonly ModifiedFileEntry[]> {
+		this._assertNotDisposed();
 		return this._entriesObs;
 	}
 	private readonly _sequencer = new Sequencer();
 
 	private _entries: ModifiedFileEntry[] = [];
 
 	get state(): IObservable<ChatEditingSessionState> {
+		this._assertNotDisposed();
 		return this._state;
 	}
 
-	private readonly _editedResources = new ResourceSet();
 	private readonly _onDidChange = new Emitter<void>();
 	get onDidChange() {
+		this._assertNotDisposed();
 		return this._onDidChange.event;
 	}
 
 	private readonly _onDidDispose = new Emitter<void>();
 	get onDidDispose() {
+		this._assertNotDisposed();
 		return this._onDidDispose.event;
 	}
 
 	get isVisible(): boolean {
+		this._assertNotDisposed();
 		return Boolean(this.editorPane && this.editorPane.isVisible());
 	}
 
@@ -608,10 +601,14 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 		super();
 	}
 
-	async accept(...uris: URI[]): Promise<void> {
-		if (this.state.get() === ChatEditingSessionState.Disposed) {
-			return;
+	private _assertNotDisposed(): void {
+		if (this._state.get() === ChatEditingSessionState.Disposed) {
+			throw new BugIndicatingError(`Cannot access a disposed editing session`);
 		}
+	}
+
+	async accept(...uris: URI[]): Promise<void> {
+		this._assertNotDisposed();
 
 		if (uris.length === 0) {
 			await Promise.all(this._entries.map(entry => entry.accept(undefined)));
@@ -628,9 +625,7 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 	}
 
 	async reject(...uris: URI[]): Promise<void> {
-		if (this.state.get() === ChatEditingSessionState.Disposed) {
-			return;
-		}
+		this._assertNotDisposed();
 
 		if (uris.length === 0) {
 			await Promise.all(this._entries.map(entry => entry.reject(undefined)));
@@ -647,6 +642,8 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 	}
 
 	async show(): Promise<void> {
+		this._assertNotDisposed();
+
 		if (this.editorPane?.isVisible()) {
 			return;
 		} else if (this.editorPane?.input) {
@@ -663,28 +660,39 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 		this.editorPane = editorPane;
 	}
 
-	override dispose() {
+	async stop(): Promise<void> {
+		this._assertNotDisposed();
+
 		// Close out all open files
-		this._editorGroupsService.groups.forEach((g) => {
-			g.editors.forEach((e) => {
+		await Promise.allSettled(this._editorGroupsService.groups.map(async (g) => {
+			return Promise.allSettled(g.editors.map(async (e) => {
 				if (e instanceof MultiDiffEditorInput || e instanceof DiffEditorInput && (e.original.resource?.scheme === ModifiedFileEntry.scheme || e.original.resource?.scheme === ChatEditingTextModelContentProvider.scheme)) {
-					g.closeEditor(e);
+					await g.closeEditor(e);
 				}
-			});
-		});
+			}));
+		}));
+
+		this.dispose();
+	}
+
+	override dispose() {
+		this._assertNotDisposed();
 
 		super.dispose();
 		this._state.set(ChatEditingSessionState.Disposed, undefined);
 		this._onDidDispose.fire();
 	}
 
 	getVirtualModel(documentId: string): ITextModel | null {
+		this._assertNotDisposed();
+
 		const entry = this._entries.find(e => e.entryId === documentId);
 		return entry?.docSnapshot ?? null;
 	}
 
 	acceptStreamingEditsStart(): void {
-		if (this.state.get() === ChatEditingSessionState.Disposed) {
+		if (this._state.get() === ChatEditingSessionState.Disposed) {
+			// we don't throw in this case because there could be a builder still connected to a disposed session
 			return;
 		}
 
@@ -693,7 +701,8 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 	}
 
 	acceptTextEdits(resource: URI, textEdits: TextEdit[]): void {
-		if (this.state.get() === ChatEditingSessionState.Disposed) {
+		if (this._state.get() === ChatEditingSessionState.Disposed) {
+			// we don't throw in this case because there could be a builder still connected to a disposed session
 			return;
 		}
 
@@ -702,7 +711,8 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 	}
 
 	resolve(): void {
-		if (this.state.get() === ChatEditingSessionState.Disposed) {
+		if (this._state.get() === ChatEditingSessionState.Disposed) {
+			// we don't throw in this case because there could be a builder still connected to a disposed session
 			return;
 		}
 
@@ -734,7 +744,6 @@ class ChatEditingSession extends Disposable implements IChatEditingSession {
 		this._register(entry);
 		this._entries = [...this._entries, entry];
 		this._entriesObs.set(this._entries, undefined);
-		this._editedResources.add(resource);
 		this._onDidChange.fire();
 
 		return entry;

src/vs/workbench/contrib/chat/browser/chatInputPart.ts

@@ -851,13 +851,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 
 		if (this._chatEditList && chatEditingSession.state.get() === ChatEditingSessionState.Idle) {
 			this._chatEditsProgress?.stop();
-			return;
+			this._chatEditsProgress?.dispose();
+			this._chatEditsProgress = undefined;
 		}
 
 		// Summary of number of files changed
 		const innerContainer = this.chatEditingSessionWidgetContainer.querySelector('.chat-editing-session-container.show-file-icons') as HTMLElement ?? dom.append(this.chatEditingSessionWidgetContainer, $('.chat-editing-session-container.show-file-icons'));
-		const editedFiles = chatEditingSession.entries.get();
-		const numberOfEditedEntries = editedFiles.length;
+		const numberOfEditedEntries = chatEditingSession.entries.get().length;
 		const overviewRegion = innerContainer.querySelector('.chat-editing-session-overview') as HTMLElement ?? dom.append(innerContainer, $('.chat-editing-session-overview'));
 		if (numberOfEditedEntries !== this._chatEditList?.object.length) {
 			const overviewText = overviewRegion.querySelector('span') ?? dom.append(overviewRegion, $('span'));
@@ -869,7 +869,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		// Chat editing session actions
 		let actionsContainer = overviewRegion.querySelector('.chat-editing-session-actions') as HTMLElement;
 		const actions = [];
-		if (!actionsContainer && editedFiles.find((e) => e.state.get() === ModifiedFileEntryState.Undecided)) {
+		if (!actionsContainer && chatEditingSession.entries.get().find((e) => e.state.get() === ModifiedFileEntryState.Undecided)) {
 			// Don't show Accept All / Discard All actions if user already selected Accept All / Discard All
 			actionsContainer = dom.append(overviewRegion, $('.chat-editing-session-actions'));
 			actions.push(
@@ -891,24 +891,22 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 			);
 
 			for (const action of actions) {
-				const button = this._register(new Button(actionsContainer, {
+				const button = this._chatEditsDisposables.add(new Button(actionsContainer, {
 					supportIcons: false,
 					secondary: action.isSecondary
 				}));
 				button.label = action.label;
-				this._register(button.onDidClick(() => {
+				this._chatEditsDisposables.add(button.onDidClick(() => {
 					this.commandService.executeCommand(action.command);
 				}));
 				dom.append(actionsContainer, button.element);
 			}
 
-			const clearButton = new Button(actionsContainer, { supportIcons: true });
-			this._chatEditsDisposables.add(clearButton);
+			const clearButton = this._chatEditsDisposables.add(new Button(actionsContainer, { supportIcons: true }));
 			clearButton.icon = Codicon.close;
-			const disp = clearButton.onDidClick((e) => {
-				disp.dispose();
-				chatEditingSession.dispose();
-			});
+			this._chatEditsDisposables.add(clearButton.onDidClick((e) => {
+				void chatEditingSession.stop();
+			}));
 			dom.append(actionsContainer, clearButton.element);
 		}
 
@@ -918,21 +916,14 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		}
 
 		// List of edited files
-		if (!editedFiles.length || editedFiles.length === this._chatEditList?.object.length) {
-			return;
-		}
-		const entries: IChatCollapsibleListItem[] = editedFiles.map((entry) => ({
-			reference: entry.modifiedURI,
-			kind: 'reference',
-		}));
 		if (!this._chatEditList) {
 			this._chatEditList = this._chatEditsListPool.get();
 			const list = this._chatEditList.object;
 			this._chatEditsDisposables.add(this._chatEditList);
 			this._chatEditsDisposables.add(list.onDidOpen((e) => {
 				if (e.element?.kind === 'reference' && URI.isUri(e.element.reference)) {
 					const modifiedFileUri = e.element.reference;
-					const editedFile = editedFiles.find((e) => e.modifiedURI.toString() === modifiedFileUri.toString());
+					const editedFile = chatEditingSession.entries.get().find((e) => e.modifiedURI.toString() === modifiedFileUri.toString());
 					if (editedFile) {
 						void this.editorService.openEditor({
 							original: { resource: URI.from(editedFile.originalURI, true) },
@@ -950,6 +941,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const list = this._chatEditList.object;
 		list.layout(height);
 		list.getHTMLElement().style.height = `${height}px`;
+		const entries: IChatCollapsibleListItem[] = chatEditingSession.entries.get().map((entry) => ({
+			reference: entry.modifiedURI,
+			kind: 'reference',
+		}));
 		list.splice(0, list.length, entries);
 	}
 

src/vs/workbench/contrib/chat/browser/chatWidget.ts

@@ -250,20 +250,25 @@ export class ChatWidget extends Disposable implements IChatWidget {
 
 		this._codeBlockModelCollection = this._register(instantiationService.createInstance(CodeBlockModelCollection));
 
-		let shouldRerenderEditingStateDisposable: IDisposable | null = null;
+		const chatEditingSessionDisposables = this._register(new DisposableStore());
+
 		this._register(this.chatEditingService.onDidCreateEditingSession((session) => {
-			if (session.chatSessionId === this.viewModel?.sessionId) {
-				shouldRerenderEditingStateDisposable = this._register(session.onDidChange(() => {
-					this.renderChatEditingSessionState(session);
-				}));
-				this.renderChatEditingSessionState(session);
+			if (session.chatSessionId !== this.viewModel?.sessionId) {
+				// this chat editing session is for a different chat widget
+				return;
 			}
-		}));
-		this._register(this.chatEditingService.onDidDisposeEditingSession((session) => {
-			if (session.chatSessionId === this.viewModel?.sessionId || !this.viewModel) {
-				shouldRerenderEditingStateDisposable?.dispose();
+			// make sure to clean up anything related to the prev session (if any)
+			chatEditingSessionDisposables.clear();
+			this.renderChatEditingSessionState(null); // this is necessary to make sure we dispose previous buttons, etc.
+
+			chatEditingSessionDisposables.add(session.onDidChange(() => {
+				this.renderChatEditingSessionState(session);
+			}));
+			chatEditingSessionDisposables.add(session.onDidDispose(() => {
+				chatEditingSessionDisposables.clear();
 				this.renderChatEditingSessionState(null);
-			}
+			}));
+			this.renderChatEditingSessionState(session);
 		}));
 
 		this._register(codeEditorService.registerCodeEditorOpenHandler(async (input: ITextResourceEditorInput, _source: ICodeEditor | null, _sideBySide?: boolean): Promise<ICodeEditor | null> => {

src/vs/workbench/contrib/chat/common/chatEditingService.ts

@@ -15,7 +15,6 @@ export interface IChatEditingService {
 	_serviceBrand: undefined;
 
 	readonly onDidCreateEditingSession: Event<IChatEditingSession>;
-	readonly onDidDisposeEditingSession: Event<IChatEditingSession>;
 
 	readonly currentEditingSession: IChatEditingSession | null;
 
@@ -25,13 +24,17 @@ export interface IChatEditingService {
 export interface IChatEditingSession {
 	readonly chatSessionId: string;
 	readonly onDidChange: Event<void>;
+	readonly onDidDispose: Event<void>;
 	readonly state: IObservable<ChatEditingSessionState>;
 	readonly entries: IObservable<readonly IModifiedFileEntry[]>;
 	readonly isVisible: boolean;
 	show(): Promise<void>;
 	accept(...uris: URI[]): Promise<void>;
 	reject(...uris: URI[]): Promise<void>;
-	dispose(): void;
+	/**
+	 * Will lead to this object getting disposed
+	 */
+	stop(): Promise<void>;
 }
 
 export const enum ModifiedFileEntryState {

src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts

@@ -160,7 +160,6 @@ suite('InteractiveChatController', function () {
 			[ICommandService, new SyncDescriptor(TestCommandService)],
 			[IChatEditingService, new class extends mock<IChatEditingService>() {
 				override onDidCreateEditingSession: Event<IChatEditingSession> = Event.None;
-				override onDidDisposeEditingSession: Event<IChatEditingSession> = Event.None;
 			}],
 			[IInlineChatSavingService, new class extends mock<IInlineChatSavingService>() {
 				override markChanged(session: Session): void {