diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index 15350fb19209daaac771952c00cbbd9f96fb35fd..96aaf02cb5e3763ccdfe85865d04d2d1fedaf526 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -1799,11 +1799,12 @@ public class BubbleController implements ConfigurationChangeListener,
         @Override
         public void removeBubble(Bubble removedBubble) {
             if (mLayerView != null) {
-                mLayerView.removeBubble(removedBubble);
-                if (!mBubbleData.hasBubbles() && !isStackExpanded()) {
-                    mLayerView.setVisibility(INVISIBLE);
-                    removeFromWindowManagerMaybe();
-                }
+                mLayerView.removeBubble(removedBubble, () -> {
+                    if (!mBubbleData.hasBubbles() && !isStackExpanded()) {
+                        mLayerView.setVisibility(INVISIBLE);
+                        removeFromWindowManagerMaybe();
+                    }
+                });
             }
         }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
index 78a41f759d96b7c36908c3810fd7cb5a3ca3a23a..42799d975e1bc51c92bed90d946f6809c13233df 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
@@ -242,13 +242,17 @@ public class BubbleBarLayerView extends FrameLayout
     }
 
     /** Removes the given {@code bubble}. */
-    public void removeBubble(Bubble bubble) {
+    public void removeBubble(Bubble bubble, Runnable endAction) {
+        Runnable cleanUp = () -> {
+            bubble.cleanupViews();
+            endAction.run();
+        };
         if (mBubbleData.getBubbles().isEmpty()) {
             // we're removing the last bubble. collapse the expanded view and cleanup bubble views
             // at the end.
-            collapse(bubble::cleanupViews);
+            collapse(cleanUp);
         } else {
-            bubble.cleanupViews();
+            cleanUp.run();
         }
     }
 
@@ -264,6 +268,9 @@ public class BubbleBarLayerView extends FrameLayout
      */
     public void collapse(@Nullable Runnable endAction) {
         if (!mIsExpanded) {
+            if (endAction != null) {
+                endAction.run();
+            }
             return;
         }
         mIsExpanded = false;