diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index d73ce6951e6da..b48bda3a6e487 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -699,10 +699,9 @@ public class BubbleData {
         if (DEBUG_BUBBLE_DATA) {
             Log.d(TAG, "setSelectedBubbleInternal: " + bubble);
         }
-        if (!mShowingOverflow && Objects.equals(bubble, mSelectedBubble)) {
+        if (Objects.equals(bubble, mSelectedBubble)) {
             return;
         }
-        // Otherwise, if we are showing the overflow menu, return to the previously selected bubble.
         boolean isOverflow = bubble != null && BubbleOverflow.KEY.equals(bubble.getKey());
         if (bubble != null
                 && !mBubbles.contains(bubble)
@@ -771,6 +770,10 @@ public class BubbleData {
                 Log.e(TAG, "Attempt to expand stack without selected bubble!");
                 return;
             }
+            if (mSelectedBubble.getKey().equals(mOverflow.getKey()) && !mBubbles.isEmpty()) {
+                // Show previously selected bubble instead of overflow menu when expanding.
+                setSelectedBubbleInternal(mBubbles.get(0));
+            }
             if (mSelectedBubble instanceof Bubble) {
                 ((Bubble) mSelectedBubble).markAsAccessedAt(mTimeSource.currentTimeMillis());
             }
@@ -779,16 +782,6 @@ public class BubbleData {
             // Apply ordering and grouping rules from expanded -> collapsed, then save
             // the result.
             mStateChange.orderChanged |= repackAll();
-            // Save the state which should be returned to when expanded (with no other changes)
-
-            if (mShowingOverflow) {
-                // Show previously selected bubble instead of overflow menu on next expansion.
-                if (!mSelectedBubble.getKey().equals(mOverflow.getKey())) {
-                    setSelectedBubbleInternal(mSelectedBubble);
-                } else {
-                    setSelectedBubbleInternal(mBubbles.get(0));
-                }
-            }
             if (mBubbles.indexOf(mSelectedBubble) > 0) {
                 // Move the selected bubble to the top while collapsed.
                 int index = mBubbles.indexOf(mSelectedBubble);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
index 3e3195fe8dc55..b0312e6d6f3c3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
@@ -869,6 +869,35 @@ public class BubbleDataTest extends ShellTestCase {
         assertNotNull(mBubbleData.getOverflowBubbleWithKey(mBubbleA2.getKey()));
     }
 
+    /**
+     * Verifies that after the stack is collapsed with the overflow selected, it will select
+     * the top bubble upon next expansion.
+     */
+    @Test
+    public void test_collapseWithOverflowSelected_nextExpansion() {
+        sendUpdatedEntryAtTime(mEntryA1, 1000);
+        sendUpdatedEntryAtTime(mEntryA2, 2000);
+        mBubbleData.setExpanded(true);
+
+        mBubbleData.setListener(mListener);
+
+        // Select the overflow
+        mBubbleData.setShowingOverflow(true);
+        mBubbleData.setSelectedBubble(mBubbleData.getOverflow());
+        verifyUpdateReceived();
+        assertSelectionChangedTo(mBubbleData.getOverflow());
+
+        // Collapse
+        mBubbleData.setExpanded(false);
+        verifyUpdateReceived();
+        assertSelectionNotChanged();
+
+        // Expand (here we should select the new bubble)
+        mBubbleData.setExpanded(true);
+        verifyUpdateReceived();
+        assertSelectionChangedTo(mBubbleA2);
+    }
+
     private void verifyUpdateReceived() {
         verify(mListener).applyUpdate(mUpdateCaptor.capture());
         reset(mListener);
@@ -902,7 +931,7 @@ public class BubbleDataTest extends ShellTestCase {
         assertWithMessage("selectionChanged").that(update.selectionChanged).isFalse();
     }
 
-    private void assertSelectionChangedTo(Bubble bubble) {
+    private void assertSelectionChangedTo(BubbleViewProvider bubble) {
         BubbleData.Update update = mUpdateCaptor.getValue();
         assertWithMessage("selectionChanged").that(update.selectionChanged).isTrue();
         assertWithMessage("selectedBubble").that(update.selectedBubble).isEqualTo(bubble);
@@ -925,7 +954,6 @@ public class BubbleDataTest extends ShellTestCase {
         assertThat(update.overflowBubbles).isEqualTo(bubbles);
     }
 
-
     private BubbleEntry createBubbleEntry(int userId, String notifKey, String packageName,
             NotificationListenerService.Ranking ranking) {
         return createBubbleEntry(userId, notifKey, packageName, ranking, 1000);