Merge pull request microsoft#254981 from microsoft/osortega/fix-cancel-search

Fix for canceling pending copilot search request

Author: osortega

src/vs/base/browser/ui/tree/asyncDataTree.ts

@@ -32,7 +32,7 @@ interface IAsyncDataTreeNode<TInput, T> {
 	readonly parent: IAsyncDataTreeNode<TInput, T> | null;
 	readonly children: IAsyncDataTreeNode<TInput, T>[];
 	readonly id?: string | null;
-	refreshPromise: Promise<void> | undefined;
+	refreshPromise: CancelablePromise<void> | undefined;
 	hasChildren: boolean;
 	stale: boolean;
 	slow: boolean;
@@ -528,7 +528,7 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 	private readonly findController?: AsyncFindController<TInput, T, TFilterData>;
 	private readonly getDefaultCollapseState: { (e: T): undefined | ObjectTreeElementCollapseState.PreserveOrCollapsed | ObjectTreeElementCollapseState.PreserveOrExpanded };
 
-	private readonly subTreeRefreshPromises = new Map<IAsyncDataTreeNode<TInput, T>, Promise<void>>();
+	private readonly subTreeRefreshPromises = new Map<IAsyncDataTreeNode<TInput, T>, CancelablePromise<void>>();
 	private readonly refreshPromises = new Map<IAsyncDataTreeNode<TInput, T>, CancelablePromise<Iterable<T>>>();
 
 	protected readonly identityProvider?: IIdentityProvider<T>;
@@ -769,8 +769,7 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 	}
 
 	async setInput(input: TInput, viewState?: IAsyncDataTreeViewState): Promise<void> {
-		this.refreshPromises.forEach(promise => promise.cancel());
-		this.refreshPromises.clear();
+		this.cancelAllRefreshPromises();
 
 		this.root.element = input!;
 
@@ -792,6 +791,14 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 		await this._updateChildren(element, recursive, rerender, undefined, options);
 	}
 
+	cancelAllRefreshPromises(): void {
+		this.refreshPromises.forEach(promise => promise.cancel());
+		this.refreshPromises.clear();
+
+		this.subTreeRefreshPromises.forEach(promise => promise.cancel());
+		this.subTreeRefreshPromises.clear();
+	}
+
 	private async _updateChildren(element: TInput | T = this.root.element, recursive = true, rerender = false, viewStateContext?: IAsyncDataTreeViewStateContext<TInput, T>, options?: IAsyncDataTreeUpdateChildrenOptions<T>): Promise<void> {
 		if (typeof this.root.element === 'undefined') {
 			throw new TreeError(this.user, 'Tree input not set');
@@ -875,7 +882,7 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 		}
 
 		if (node.refreshPromise) {
-			await this.root.refreshPromise;
+			await node.refreshPromise;
 			await Event.toPromise(this._onDidRender.event);
 		}
 
@@ -886,7 +893,7 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 		const result = this.tree.expand(node === this.root ? null : node, recursive);
 
 		if (node.refreshPromise) {
-			await this.root.refreshPromise;
+			await node.refreshPromise;
 			await Event.toPromise(this._onDidRender.event);
 		}
 
@@ -1088,28 +1095,26 @@ export class AsyncDataTree<TInput, T, TFilterData = void> implements IDisposable
 				return;
 			}
 		}
