diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
index 4e19e06..d54be0e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
@@ -1665,6 +1665,7 @@
         }
         beforeExpandedViewAnimation();
 
+        updateBadgesAndZOrder(false /* setBadgeForCollapsedStack */);
         mBubbleContainer.setActiveController(mExpandedAnimationController);
         updateOverflowVisibility();
         updatePointerPosition();
@@ -1878,7 +1879,7 @@
                                 mExpandedBubble));
                     }
                     updateOverflowVisibility();
-
+                    updateBadgesAndZOrder(true /* setBadgeForCollapsedStack */);
                     afterExpandedViewAnimation();
                     if (previouslySelected != null) {
                         previouslySelected.setContentVisibility(false);
@@ -2626,7 +2627,6 @@
         }
 
         mStackOnLeftOrWillBe = mStackAnimationController.isStackOnLeftSide();
-        updateBadgesAndZOrder(false /* setBadgeForCollapsedStack */);
     }
 
     /**