diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index 079a11c..59bdf4b 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -362,18 +362,30 @@
     public void updateIconBadges(Set<PackageUserKey> updatedBadges) {
         final PackageUserKey packageUserKey = new PackageUserKey(null, null);
         for (int j = 0; j < mAH.length; j++) {
-            if (mAH[j].recyclerView != null) {
-                final int n = mAH[j].recyclerView.getChildCount();
-                for (int i = 0; i < n; i++) {
-                    View child = mAH[j].recyclerView.getChildAt(i);
-                    if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
-                        continue;
-                    }
-                    ItemInfo info = (ItemInfo) child.getTag();
-                    if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
-                        ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
-                    }
-                }
+            updateIconBadges(updatedBadges, packageUserKey, mAH[j].recyclerView);
+        }
+        if (mHeader != null) {
+            updateIconBadges(updatedBadges, packageUserKey, mHeader.getPredictionRow());
+        }
+    }
+
+    private void updateIconBadges(Set<PackageUserKey> updatedBadges, PackageUserKey packageUserKey,
+            ViewGroup parent) {
+        if (parent == null) {
+            return;
+        }
+        final int n = parent.getChildCount();
+        for (int i = 0; i < n; i++) {
+            View child = parent.getChildAt(i);
+            if (child instanceof PredictionRowView) {
+                updateIconBadges(updatedBadges, packageUserKey, (PredictionRowView) child);
+            }
+            if (!(child instanceof BubbleTextView) || !(child.getTag() instanceof ItemInfo)) {
+                continue;
+            }
+            ItemInfo info = (ItemInfo) child.getTag();
+            if (packageUserKey.updateFromItemInfo(info) && updatedBadges.contains(packageUserKey)) {
+                ((BubbleTextView) child).applyBadgeState(info, true /* animate */);
             }
         }
     }
diff --git a/src/com/android/launcher3/notification/NotificationListener.java b/src/com/android/launcher3/notification/NotificationListener.java
index 7b70df7..114b2b8 100644
--- a/src/com/android/launcher3/notification/NotificationListener.java
+++ b/src/com/android/launcher3/notification/NotificationListener.java
@@ -62,6 +62,7 @@
 
     private static NotificationListener sNotificationListenerInstance = null;
     private static NotificationsChangedListener sNotificationsChangedListener;
+    private static StatusBarNotificationsChangedListener sStatusBarNotificationsChangedListener;
     private static boolean sIsConnected;
     private static boolean sIsCreated;
 
@@ -180,10 +181,19 @@
         }
     }
 
+    public static void setStatusBarNotificationsChangedListener
+            (StatusBarNotificationsChangedListener listener) {
+        sStatusBarNotificationsChangedListener = listener;
+    }
+
     public static void removeNotificationsChangedListener() {
         sNotificationsChangedListener = null;
     }
 
+    public static void removeStatusBarNotificationsChangedListener() {
+        sStatusBarNotificationsChangedListener = null;
+    }
+
     @Override
     public void onListenerConnected() {
         super.onListenerConnected();
@@ -205,7 +215,10 @@
     public void onNotificationPosted(final StatusBarNotification sbn) {
         super.onNotificationPosted(sbn);
         mWorkerHandler.obtainMessage(MSG_NOTIFICATION_POSTED, new NotificationPostedMsg(sbn))
-                .sendToTarget();
+            .sendToTarget();
+        if (sStatusBarNotificationsChangedListener != null) {
+            sStatusBarNotificationsChangedListener.onNotificationPosted(sbn);
+        }
     }
 
     /**
@@ -227,10 +240,13 @@
     public void onNotificationRemoved(final StatusBarNotification sbn) {
         super.onNotificationRemoved(sbn);
         Pair<PackageUserKey, NotificationKeyData> packageUserKeyAndNotificationKey
-                = new Pair<>(PackageUserKey.fromNotification(sbn),
-                        NotificationKeyData.fromNotification(sbn));
+            = new Pair<>(PackageUserKey.fromNotification(sbn),
+            NotificationKeyData.fromNotification(sbn));
         mWorkerHandler.obtainMessage(MSG_NOTIFICATION_REMOVED, packageUserKeyAndNotificationKey)
-                .sendToTarget();
+            .sendToTarget();
+        if (sStatusBarNotificationsChangedListener != null) {
+            sStatusBarNotificationsChangedListener.onNotificationRemoved(sbn);
+        }
 
         NotificationGroup notificationGroup = mNotificationGroupMap.get(sbn.getGroupKey());
         if (notificationGroup != null) {
@@ -318,4 +334,9 @@
                 NotificationKeyData notificationKey);
         void onNotificationFullRefresh(List<StatusBarNotification> activeNotifications);
     }
+
+    public interface StatusBarNotificationsChangedListener {
+        void onNotificationPosted(StatusBarNotification sbn);
+        void onNotificationRemoved(StatusBarNotification sbn);
+    }
 }
diff --git a/src/com/android/launcher3/popup/PopupContainerWithArrow.java b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
index cedf291..f90abb4 100644
--- a/src/com/android/launcher3/popup/PopupContainerWithArrow.java
+++ b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
@@ -68,7 +68,6 @@
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.dragndrop.DragView;
-import com.android.launcher3.graphics.IconPalette;
 import com.android.launcher3.graphics.TriangleShape;
 import com.android.launcher3.logging.LoggerUtils;
 import com.android.launcher3.notification.NotificationInfo;
@@ -650,7 +649,7 @@
                 // reopen the container to ensure measurements etc. all work out. While this could
                 // be quite janky, in practice the user would typically see a small flicker as the
                 // animation restarts partway through, and this is a very rare edge case anyway.
-                ((PopupContainerWithArrow) getParent()).close(false);
+                close(false);
                 PopupContainerWithArrow.showForIcon(mOriginalIcon);
             }
         } else if (onClickListener == null && widgetsView != null) {
@@ -658,7 +657,7 @@
             if (mSystemShortcutContainer != this) {
                 mSystemShortcutContainer.removeView(widgetsView);
             } else {
-                ((PopupContainerWithArrow) getParent()).close(false);
+                close(false);
                 PopupContainerWithArrow.showForIcon(mOriginalIcon);
             }
         }