-
 		return this.doRefreshSubTree(node, recursive, viewStateContext);
 	}
 
 	private async doRefreshSubTree(node: IAsyncDataTreeNode<TInput, T>, recursive: boolean, viewStateContext?: IAsyncDataTreeViewStateContext<TInput, T>): Promise<void> {
-		let done: () => void;
-		node.refreshPromise = new Promise(c => done = c);
-		this.subTreeRefreshPromises.set(node, node.refreshPromise);
+		const cancelablePromise = createCancelablePromise(async () => {
+			const childrenToRefresh = await this.doRefreshNode(node, recursive, viewStateContext);
+			node.stale = false;
 
-		node.refreshPromise.finally(() => {
+			await Promises.settled(childrenToRefresh.map(child => this.doRefreshSubTree(child, recursive, viewStateContext)));
+		});
+
+		node.refreshPromise = cancelablePromise;
+		this.subTreeRefreshPromises.set(node, cancelablePromise);
+
+		cancelablePromise.finally(() => {
 			node.refreshPromise = undefined;
 			this.subTreeRefreshPromises.delete(node);
 		});
 
-		try {
-			const childrenToRefresh = await this.doRefreshNode(node, recursive, viewStateContext);
-			node.stale = false;
-
-			await Promises.settled(childrenToRefresh.map(child => this.doRefreshSubTree(child, recursive, viewStateContext)));
-		} finally {
-			done!();
-		}
+		return cancelablePromise;
 	}
 
 	private async doRefreshNode(node: IAsyncDataTreeNode<TInput, T>, recursive: boolean, viewStateContext?: IAsyncDataTreeViewStateContext<TInput, T>): Promise<IAsyncDataTreeNode<TInput, T>[]> {

src/vs/workbench/contrib/search/browser/searchView.ts

@@ -543,9 +543,9 @@ export class SearchView extends ViewPane {
 		// Subscribe to AI search result changes and update the tree when new AI results are reported
 		this._onAIResultChangedDisposable?.dispose();
 		this._onAIResultChangedDisposable = this._register(
-			this.viewModel.searchResult.aiTextSearchResult.onChange(() => {
+			this.viewModel.searchResult.aiTextSearchResult.onChange((e) => {
 				// Only refresh the AI node, not the whole tree
-				if (this.tree && this.tree.hasNode(this.searchResult.aiTextSearchResult)) {
+				if (this.tree && this.tree.hasNode(this.searchResult.aiTextSearchResult) && !e.removed) {
 					this.tree.updateChildren(this.searchResult.aiTextSearchResult);
 				}
 			})
@@ -1349,6 +1349,7 @@ export class SearchView extends ViewPane {
 			this.searchWidget.clear();
 		}
 		this.viewModel.cancelSearch();
+		this.viewModel.cancelAISearch();
 		this.tree.ariaLabel = nls.localize('emptySearch', "Empty Search");
 
 		this.accessibilitySignalService.playSignal(AccessibilitySignal.clear);
@@ -1858,17 +1859,17 @@ export class SearchView extends ViewPane {
 
 	public clearAIResults() {
 		this.model.searchResult.aiTextSearchResult.hidden = true;
-		if (!this._pendingSemanticSearchPromise) {
-			this._cachedResults = undefined;
-			this._cachedKeywords = [];
-			this.model.cancelAISearch(true);
-			this.model.clearAiSearchResults();
-		}
+		this.refreshTreeController.clearAllPending();
+		this._pendingSemanticSearchPromise = undefined;
+		this._cachedResults = undefined;
+		this._cachedKeywords = [];
+		this.model.cancelAISearch(true);
+		this.model.clearAiSearchResults();
 	}
 
 	public async requestAIResults() {
 		this.logService.info(`SearchView: Requesting semantic results from keybinding. Cached: ${!!this.cachedResults}`);
-		if (!this.cachedResults || this.cachedResults.results.length === 0) {
+		if ((!this.cachedResults || this.cachedResults.results.length === 0) && !this._pendingSemanticSearchPromise) {
 			this.clearAIResults();
 		}
 		this.model.searchResult.aiTextSearchResult.hidden = false;
@@ -2646,6 +2647,10 @@ class RefreshTreeController extends Disposable {
 
 	private queuedIChangeEvents: IChangeEvent[] = [];
 
+	public clearAllPending(): void {
+		this.searchView.getControl().cancelAllRefreshPromises();
+	}
+
 	public async queue(e?: IChangeEvent): Promise<void> {
 		if (e) {
 			this.queuedIChangeEvents.push(e);