From 9b9508dfc3b0c91f1e00a763f8f387864c8ce963 Mon Sep 17 00:00:00 2001
From: andreamah <andrea_mah22@hotmail.com>
Date: Thu, 26 Sep 2024 13:22:27 -0700
Subject: [PATCH] await on expand Fixes #229878

---
 .../search/browser/searchActionsRemoveReplace.ts      | 11 +++++++----
 .../contrib/search/test/browser/searchActions.test.ts |  8 ++++----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/vs/workbench/contrib/search/browser/searchActionsRemoveReplace.ts b/src/vs/workbench/contrib/search/browser/searchActionsRemoveReplace.ts
index 66211a778d0..4ce73c9cd8f 100644
--- a/src/vs/workbench/contrib/search/browser/searchActionsRemoveReplace.ts
+++ b/src/vs/workbench/contrib/search/browser/searchActionsRemoveReplace.ts
@@ -126,7 +126,7 @@ registerAction2(class RemoveAction extends Action2 {
 
 		if (focusElement && shouldRefocusMatch) {
 			if (!nextFocusElement) {
-				nextFocusElement = getLastNodeFromSameType(viewer, focusElement);
+				nextFocusElement = await getLastNodeFromSameType(viewer, focusElement);
 			}
 
 			if (nextFocusElement && !arrayContainsElementOrParent(nextFocusElement, elementsToRemove)) {
@@ -294,7 +294,7 @@ async function performReplace(accessor: ServicesAccessor,
 
 	if (focusElement) {
 		if (!nextFocusElement) {
-			nextFocusElement = getLastNodeFromSameType(viewer, focusElement);
+			nextFocusElement = await getLastNodeFromSameType(viewer, focusElement);
 		}
 
 		if (nextFocusElement) {
@@ -389,13 +389,16 @@ export function getElementToFocusAfterRemoved(viewer: WorkbenchCompressibleAsync
 /***
  * Finds the last element in the tree with the same type as `element`
  */
-export function getLastNodeFromSameType(viewer: WorkbenchCompressibleAsyncDataTree<SearchResult, RenderableMatch>, element: RenderableMatch): RenderableMatch | undefined {
+export async function getLastNodeFromSameType(viewer: WorkbenchCompressibleAsyncDataTree<SearchResult, RenderableMatch>, element: RenderableMatch): Promise<RenderableMatch | undefined> {
 	let lastElem: RenderableMatch | null = viewer.lastVisibleElement ?? null;
 
 	while (lastElem) {
 		const compareVal = compareLevels(element, lastElem);
 		if (compareVal === -1) {
-			viewer.expand(lastElem);
+			const expanded = await viewer.expand(lastElem);
+			if (!expanded) {
+				return lastElem;
+			}
 			lastElem = viewer.lastVisibleElement;
 		} else if (compareVal === 1) {
 			const potentialLastElem = viewer.getParentElement(lastElem);
diff --git a/src/vs/workbench/contrib/search/test/browser/searchActions.test.ts b/src/vs/workbench/contrib/search/test/browser/searchActions.test.ts
index 111f325eff7..b6f116f4f0c 100644
--- a/src/vs/workbench/contrib/search/test/browser/searchActions.test.ts
+++ b/src/vs/workbench/contrib/search/test/browser/searchActions.test.ts
@@ -75,25 +75,25 @@ suite('Search Actions', () => {
 		assert.strictEqual(data[4], actual);
 	});
 
-	test('Find last FileMatch in Tree', function () {
+	test('Find last FileMatch in Tree', async function () {
 		const fileMatch1 = aFileMatch();
 		const fileMatch2 = aFileMatch();
 		const fileMatch3 = aFileMatch();
 		const data = [fileMatch1, aMatch(fileMatch1), fileMatch2, aMatch(fileMatch2), fileMatch3, aMatch(fileMatch3)];
 		const tree = aTree(data);
 
-		const actual = getLastNodeFromSameType(tree, fileMatch1);
+		const actual = await getLastNodeFromSameType(tree, fileMatch1);
 		assert.strictEqual(fileMatch3, actual);
 	});
 
-	test('Find last Match in Tree', function () {
+	test('Find last Match in Tree', async function () {
 		const fileMatch1 = aFileMatch();
 		const fileMatch2 = aFileMatch();
 		const fileMatch3 = aFileMatch();
 		const data = [fileMatch1, aMatch(fileMatch1), fileMatch2, aMatch(fileMatch2), fileMatch3, aMatch(fileMatch3)];
 		const tree = aTree(data);
 
-		const actual = getLastNodeFromSameType(tree, aMatch(fileMatch1));
+		const actual = await getLastNodeFromSameType(tree, aMatch(fileMatch1));
 		assert.strictEqual(data[5], actual);
 	});