diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index fb2f177..3cc91de 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -560,7 +560,7 @@
 
                 // Need to check for !appCancel here because the notification may have
                 // previously been dismissed & entry.isRowDismissed would still be true
-                boolean userRemovedNotif = (entry.isRowDismissed() && !isAppCancel)
+                boolean userRemovedNotif = (entry != null && entry.isRowDismissed() && !isAppCancel)
                         || isClearAll || isUserDimiss || isSummaryCancel;
 
                 if (isSummaryOfBubbles) {
@@ -570,7 +570,7 @@
                 // The bubble notification sticks around in the data as long as the bubble is
                 // not dismissed and the app hasn't cancelled the notification.
                 Bubble bubble = mBubbleData.getBubbleWithKey(key);
-                boolean bubbleExtended = entry.isBubble() && userRemovedNotif;
+                boolean bubbleExtended = entry != null && entry.isBubble() && userRemovedNotif;
                 if (bubbleExtended) {
                     bubble.setShowInShadeWhenBubble(false);
                     bubble.setShowBubbleDot(false);
@@ -579,7 +579,7 @@
                     }
                     mNotificationEntryManager.updateNotifications();
                     return true;
-                } else if (!userRemovedNotif) {
+                } else if (!userRemovedNotif && entry != null) {
                     // This wasn't a user removal so we should remove the bubble as well
                     mBubbleData.notificationEntryRemoved(entry, DISMISS_NOTIF_CANCEL);
                     return false;
@@ -939,13 +939,11 @@
             final Bubble expandedBubble = mStackView != null
                     ? mStackView.getExpandedBubble()
                     : null;
-            if (expandedBubble == null) {
-                return;
-            }
-            if (expandedBubble.getDisplayId() == displayId) {
+            int expandedId = expandedBubble != null ? expandedBubble.getDisplayId() : -1;
+            if (mStackView != null && mStackView.isExpanded() && expandedId == displayId) {
                 mBubbleData.setExpanded(false);
-                expandedBubble.getExpandedView().notifyDisplayEmpty();
             }
+            mBubbleData.notifyDisplayEmpty(displayId);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
index 81c8da8..eb826e5 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -392,6 +392,22 @@
         dispatchPendingChanges();
     }
 
+    /**
+     * Indicates that the provided display is no longer in use and should be cleaned up.
+     *
+     * @param displayId the id of the display to clean up.
+     */
+    void notifyDisplayEmpty(int displayId) {
+        for (Bubble b : mBubbles) {
+            if (b.getDisplayId() == displayId) {
+                if (b.getExpandedView() != null) {
+                    b.getExpandedView().notifyDisplayEmpty();
+                }
+                return;
+            }
+        }
+    }
+
     private void dispatchPendingChanges() {
         if (mListener != null && mStateChange.anythingChanged()) {
             mListener.applyUpdate(mStateChange);