Don't collect codeblocks after chat part disposal (microsoft#242392)

Fix microsoft#241859

Author: roblourens

src/vs/base/common/lifecycle.ts

@@ -827,3 +827,19 @@ export class DisposableMap<K, V extends IDisposable = IDisposable> implements ID
 		return this._store[Symbol.iterator]();
 	}
 }
+
+/**
+ * Call `then` on a Promise, unless the returned disposable is disposed.
+ */
+export function thenIfNotDisposed<T>(promise: Promise<T>, then: (result: T) => void): IDisposable {
+	let disposed = false;
+	promise.then(result => {
+		if (disposed) {
+			return;
+		}
+		then(result);
+	});
+	return toDisposable(() => {
+		disposed = true;
+	});
+}

src/vs/base/test/common/lifecycle.test.ts

@@ -5,7 +5,7 @@
 
 import assert from 'assert';
 import { Emitter } from '../../common/event.js';
-import { DisposableStore, dispose, IDisposable, markAsSingleton, ReferenceCollection, SafeDisposable, toDisposable } from '../../common/lifecycle.js';
+import { DisposableStore, dispose, IDisposable, markAsSingleton, ReferenceCollection, SafeDisposable, thenIfNotDisposed, toDisposable } from '../../common/lifecycle.js';
 import { ensureNoDisposablesAreLeakedInTestSuite, throwIfDisposablesAreLeaked } from './utils.js';
 
 class Disposable implements IDisposable {
@@ -328,4 +328,30 @@ suite('No Leakage Utilities', () => {
 			toDisposable(() => { }).dispose();
 		});
 	});
+
+	suite('thenIfNotDisposed', () => {
+		const store = ensureNoDisposablesAreLeakedInTestSuite();
+
+		test('normal case', async () => {
+			let called = false;
+			store.add(thenIfNotDisposed(Promise.resolve(123), (result: number) => {
+				assert.strictEqual(result, 123);
+				called = true;
+			}));
+
+			await new Promise(resolve => setTimeout(resolve, 0));
+			assert.strictEqual(called, true);
+		});
+
+		test('disposed before promise resolves', async () => {
+			let called = false;
+			const disposable = thenIfNotDisposed(Promise.resolve(123), () => {
+				called = true;
+			});
+
+			disposable.dispose();
+			await new Promise(resolve => setTimeout(resolve, 0));
+			assert.strictEqual(called, false);
+		});
+	});
 });

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

@@ -19,7 +19,7 @@ import { FuzzyScore } from '../../../../base/common/filters.js';
 import { MarkdownString } from '../../../../base/common/htmlContent.js';
 import { Iterable } from '../../../../base/common/iterator.js';
 import { KeyCode } from '../../../../base/common/keyCodes.js';
-import { Disposable, DisposableStore, IDisposable, dispose, toDisposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, IDisposable, dispose, thenIfNotDisposed, toDisposable } from '../../../../base/common/lifecycle.js';
 import { ResourceMap } from '../../../../base/common/map.js';
 import { FileAccess } from '../../../../base/common/network.js';
 import { clamp } from '../../../../base/common/numbers.js';
@@ -986,7 +986,7 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 
 		part.codeblocks?.forEach((info, i) => {
 			codeBlocksByResponseId[codeBlockStartIndex + i] = info;
-			info.uriPromise.then(uri => {
+			part.addDisposable!(thenIfNotDisposed(info.uriPromise, uri => {
 				if (!uri) {
 					return;
 				}
@@ -998,7 +998,7 @@ export class ChatListItemRenderer extends Disposable implements ITreeRenderer<Ch
 						this.codeBlocksByEditorUri.delete(uri);
 					}
 				}));
-			});
+			}));
 		});
 
 	}

src/vs/workbench/contrib/mergeEditor/browser/utils.ts

@@ -5,7 +5,7 @@
 
 import { ArrayQueue, CompareResult } from '../../../../base/common/arrays.js';
 import { onUnexpectedError } from '../../../../base/common/errors.js';
-import { DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
+import { DisposableStore, IDisposable } from '../../../../base/common/lifecycle.js';
 import { IObservable, autorunOpts } from '../../../../base/common/observable.js';
 import { CodeEditorWidget } from '../../../../editor/browser/widget/codeEditor/codeEditorWidget.js';
 import { IModelDeltaDecoration } from '../../../../editor/common/model.js';
@@ -85,19 +85,6 @@ export function elementAtOrUndefined<T>(arr: T[], index: number): T | undefined
 	return arr[index];
 }
 
-export function thenIfNotDisposed<T>(promise: Promise<T>, then: () => void): IDisposable {
-	let disposed = false;
-	promise.then(() => {
-		if (disposed) {
-			return;
-		}
-		then();
-	});
-	return toDisposable(() => {
-		disposed = true;
-	});
-}
-
 export function setFields<T extends {}>(obj: T, fields: Partial<T>): T {
 	return Object.assign(obj, fields);
 }
@@ -154,4 +141,3 @@ export class PersistentStore<T> {
 		);
 	}
 }
-

src/vs/workbench/contrib/mergeEditor/browser/view/mergeEditor.ts

@@ -10,7 +10,7 @@ import { CancellationToken } from '../../../../../base/common/cancellation.js';
 import { Color } from '../../../../../base/common/color.js';
 import { BugIndicatingError, onUnexpectedError } from '../../../../../base/common/errors.js';
 import { Emitter, Event } from '../../../../../base/common/event.js';
-import { Disposable, DisposableStore, IDisposable, MutableDisposable, toDisposable } from '../../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, IDisposable, MutableDisposable, thenIfNotDisposed, toDisposable } from '../../../../../base/common/lifecycle.js';
 import { autorun, autorunWithStore, IObservable, IReader, observableValue, transaction } from '../../../../../base/common/observable.js';
 import { basename, isEqual } from '../../../../../base/common/resources.js';
 import { isDefined } from '../../../../../base/common/types.js';
@@ -39,7 +39,7 @@ import { readTransientState, writeTransientState } from '../../../codeEditor/bro
 import { MergeEditorInput } from '../mergeEditorInput.js';
 import { IMergeEditorInputModel } from '../mergeEditorInputModel.js';
 import { MergeEditorModel } from '../model/mergeEditorModel.js';
-import { deepMerge, PersistentStore, thenIfNotDisposed } from '../utils.js';
+import { deepMerge, PersistentStore } from '../utils.js';
 import { observableConfigValue } from '../../../../../platform/observable/common/platformObservableUtils.js';
 import { BaseCodeEditorView } from './editors/baseCodeEditorView.js';
 import { ScrollSynchronizer } from './scrollSynchronizer.js';