Adopt thenRegisterOrDispose helper (microsoft#261287)

Adopt thenRegisterOrDispose helper

Fixes microsoft#261211

Author: Tyriar

src/vs/base/common/lifecycle.ts

@@ -824,12 +824,13 @@ export function thenIfNotDisposed<T>(promise: Promise<T>, then: (result: T) => v
  * disposable or register it to the {@link DisposableStore}, depending on whether the store is
  * disposed or not.
  */
-export function thenRegisterOrDispose<T extends IDisposable>(promise: Promise<T>, store: DisposableStore): void {
-	promise.then(ref => {
+export function thenRegisterOrDispose<T extends IDisposable>(promise: Promise<T>, store: DisposableStore): Promise<T> {
+	return promise.then(disposable => {
 		if (store.isDisposed) {
-			ref.dispose();
+			disposable.dispose();
 		} else {
-			store.add(ref);
+			store.add(disposable);
 		}
+		return disposable;
 	});
 }

src/vs/platform/files/node/watcher/nodejs/nodejsWatcherLib.ts

@@ -7,7 +7,7 @@ import { watch, promises } from 'fs';
 import { RunOnceWorker, ThrottledWorker } from '../../../../../base/common/async.js';
 import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js';
 import { isEqual, isEqualOrParent } from '../../../../../base/common/extpath.js';
-import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, IDisposable, thenRegisterOrDispose, toDisposable } from '../../../../../base/common/lifecycle.js';
 import { normalizeNFC } from '../../../../../base/common/normalization.js';
 import { basename, dirname, join } from '../../../../../base/common/path.js';
 import { isLinux, isMacintosh } from '../../../../../base/common/platform.js';
@@ -160,12 +160,7 @@ export class NodeJSFileWatcherLibrary extends Disposable {
 			}
 
 			if (error) {
-				const watchDisposable = await this.doWatch(isDirectory);
-				if (!disposables.isDisposed) {
-					disposables.add(watchDisposable);
-				} else {
-					watchDisposable.dispose();
-				}
+				await thenRegisterOrDispose(this.doWatch(isDirectory), disposables);
 			} else if (change) {
 				if (typeof change.cId === 'number' || typeof this.request.correlationId === 'number') {
 					// Re-emit this change with the correlation id of the request

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookEntry.ts

@@ -6,7 +6,7 @@
 import { streamToBuffer } from '../../../../../base/common/buffer.js';
 import { CancellationToken } from '../../../../../base/common/cancellation.js';
 import { StringSHA1 } from '../../../../../base/common/hash.js';
-import { DisposableStore, IReference } from '../../../../../base/common/lifecycle.js';
+import { DisposableStore, IReference, thenRegisterOrDispose } from '../../../../../base/common/lifecycle.js';
 import { ResourceMap, ResourceSet } from '../../../../../base/common/map.js';
 import { Schemas } from '../../../../../base/common/network.js';
 import { ITransaction, IObservable, observableValue, autorun, transaction, ObservablePromise } from '../../../../../base/common/observable.js';
@@ -974,12 +974,7 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 			this.cellTextModelMap.set(cell.uri, model);
 			return model;
 		} else {
-			const textEditorModel = await this.textModelService.createModelReference(cell.uri);
-			if (this._store.isDisposed) {
-				textEditorModel.dispose();
-			} else {
-				this._register(textEditorModel);
-			}
+			const textEditorModel = await thenRegisterOrDispose(this.textModelService.createModelReference(cell.uri), this._store);
 			const model = textEditorModel.object.textEditorModel;
 			this.cellTextModelMap.set(cell.uri, model);
 			return model;

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingTimeline.ts

@@ -8,7 +8,7 @@
 import { equals as arraysEqual, binarySearch2 } from '../../../../../base/common/arrays.js';
 import { findLast } from '../../../../../base/common/arraysFind.js';
 import { Iterable } from '../../../../../base/common/iterator.js';
-import { DisposableStore } from '../../../../../base/common/lifecycle.js';
+import { DisposableStore, thenRegisterOrDispose } from '../../../../../base/common/lifecycle.js';
 import { ResourceMap } from '../../../../../base/common/map.js';
 import { equals as objectsEqual } from '../../../../../base/common/objects.js';
 import { derived, derivedOpts, IObservable, ITransaction, ObservablePromise, observableValue, transaction } from '../../../../../base/common/observable.js';
@@ -416,14 +416,9 @@ export class ChatEditingTimeline {
 			}
 			const store = new DisposableStore();
 			reader.store.add(store);
-			const referencesPromise = Promise.all([firstSnapshotUri, lastSnapshotUri].map(u => this._textModelService.createModelReference(u))).then(refs => {
-				if (store.isDisposed) {
-					refs.forEach(ref => ref.dispose());
-				} else {
-					refs.forEach(ref => store.add(ref));
-				}
-				return refs;
-			});
+			const referencesPromise = Promise.all([firstSnapshotUri, lastSnapshotUri].map(u => {
+				return thenRegisterOrDispose(this._textModelService.createModelReference(u), store);
+			}));
 			return new ObservablePromise(referencesPromise);
 		});
 		const diff = derived((reader): ObservablePromise<IEditSessionEntryDiff> | undefined => {