decouple Bubble from NotificationEntry

1. decouple Bubble from NotificationEntry
2. save title from Bubble into BubbleEntity
3. copied boolean values from NotificationEntry into Bubble for UI
variants (e.g. isVisuallyInterruptive, isClearable,
shouldSuppressNotificationDot... e.t.c)

Bug: 151474524
Change-Id: I606c6ff93b3dc3867b4d0a6129d7117d9999c170
Test: manual
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index 7f78ddf..6da7bc8 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -15,7 +15,6 @@
  */
 package com.android.systemui.bubbles;
 
-import static android.app.Notification.FLAG_BUBBLE;
 import static android.os.AsyncTask.Status.FINISHED;
 import static android.view.Display.INVALID_DISPLAY;
 
@@ -29,21 +28,19 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ApplicationInfo;
-import android.content.pm.LauncherApps;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
 import android.graphics.Path;
-import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
-import android.os.Bundle;
+import android.graphics.drawable.Icon;
 import android.os.UserHandle;
 import android.provider.Settings;
-import android.service.notification.StatusBarNotification;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.logging.InstanceId;
 import com.android.systemui.shared.system.SysUiStatsLog;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
@@ -57,17 +54,12 @@
 class Bubble implements BubbleViewProvider {
     private static final String TAG = "Bubble";
 
-    /**
-     * NotificationEntry associated with the bubble. A null value implies this bubble is loaded
-     * from disk.
-     */
-    @Nullable
-    private NotificationEntry mEntry;
     private final String mKey;
 
     private long mLastUpdated;
     private long mLastAccessed;
 
+    @Nullable
     private BubbleController.NotificationSuppressionChangedListener mSuppressionListener;
 
     /** Whether the bubble should show a dot for the notification indicating updated content. */
@@ -75,8 +67,6 @@
 
     /** Whether flyout text should be suppressed, regardless of any other flags or state. */
     private boolean mSuppressFlyout;
-    /** Whether this bubble should auto expand regardless of the normal flag, used for overflow. */
-    private boolean mShouldAutoExpand;
 
     // Items that are typically loaded later
     private String mAppName;
@@ -92,6 +82,7 @@
      * Presentational info about the flyout.
      */
     public static class FlyoutMessage {
+        @Nullable public Icon senderIcon;
         @Nullable public Drawable senderAvatar;
         @Nullable public CharSequence senderName;
         @Nullable public CharSequence message;
@@ -109,16 +100,39 @@
     private UserHandle mUser;
     @NonNull
     private String mPackageName;
+    @Nullable
+    private String mTitle;
+    @Nullable
+    private Icon mIcon;
+    private boolean mIsBubble;
+    private boolean mIsVisuallyInterruptive;
+    private boolean mIsClearable;
+    private boolean mShouldSuppressNotificationDot;
+    private boolean mShouldSuppressNotificationList;
+    private boolean mShouldSuppressPeek;
     private int mDesiredHeight;
     @DimenRes
     private int mDesiredHeightResId;
 
+    /** for logging **/
+    @Nullable
+    private InstanceId mInstanceId;
+    @Nullable
+    private String mChannelId;
+    private int mNotificationId;
+    private int mAppUid = -1;
+
+    @Nullable
+    private PendingIntent mIntent;
+    @Nullable
+    private PendingIntent mDeleteIntent;
+
     /**
      * Create a bubble with limited information based on given {@link ShortcutInfo}.
      * Note: Currently this is only being used when the bubble is persisted to disk.
      */
     Bubble(@NonNull final String key, @NonNull final ShortcutInfo shortcutInfo,
-            final int desiredHeight, final int desiredHeightResId) {
+            final int desiredHeight, final int desiredHeightResId, @Nullable final String title) {
         Objects.requireNonNull(key);
         Objects.requireNonNull(shortcutInfo);
         mShortcutInfo = shortcutInfo;
@@ -126,8 +140,10 @@
         mFlags = 0;
         mUser = shortcutInfo.getUserHandle();
         mPackageName = shortcutInfo.getPackage();
+        mIcon = shortcutInfo.getIcon();
         mDesiredHeight = desiredHeight;
         mDesiredHeightResId = desiredHeightResId;
+        mTitle = title;
     }
 
     /** Used in tests when no UI is required. */
@@ -145,12 +161,6 @@
         return mKey;
     }
 
-    @Nullable
-    public NotificationEntry getEntry() {
-        return mEntry;
-    }
-
-    @NonNull
     public UserHandle getUser() {
         return mUser;
     }
@@ -203,14 +213,7 @@
 
     @Nullable
     public String getTitle() {
-        final CharSequence titleCharSeq;
-        if (mEntry == null) {
-            titleCharSeq = null;
-        } else {
-            titleCharSeq = mEntry.getSbn().getNotification().extras.getCharSequence(
-                    Notification.EXTRA_TITLE);
-        }
-        return titleCharSeq != null ? titleCharSeq.toString() : null;
+        return mTitle;
     }
 
     /**
@@ -331,17 +334,44 @@
     void setEntry(@NonNull final NotificationEntry entry) {
         Objects.requireNonNull(entry);
         Objects.requireNonNull(entry.getSbn());
-        mEntry = entry;
         mLastUpdated = entry.getSbn().getPostTime();
-        mFlags = entry.getSbn().getNotification().flags;
+        mIsBubble = entry.getSbn().getNotification().isBubbleNotification();
         mPackageName = entry.getSbn().getPackageName();
         mUser = entry.getSbn().getUser();
+        mTitle = getTitle(entry);
+        mIsClearable = entry.isClearable();
+        mShouldSuppressNotificationDot = entry.shouldSuppressNotificationDot();
+        mShouldSuppressNotificationList = entry.shouldSuppressNotificationList();
+        mShouldSuppressPeek = entry.shouldSuppressPeek();
+        mChannelId = entry.getSbn().getNotification().getChannelId();
+        mNotificationId = entry.getSbn().getId();
+        mAppUid = entry.getSbn().getUid();
+        mInstanceId = entry.getSbn().getInstanceId();
+        mFlyoutMessage = BubbleViewInfoTask.extractFlyoutMessage(entry);
+        if (entry.getRanking() != null) {
+            mShortcutInfo = entry.getRanking().getShortcutInfo() != null
+                    ? entry.getRanking().getShortcutInfo() : mShortcutInfo;
+            mIsVisuallyInterruptive = entry.getRanking().visuallyInterruptive();
+        }
         if (entry.getBubbleMetadata() != null) {
+            mFlags = entry.getBubbleMetadata().getFlags();
             mDesiredHeight = entry.getBubbleMetadata().getDesiredHeight();
             mDesiredHeightResId = entry.getBubbleMetadata().getDesiredHeightResId();
+            mIcon = entry.getBubbleMetadata().getIcon();
+            mIntent = entry.getBubbleMetadata().getIntent();
+            mDeleteIntent = entry.getBubbleMetadata().getDeleteIntent();
         }
     }
 
+    @Nullable
+    Icon getIcon() {
+        return mIcon;
+    }
+
+    boolean isVisuallyInterruptive() {
+        return mIsVisuallyInterruptive;
+    }
+
     /**
      * @return the last time this bubble was updated or accessed, whichever is most recent.
      */
@@ -364,6 +394,19 @@
         return mExpandedView != null ? mExpandedView.getVirtualDisplayId() : INVALID_DISPLAY;
     }
 
+    public InstanceId getInstanceId() {
+        return mInstanceId;
+    }
+
+    @Nullable
+    public String getChannelId() {
+        return mChannelId;
+    }
+
+    public int getNotificationId() {
+        return mNotificationId;
+    }
+
     /**
      * Should be invoked whenever a Bubble is accessed (selected while expanded).
      */
@@ -384,24 +427,19 @@
      * Whether this notification should be shown in the shade.
      */
     boolean showInShade() {
-        if (mEntry == null) return false;
-        return !shouldSuppressNotification() || !mEntry.isClearable();
+        return !shouldSuppressNotification() || !mIsClearable;
     }
 
     /**
      * Sets whether this notification should be suppressed in the shade.
      */
     void setSuppressNotification(boolean suppressNotification) {
-        if (mEntry == null) return;
         boolean prevShowInShade = showInShade();
-        Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
-        int flags = data.getFlags();
         if (suppressNotification) {
-            flags |= Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
+            mFlags |= Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
         } else {
-            flags &= ~Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
+            mFlags &= ~Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION;
         }
-        data.setFlags(flags);
 
         if (showInShade() != prevShowInShade && mSuppressionListener != null) {
             mSuppressionListener.onBubbleNotificationSuppressionChange(this);
@@ -424,9 +462,8 @@
      */
     @Override
     public boolean showDot() {
-        if (mEntry == null) return false;
         return mShowBubbleUpdateDot
-                && !mEntry.shouldSuppressNotificationDot()
+                && !mShouldSuppressNotificationDot
                 && !shouldSuppressNotification();
     }
 
@@ -434,10 +471,9 @@
      * Whether the flyout for the bubble should be shown.
      */
     boolean showFlyout() {
-        if (mEntry == null) return false;
-        return !mSuppressFlyout && !mEntry.shouldSuppressPeek()
+        return !mSuppressFlyout && !mShouldSuppressPeek
                 && !shouldSuppressNotification()
-                && !mEntry.shouldSuppressNotificationList();
+                && !mShouldSuppressNotificationList;
     }
 
     /**
@@ -480,25 +516,14 @@
         }
     }
 
-    /**
-     * Whether shortcut information should be used to populate the bubble.
-     * <p>
-     * To populate the activity use {@link LauncherApps#startShortcut(ShortcutInfo, Rect, Bundle)}.
-     * To populate the icon use {@link LauncherApps#getShortcutIconDrawable(ShortcutInfo, int)}.
-     */
-    boolean usingShortcutInfo() {
-        return mEntry != null && mEntry.getBubbleMetadata().getShortcutId() != null
-                || mShortcutInfo != null;
+    @Nullable
+    PendingIntent getBubbleIntent() {
+        return mIntent;
     }
 
     @Nullable
-    PendingIntent getBubbleIntent() {
-        if (mEntry == null) return null;
-        Notification.BubbleMetadata data = mEntry.getBubbleMetadata();
-        if (data != null) {
-            return data.getIntent();
-        }
-        return null;
+    PendingIntent getDeleteIntent() {
+        return mDeleteIntent;
     }
 
     Intent getSettingsIntent(final Context context) {
@@ -514,8 +539,12 @@
         return intent;
     }
 
+    public int getAppUid() {
+        return mAppUid;
+    }
+
     private int getUid(final Context context) {
-        if (mEntry != null) return mEntry.getSbn().getUid();
+        if (mAppUid != -1) return mAppUid;
         final PackageManager pm = context.getPackageManager();
         if (pm == null) return -1;
         try {
@@ -548,24 +577,27 @@
     }
 
     private boolean shouldSuppressNotification() {
-        if (mEntry == null) return true;
-        return mEntry.getBubbleMetadata() != null
-                && mEntry.getBubbleMetadata().isNotificationSuppressed();
+        return isEnabled(Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION);
     }
 
-    boolean shouldAutoExpand() {
-        if (mEntry == null) return mShouldAutoExpand;
-        Notification.BubbleMetadata metadata = mEntry.getBubbleMetadata();
-        return (metadata != null && metadata.getAutoExpandBubble()) ||  mShouldAutoExpand;
+    public boolean shouldAutoExpand() {
+        return isEnabled(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
     }
 
     void setShouldAutoExpand(boolean shouldAutoExpand) {
-        mShouldAutoExpand = shouldAutoExpand;
+        if (shouldAutoExpand) {
+            enable(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
+        } else {
+            disable(Notification.BubbleMetadata.FLAG_AUTO_EXPAND_BUBBLE);
+        }
+    }
+
+    public void setIsBubble(final boolean isBubble) {
+        mIsBubble = isBubble;
     }
 
     public boolean isBubble() {
-        if (mEntry == null) return (mFlags & FLAG_BUBBLE) != 0;
-        return (mEntry.getSbn().getNotification().flags & FLAG_BUBBLE) != 0;
+        return mIsBubble;
     }
 
     public void enable(int option) {
@@ -576,6 +608,10 @@
         mFlags &= ~option;
     }
 
+    public boolean isEnabled(int option) {
+        return (mFlags & option) != 0;
+    }
+
     @Override
     public String toString() {
         return "Bubble{" + mKey + '}';
@@ -610,34 +646,24 @@
 
     @Override
     public void logUIEvent(int bubbleCount, int action, float normalX, float normalY, int index) {
-        if (this.getEntry() == null
-                || this.getEntry().getSbn() == null) {
-            SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
-                    null /* package name */,
-                    null /* notification channel */,
-                    0 /* notification ID */,
-                    0 /* bubble position */,
-                    bubbleCount,
-                    action,
-                    normalX,
-                    normalY,
-                    false /* unread bubble */,
-                    false /* on-going bubble */,
-                    false /* isAppForeground (unused) */);
-        } else {
-            StatusBarNotification notification = this.getEntry().getSbn();
-            SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
-                    notification.getPackageName(),
-                    notification.getNotification().getChannelId(),
-                    notification.getId(),
-                    index,
-                    bubbleCount,
-                    action,
-                    normalX,
-                    normalY,
-                    this.showInShade(),
-                    false /* isOngoing (unused) */,
-                    false /* isAppForeground (unused) */);
-        }
+        SysUiStatsLog.write(SysUiStatsLog.BUBBLE_UI_CHANGED,
+                mPackageName,
+                mChannelId,
+                mNotificationId,
+                index,
+                bubbleCount,
+                action,
+                normalX,
+                normalY,
+                showInShade(),
+                false /* isOngoing (unused) */,
+                false /* isAppForeground (unused) */);
+    }
+
+    @Nullable
+    private static String getTitle(@NonNull final NotificationEntry e) {
+        final CharSequence titleCharSeq = e.getSbn().getNotification().extras.getCharSequence(
+                Notification.EXTRA_TITLE);
+        return titleCharSeq == null ? null : titleCharSeq.toString();
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index c4c5da4..b2c5402 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -505,8 +505,7 @@
         addNotifCallback(new NotifCallback() {
             @Override
             public void removeNotification(NotificationEntry entry, int reason) {
-                mNotificationEntryManager.performRemoveNotification(entry.getSbn(),
-                        reason);
+                mNotificationEntryManager.performRemoveNotification(entry.getSbn(), reason);
             }
 
             @Override
@@ -637,8 +636,13 @@
                 mStackView.setExpandListener(mExpandListener);
             }
 
-            mStackView.setUnbubbleConversationCallback(notificationEntry ->
-                    onUserChangedBubble(notificationEntry, false /* shouldBubble */));
+            mStackView.setUnbubbleConversationCallback(key -> {
+                final NotificationEntry entry =
+                        mNotificationEntryManager.getPendingOrActiveNotif(key);
+                if (entry != null) {
+                    onUserChangedBubble(entry, false /* shouldBubble */);
+                }
+            });
         }
 
         addToWindowManagerMaybe();
@@ -1024,10 +1028,7 @@
      * @param entry the notification to change bubble state for.
      * @param shouldBubble whether the notification should show as a bubble or not.
      */
-    public void onUserChangedBubble(@Nullable final NotificationEntry entry, boolean shouldBubble) {
-        if (entry == null) {
-            return;
-        }
+    public void onUserChangedBubble(@NonNull final NotificationEntry entry, boolean shouldBubble) {
         NotificationChannel channel = entry.getChannel();
         final String appPkg = entry.getSbn().getPackageName();
         final int appUid = entry.getSbn().getUid();
@@ -1103,7 +1104,8 @@
             mBubbleData.removeSuppressedSummary(groupKey);
 
             // Remove any associated bubble children with the summary
-            final List<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(groupKey);
+            final List<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(
+                    groupKey, mNotificationEntryManager);
             for (int i = 0; i < bubbleChildren.size(); i++) {
                 removeBubble(bubbleChildren.get(i).getKey(), DISMISS_GROUP_CANCELLED);
             }
@@ -1161,21 +1163,18 @@
 
     private void setIsBubble(@NonNull final Bubble b, final boolean isBubble) {
         Objects.requireNonNull(b);
-        if (isBubble) {
-            b.enable(FLAG_BUBBLE);
-        } else {
-            b.disable(FLAG_BUBBLE);
-        }
-        if (b.getEntry() != null) {
+        b.setIsBubble(isBubble);
+        final NotificationEntry entry = mNotificationEntryManager
+                .getPendingOrActiveNotif(b.getKey());
+        if (entry != null) {
             // Updating the entry to be a bubble will trigger our normal update flow
-            setIsBubble(b.getEntry(), isBubble, b.shouldAutoExpand());
+            setIsBubble(entry, isBubble, b.shouldAutoExpand());
         } else if (isBubble) {
-            // If we have no entry to update, it's a persisted bubble so
-            // we need to add it to the stack ourselves
+            // If bubble doesn't exist, it's a persisted bubble so we need to add it to the
+            // stack ourselves
             Bubble bubble = mBubbleData.getOrCreateBubble(null, b /* persistedBubble */);
             inflateAndAdd(bubble, bubble.shouldAutoExpand() /* suppressFlyout */,
                     !bubble.shouldAutoExpand() /* showInShade */);
-
         }
     }
 
@@ -1214,6 +1213,8 @@
                 if (reason == DISMISS_NOTIF_CANCEL) {
                     bubblesToBeRemovedFromRepository.add(bubble);
                 }
+                final NotificationEntry entry = mNotificationEntryManager.getPendingOrActiveNotif(
+                        bubble.getKey());
                 if (!mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
                     if (!mBubbleData.hasOverflowBubbleWithKey(bubble.getKey())
                         && (!bubble.showInShade()
@@ -1222,26 +1223,27 @@
                         // The bubble is now gone & the notification is hidden from the shade, so
                         // time to actually remove it
                         for (NotifCallback cb : mCallbacks) {
-                            if (bubble.getEntry() != null) {
-                                cb.removeNotification(bubble.getEntry(), REASON_CANCEL);
+                            if (entry != null) {
+                                cb.removeNotification(entry, REASON_CANCEL);
                             }
                         }
                     } else {
                         if (bubble.isBubble()) {
                             setIsBubble(bubble, false /* isBubble */);
                         }
-                        if (bubble.getEntry() != null && bubble.getEntry().getRow() != null) {
-                            bubble.getEntry().getRow().updateBubbleButton();
+                        if (entry != null && entry.getRow() != null) {
+                            entry.getRow().updateBubbleButton();
                         }
                     }
 
                 }
-                if (bubble.getEntry() != null) {
-                    final String groupKey = bubble.getEntry().getSbn().getGroupKey();
-                    if (mBubbleData.getBubblesInGroup(groupKey).isEmpty()) {
+                if (entry != null) {
+                    final String groupKey = entry.getSbn().getGroupKey();
+                    if (mBubbleData.getBubblesInGroup(
+                            groupKey, mNotificationEntryManager).isEmpty()) {
                         // Time to potentially remove the summary
                         for (NotifCallback cb : mCallbacks) {
-                            cb.maybeCancelSummary(bubble.getEntry());
+                            cb.maybeCancelSummary(entry);
                         }
                     }
                 }
@@ -1266,9 +1268,12 @@
 
             if (update.selectionChanged && mStackView != null) {
                 mStackView.setSelectedBubble(update.selectedBubble);
-                if (update.selectedBubble != null && update.selectedBubble.getEntry() != null) {
-                    mNotificationGroupManager.updateSuppression(
-                            update.selectedBubble.getEntry());
+                if (update.selectedBubble != null) {
+                    final NotificationEntry entry = mNotificationEntryManager
+                            .getPendingOrActiveNotif(update.selectedBubble.getKey());
+                    if (entry != null) {
+                        mNotificationGroupManager.updateSuppression(entry);
+                    }
                 }
             }
 
@@ -1341,7 +1346,8 @@
         }
 
         String groupKey = entry.getSbn().getGroupKey();
-        ArrayList<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(groupKey);
+        ArrayList<Bubble> bubbleChildren = mBubbleData.getBubblesInGroup(
+                groupKey, mNotificationEntryManager);
         boolean isSuppressedSummary = (mBubbleData.isSummarySuppressed(groupKey)
                 && mBubbleData.getSummaryKey(groupKey).equals(entry.getKey()));
         boolean isSummary = entry.getSbn().getNotification().isGroupSummary();
@@ -1361,9 +1367,15 @@
                     // As far as group manager is concerned, once a child is no longer shown
                     // in the shade, it is essentially removed.
                     Bubble bubbleChild = mBubbleData.getAnyBubbleWithkey(child.getKey());
-                    mNotificationGroupManager.onEntryRemoved(bubbleChild.getEntry());
-                    bubbleChild.setSuppressNotification(true);
-                    bubbleChild.setShowDot(false /* show */);
+                    if (bubbleChild != null) {
+                        final NotificationEntry entry = mNotificationEntryManager
+                                .getPendingOrActiveNotif(bubbleChild.getKey());
+                        if (entry != null) {
+                            mNotificationGroupManager.onEntryRemoved(entry);
+                        }
+                        bubbleChild.setSuppressNotification(true);
+                        bubbleChild.setShowDot(false /* show */);
+                    }
                 } else {
                     // non-bubbled children can be removed
                     for (NotifCallback cb : mCallbacks) {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
index 20a9a8c..c870612 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -22,7 +22,6 @@
 import static com.android.systemui.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
 
 import android.annotation.NonNull;
-import android.app.Notification;
 import android.app.PendingIntent;
 import android.content.Context;
 import android.util.Log;
@@ -34,6 +33,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.R;
 import com.android.systemui.bubbles.BubbleController.DismissReason;
+import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
 import java.io.FileDescriptor;
@@ -256,8 +256,7 @@
         }
         mPendingBubbles.remove(bubble.getKey()); // No longer pending once we're here
         Bubble prevBubble = getBubbleInStackWithKey(bubble.getKey());
-        suppressFlyout |= bubble.getEntry() == null
-                || !bubble.getEntry().getRanking().visuallyInterruptive();
+        suppressFlyout |= !bubble.isVisuallyInterruptive();
 
         if (prevBubble == null) {
             // Create a new bubble
@@ -335,13 +334,15 @@
      * Retrieves any bubbles that are part of the notification group represented by the provided
      * group key.
      */
-    ArrayList<Bubble> getBubblesInGroup(@Nullable String groupKey) {
+    ArrayList<Bubble> getBubblesInGroup(@Nullable String groupKey, @NonNull
+            NotificationEntryManager nem) {
         ArrayList<Bubble> bubbleChildren = new ArrayList<>();
         if (groupKey == null) {
             return bubbleChildren;
         }
         for (Bubble b : mBubbles) {
-            if (b.getEntry() != null && groupKey.equals(b.getEntry().getSbn().getGroupKey())) {
+            final NotificationEntry entry = nem.getPendingOrActiveNotif(b.getKey());
+            if (entry != null && groupKey.equals(entry.getSbn().getGroupKey())) {
                 bubbleChildren.add(b);
             }
         }
@@ -439,9 +440,7 @@
             Bubble newSelected = mBubbles.get(newIndex);
             setSelectedBubbleInternal(newSelected);
         }
-        if (bubbleToRemove.getEntry() != null) {
-            maybeSendDeleteIntent(reason, bubbleToRemove.getEntry());
-        }
+        maybeSendDeleteIntent(reason, bubbleToRemove);
     }
 
     void overflowBubble(@DismissReason int reason, Bubble bubble) {
@@ -611,21 +610,14 @@
         return true;
     }
 
-    private void maybeSendDeleteIntent(@DismissReason int reason,
-            @NonNull final NotificationEntry entry) {
-        if (reason == BubbleController.DISMISS_USER_GESTURE) {
-            Notification.BubbleMetadata bubbleMetadata = entry.getBubbleMetadata();
-            PendingIntent deleteIntent = bubbleMetadata != null
-                    ? bubbleMetadata.getDeleteIntent()
-                    : null;
-            if (deleteIntent != null) {
-                try {
-                    deleteIntent.send();
-                } catch (PendingIntent.CanceledException e) {
-                    Log.w(TAG, "Failed to send delete intent for bubble with key: "
-                            + entry.getKey());
-                }
-            }
+    private void maybeSendDeleteIntent(@DismissReason int reason, @NonNull final Bubble bubble) {
+        if (reason != BubbleController.DISMISS_USER_GESTURE) return;
+        PendingIntent deleteIntent = bubble.getDeleteIntent();
+        if (deleteIntent == null) return;
+        try {
+            deleteIntent.send();
+        } catch (PendingIntent.CanceledException e) {
+            Log.w(TAG, "Failed to send delete intent for bubble with key: " + bubble.getKey());
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
index d20f405..0c25d14 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDataRepository.kt
@@ -74,11 +74,15 @@
 
     private fun transform(userId: Int, bubbles: List<Bubble>): List<BubbleEntity> {
         return bubbles.mapNotNull { b ->
-            var shortcutId = b.shortcutInfo?.id
-            if (shortcutId == null) shortcutId = b.entry?.bubbleMetadata?.shortcutId
-            if (shortcutId == null) return@mapNotNull null
-            BubbleEntity(userId, b.packageName, shortcutId, b.key, b.rawDesiredHeight,
-                    b.rawDesiredHeightResId)
+            BubbleEntity(
+                    userId,
+                    b.packageName,
+                    b.shortcutInfo?.id ?: return@mapNotNull null,
+                    b.key,
+                    b.rawDesiredHeight,
+                    b.rawDesiredHeightResId,
+                    b.title
+            )
         }
     }
 
@@ -159,8 +163,13 @@
         val bubbles = entities.mapNotNull { entity ->
             shortcutMap[ShortcutKey(entity.userId, entity.packageName)]
                     ?.first { shortcutInfo -> entity.shortcutId == shortcutInfo.id }
-                    ?.let { shortcutInfo -> Bubble(entity.key, shortcutInfo, entity.desiredHeight,
-                            entity.desiredHeightResId) }
+                    ?.let { shortcutInfo -> Bubble(
+                            entity.key,
+                            shortcutInfo,
+                            entity.desiredHeight,
+                            entity.desiredHeightResId,
+                            entity.title
+                    ) }
         }
         uiScope.launch { cb(bubbles) }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
index 6dcc9dc..471a769 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
@@ -169,7 +169,7 @@
                             return;
                         }
                         try {
-                            if (!mIsOverflow && mBubble.usingShortcutInfo()) {
+                            if (!mIsOverflow && mBubble.getShortcutInfo() != null) {
                                 options.setApplyActivityFlagsForBubbles(true);
                                 mActivityView.startShortcutActivity(mBubble.getShortcutInfo(),
                                         options, null /* sourceBounds */);
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
index 8c76cda..1fa3aaa 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleFlyoutView.java
@@ -31,6 +31,7 @@
 import android.graphics.Path;
 import android.graphics.PointF;
 import android.graphics.RectF;
+import android.graphics.drawable.Drawable;
 import android.graphics.drawable.ShapeDrawable;
 import android.text.TextUtils;
 import android.view.LayoutInflater;
@@ -223,9 +224,10 @@
             float[] dotCenter,
             boolean hideDot) {
 
-        if (flyoutMessage.senderAvatar != null && flyoutMessage.isGroupChat) {
+        final Drawable senderAvatar = flyoutMessage.senderAvatar;
+        if (senderAvatar != null && flyoutMessage.isGroupChat) {
             mSenderAvatar.setVisibility(VISIBLE);
-            mSenderAvatar.setImageDrawable(flyoutMessage.senderAvatar);
+            mSenderAvatar.setImageDrawable(senderAvatar);
         } else {
             mSenderAvatar.setVisibility(GONE);
             mSenderAvatar.setTranslationX(0);
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
index 74231c6..a799f2d 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleIconFactory.java
@@ -15,7 +15,8 @@
  */
 package com.android.systemui.bubbles;
 
-import android.app.Notification;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherApps;
@@ -50,15 +51,14 @@
     /**
      * Returns the drawable that the developer has provided to display in the bubble.
      */
-    Drawable getBubbleDrawable(Context context, ShortcutInfo shortcutInfo,
-            Notification.BubbleMetadata metadata) {
+    Drawable getBubbleDrawable(@NonNull final Context context,
+            @Nullable final ShortcutInfo shortcutInfo, @Nullable final Icon ic) {
         if (shortcutInfo != null) {
             LauncherApps launcherApps =
                     (LauncherApps) context.getSystemService(Context.LAUNCHER_APPS_SERVICE);
             int density = context.getResources().getConfiguration().densityDpi;
             return launcherApps.getShortcutIconDrawable(shortcutInfo, density);
         } else {
-            Icon ic = metadata.getIcon();
             if (ic != null) {
                 if (ic.getType() == Icon.TYPE_URI
                         || ic.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
index c5faae0..c1dd8c3 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleLoggerImpl.java
@@ -16,8 +16,6 @@
 
 package com.android.systemui.bubbles;
 
-import android.service.notification.StatusBarNotification;
-
 import com.android.internal.logging.UiEventLoggerImpl;
 
 /**
@@ -32,12 +30,11 @@
      * @param e UI event
      */
     public void log(Bubble b, UiEventEnum e) {
-        if (b.getEntry() == null) {
+        if (b.getInstanceId() == null) {
             // Added from persistence -- TODO log this with specific event?
             return;
         }
-        StatusBarNotification sbn = b.getEntry().getSbn();
-        logWithInstanceId(e, sbn.getUid(), sbn.getPackageName(), sbn.getInstanceId());
+        logWithInstanceId(e, b.getAppUid(), b.getPackageName(), b.getInstanceId());
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
index b4672c1..ea324af 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflowActivity.java
@@ -300,9 +300,7 @@
                 });
 
         // If the bubble was persisted, the entry is null but it should have shortcut info
-        ShortcutInfo info = b.getEntry() == null
-                ? b.getShortcutInfo()
-                : b.getEntry().getRanking().getShortcutInfo();
+        ShortcutInfo info = b.getShortcutInfo();
         if (info == null) {
             Log.d(TAG, "ShortcutInfo required to bubble but none found for " + b);
         } else {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index d5fe9b2..a59d2c0 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -91,7 +91,6 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.SysUiStatsLog;
-import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.phone.CollapsedStatusBarFragment;
 import com.android.systemui.util.DismissCircleView;
 import com.android.systemui.util.FloatingContentCoordinator;
@@ -287,7 +286,7 @@
     private BubbleController.BubbleExpandListener mExpandListener;
 
     /** Callback to run when we want to unbubble the given notification's conversation. */
-    private Consumer<NotificationEntry> mUnbubbleConversationCallback;
+    private Consumer<String> mUnbubbleConversationCallback;
 
     private SysUiState mSysUiState;
 
@@ -997,10 +996,7 @@
         mManageMenu.findViewById(R.id.bubble_manage_menu_dont_bubble_container).setOnClickListener(
                 view -> {
                     showManageMenu(false /* show */);
-                    final Bubble bubble = mBubbleData.getSelectedBubble();
-                    if (bubble != null && mBubbleData.hasBubbleInStackWithKey(bubble.getKey())) {
-                        mUnbubbleConversationCallback.accept(bubble.getEntry());
-                    }
+                    mUnbubbleConversationCallback.accept(mBubbleData.getSelectedBubble().getKey());
                 });
 
         mManageMenu.findViewById(R.id.bubble_manage_menu_settings_container).setOnClickListener(
@@ -1348,7 +1344,7 @@
 
     /** Sets the function to call to un-bubble the given conversation. */
     public void setUnbubbleConversationCallback(
-            Consumer<NotificationEntry> unbubbleConversationCallback) {
+            Consumer<String> unbubbleConversationCallback) {
         mUnbubbleConversationCallback = unbubbleConversationCallback;
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
index 525d5b5..3e4ff52 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewInfoTask.java
@@ -37,8 +37,6 @@
 import android.graphics.drawable.Icon;
 import android.os.AsyncTask;
 import android.os.Parcelable;
-import android.os.UserHandle;
-import android.service.notification.StatusBarNotification;
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.PathParser;
@@ -53,6 +51,7 @@
 
 import java.lang.ref.WeakReference;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Simple task to inflate views & load necessary info to display a bubble.
@@ -129,35 +128,10 @@
         @Nullable
         static BubbleViewInfo populate(Context c, BubbleStackView stackView,
                 BubbleIconFactory iconFactory, Bubble b, boolean skipInflation) {
-            final NotificationEntry entry = b.getEntry();
-            if (entry == null) {
-                // populate from ShortcutInfo when NotificationEntry is not available
-                final ShortcutInfo s = b.getShortcutInfo();
-                return populate(c, stackView, iconFactory, skipInflation || b.isInflated(),
-                        s.getPackage(), s.getUserHandle(), s, null);
-            }
-            final StatusBarNotification sbn = entry.getSbn();
-            final String bubbleShortcutId =  entry.getBubbleMetadata().getShortcutId();
-            final ShortcutInfo si = bubbleShortcutId == null
-                    ? null : entry.getRanking().getShortcutInfo();
-            return populate(
-                    c, stackView, iconFactory, skipInflation || b.isInflated(),
-                    sbn.getPackageName(), sbn.getUser(), si, entry);
-        }
-
-        private static BubbleViewInfo populate(
-                @NonNull final Context c,
-                @NonNull final BubbleStackView stackView,
-                @NonNull final BubbleIconFactory iconFactory,
-                final boolean isInflated,
-                @NonNull final String packageName,
-                @NonNull final UserHandle user,
-                @Nullable final ShortcutInfo shortcutInfo,
-                @Nullable final NotificationEntry entry) {
             BubbleViewInfo info = new BubbleViewInfo();
 
             // View inflation: only should do this once per bubble
-            if (!isInflated) {
+            if (!skipInflation && !b.isInflated()) {
                 LayoutInflater inflater = LayoutInflater.from(c);
                 info.imageView = (BadgedImageView) inflater.inflate(
                         R.layout.bubble_view, stackView, false /* attachToRoot */);
@@ -167,8 +141,8 @@
                 info.expandedView.setStackView(stackView);
             }
 
-            if (shortcutInfo != null) {
-                info.shortcutInfo = shortcutInfo;
+            if (b.getShortcutInfo() != null) {
+                info.shortcutInfo = b.getShortcutInfo();
             }
 
             // App name & app icon
@@ -178,7 +152,7 @@
             Drawable appIcon;
             try {
                 appInfo = pm.getApplicationInfo(
-                        packageName,
+                        b.getPackageName(),
                         PackageManager.MATCH_UNINSTALLED_PACKAGES
                                 | PackageManager.MATCH_DISABLED_COMPONENTS
                                 | PackageManager.MATCH_DIRECT_BOOT_UNAWARE
@@ -186,17 +160,17 @@
                 if (appInfo != null) {
                     info.appName = String.valueOf(pm.getApplicationLabel(appInfo));
                 }
-                appIcon = pm.getApplicationIcon(packageName);
-                badgedIcon = pm.getUserBadgedIcon(appIcon, user);
+                appIcon = pm.getApplicationIcon(b.getPackageName());
+                badgedIcon = pm.getUserBadgedIcon(appIcon, b.getUser());
             } catch (PackageManager.NameNotFoundException exception) {
                 // If we can't find package... don't think we should show the bubble.
-                Log.w(TAG, "Unable to find package: " + packageName);
+                Log.w(TAG, "Unable to find package: " + b.getPackageName());
                 return null;
             }
 
             // Badged bubble image
             Drawable bubbleDrawable = iconFactory.getBubbleDrawable(c, info.shortcutInfo,
-                    entry == null ? null : entry.getBubbleMetadata());
+                    b.getIcon());
             if (bubbleDrawable == null) {
                 // Default to app icon
                 bubbleDrawable = appIcon;
@@ -222,8 +196,10 @@
                     Color.WHITE, WHITE_SCRIM_ALPHA);
 
             // Flyout
-            if (entry != null) {
-                info.flyoutMessage = extractFlyoutMessage(c, entry);
+            info.flyoutMessage = b.getFlyoutMessage();
+            if (info.flyoutMessage != null) {
+                info.flyoutMessage.senderAvatar =
+                        loadSenderAvatar(c, info.flyoutMessage.senderIcon);
             }
             return info;
         }
@@ -235,8 +211,8 @@
      * notification, based on its type. Returns null if there should not be an update message.
      */
     @NonNull
-    static Bubble.FlyoutMessage extractFlyoutMessage(Context context,
-            NotificationEntry entry) {
+    static Bubble.FlyoutMessage extractFlyoutMessage(NotificationEntry entry) {
+        Objects.requireNonNull(entry);
         final Notification underlyingNotif = entry.getSbn().getNotification();
         final Class<? extends Notification.Style> style = underlyingNotif.getNotificationStyle();
 
@@ -264,20 +240,9 @@
                 if (latestMessage != null) {
                     bubbleMessage.message = latestMessage.getText();
                     Person sender = latestMessage.getSenderPerson();
-                    bubbleMessage.senderName = sender != null
-                            ? sender.getName()
-                            : null;
-
+                    bubbleMessage.senderName = sender != null ? sender.getName() : null;
                     bubbleMessage.senderAvatar = null;
-                    if (sender != null && sender.getIcon() != null) {
-                        if (sender.getIcon().getType() == Icon.TYPE_URI
-                                || sender.getIcon().getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
-                            context.grantUriPermission(context.getPackageName(),
-                                    sender.getIcon().getUri(),
-                                    Intent.FLAG_GRANT_READ_URI_PERMISSION);
-                        }
-                        bubbleMessage.senderAvatar = sender.getIcon().loadDrawable(context);
-                    }
+                    bubbleMessage.senderIcon = sender != null ? sender.getIcon() : null;
                     return bubbleMessage;
                 }
             } else if (Notification.InboxStyle.class.equals(style)) {
@@ -306,4 +271,15 @@
 
         return bubbleMessage;
     }
+
+    @Nullable
+    static Drawable loadSenderAvatar(@NonNull final Context context, @Nullable final Icon icon) {
+        Objects.requireNonNull(context);
+        if (icon == null) return null;
+        if (icon.getType() == Icon.TYPE_URI || icon.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
+            context.grantUriPermission(context.getPackageName(),
+                    icon.getUri(), Intent.FLAG_GRANT_READ_URI_PERMISSION);
+        }
+        return icon.loadDrawable(context);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
index 355c4b1..24768cd 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleEntity.kt
@@ -24,5 +24,6 @@
     val shortcutId: String,
     val key: String,
     val desiredHeight: Int,
-    @DimenRes val desiredHeightResId: Int
+    @DimenRes val desiredHeightResId: Int,
+    val title: String? = null
 )
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
index a8faf25..66fff338 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/storage/BubbleXmlHelper.kt
@@ -33,6 +33,7 @@
 private const val ATTR_KEY = "key"
 private const val ATTR_DESIRED_HEIGHT = "h"
 private const val ATTR_DESIRED_HEIGHT_RES_ID = "hid"
+private const val ATTR_TITLE = "t"
 
 /**
  * Writes the bubbles in xml format into given output stream.
@@ -63,6 +64,7 @@
         serializer.attribute(null, ATTR_KEY, bubble.key)
         serializer.attribute(null, ATTR_DESIRED_HEIGHT, bubble.desiredHeight.toString())
         serializer.attribute(null, ATTR_DESIRED_HEIGHT_RES_ID, bubble.desiredHeightResId.toString())
+        bubble.title?.let { serializer.attribute(null, ATTR_TITLE, it) }
         serializer.endTag(null, TAG_BUBBLE)
     } catch (e: IOException) {
         throw RuntimeException(e)
@@ -92,7 +94,8 @@
             parser.getAttributeWithName(ATTR_SHORTCUT_ID) ?: return null,
             parser.getAttributeWithName(ATTR_KEY) ?: return null,
             parser.getAttributeWithName(ATTR_DESIRED_HEIGHT)?.toInt() ?: return null,
-            parser.getAttributeWithName(ATTR_DESIRED_HEIGHT_RES_ID)?.toInt() ?: return null
+            parser.getAttributeWithName(ATTR_DESIRED_HEIGHT_RES_ID)?.toInt() ?: return null,
+            parser.getAttributeWithName(ATTR_TITLE)
     )
 }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
index 59f8c4e..36398a6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -304,6 +304,10 @@
     public void testPromoteBubble_autoExpand() throws Exception {
         mBubbleController.updateBubble(mRow2.getEntry());
         mBubbleController.updateBubble(mRow.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
+                .thenReturn(mRow.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
+                .thenReturn(mRow2.getEntry());
         mBubbleController.removeBubble(
                 mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
@@ -331,6 +335,10 @@
         mBubbleController.updateBubble(mRow2.getEntry());
         mBubbleController.updateBubble(mRow.getEntry(), /* suppressFlyout */
                 false, /* showInShade */ true);
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
+                .thenReturn(mRow.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
+                .thenReturn(mRow2.getEntry());
         mBubbleController.removeBubble(
                 mRow.getEntry().getKey(), BubbleController.DISMISS_USER_GESTURE);
 
@@ -433,15 +441,16 @@
         assertTrue(mSysUiStateBubblesExpanded);
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry(), mBubbleData.getSelectedBubble().getEntry());
+        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getSelectedBubble().getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Switch which bubble is expanded
-        mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
+        mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(
+                mRow.getEntry().getKey()));
         mBubbleData.setExpanded(true);
-        assertEquals(mRow.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow.getEntry()));
 
@@ -543,27 +552,27 @@
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
-                        .getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(
+                        stackView.getExpandedBubble().getKey()).getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
         // Make sure first bubble is selected
-        assertEquals(mRow.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getEntry().getKey());
 
         // Dismiss that one
         mBubbleController.removeBubble(
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey())
-                        .getEntry().getKey(),
+                mBubbleData.getBubbleInStackWithKey(
+                        stackView.getExpandedBubble().getKey()).getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -839,6 +848,12 @@
                 mRow2.getEntry(), /* suppressFlyout */ false, /* showInShade */ false);
         mBubbleController.updateBubble(
                 mRow3.getEntry(), /* suppressFlyout */ false, /* showInShade */ false);
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
+                .thenReturn(mRow.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow2.getEntry().getKey()))
+                .thenReturn(mRow2.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow3.getEntry().getKey()))
+                .thenReturn(mRow3.getEntry());
         assertEquals(mBubbleData.getBubbles().size(), 3);
 
         mBubbleData.setMaxOverflowBubbles(1);
@@ -908,6 +923,8 @@
         // GIVEN a group summary with a bubble child
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
@@ -927,6 +944,8 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -948,6 +967,8 @@
         // GIVEN a group summary with two (non-bubble) children and one bubble child
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(2);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         mEntryListener.onPendingEntryAdded(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
index 72f816f..be03923 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleTest.java
@@ -86,8 +86,7 @@
         final String msg = "Hello there!";
         doReturn(Notification.Style.class).when(mNotif).getNotificationStyle();
         mExtras.putCharSequence(Notification.EXTRA_TEXT, msg);
-        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mContext,
-                mEntry).message);
+        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
     }
 
     @Test
@@ -98,8 +97,7 @@
         mExtras.putCharSequence(Notification.EXTRA_BIG_TEXT, msg);
 
         // Should be big text, not the small text.
-        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mContext,
-                mEntry).message);
+        assertEquals(msg, BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
     }
 
     @Test
@@ -107,8 +105,7 @@
         doReturn(Notification.MediaStyle.class).when(mNotif).getNotificationStyle();
 
         // Media notifs don't get update messages.
-        assertNull(BubbleViewInfoTask.extractFlyoutMessage(mContext,
-                mEntry).message);
+        assertNull(BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
     }
 
     @Test
@@ -124,7 +121,7 @@
 
         // Should be the last one only.
         assertEquals("Really? I prefer them that way.",
-                BubbleViewInfoTask.extractFlyoutMessage(mContext, mEntry).message);
+                BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
     }
 
     @Test
@@ -139,11 +136,8 @@
                                 "Oh, hello!", 0, "Mady").toBundle()});
 
         // Should be the last one only.
-        assertEquals("Oh, hello!",
-                BubbleViewInfoTask.extractFlyoutMessage(mContext, mEntry).message);
-        assertEquals("Mady",
-                BubbleViewInfoTask.extractFlyoutMessage(mContext,
-                        mEntry).senderName);
+        assertEquals("Oh, hello!", BubbleViewInfoTask.extractFlyoutMessage(mEntry).message);
+        assertEquals("Mady", BubbleViewInfoTask.extractFlyoutMessage(mEntry).senderName);
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
index 58e06b5..1c70db3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/NewNotifPipelineBubbleControllerTest.java
@@ -302,6 +302,8 @@
     public void testRemoveBubble_withDismissedNotif_notInOverflow() {
         mEntryListener.onEntryAdded(mRow.getEntry());
         mBubbleController.updateBubble(mRow.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(mRow.getEntry().getKey()))
+                .thenReturn(mRow.getEntry());
 
         assertTrue(mBubbleController.hasBubbles());
         assertFalse(mBubbleController.isBubbleNotificationSuppressedFromShade(mRow.getEntry()));
@@ -388,14 +390,14 @@
                 true, mRow.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry(), mBubbleData.getSelectedBubble().getEntry());
+        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getSelectedBubble().getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(mRow2.getEntry()));
 
         // Switch which bubble is expanded
         mBubbleData.setSelectedBubble(mBubbleData.getBubbleInStackWithKey(mRow.getEntry().getKey()));
         mBubbleData.setExpanded(true);
-        assertEquals(mRow.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow.getEntry()));
 
@@ -488,27 +490,27 @@
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow2.getEntry().getKey());
 
         // Last added is the one that is expanded
-        assertEquals(mRow2.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow2.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         assertTrue(mBubbleController.isBubbleNotificationSuppressedFromShade(
                 mRow2.getEntry()));
 
         // Dismiss currently expanded
         mBubbleController.removeBubble(
                 mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
+                        stackView.getExpandedBubble().getKey()).getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
         verify(mBubbleExpandListener).onBubbleExpandChanged(false, mRow2.getEntry().getKey());
 
         // Make sure first bubble is selected
-        assertEquals(mRow.getEntry(),
-                mBubbleData.getBubbleInStackWithKey(stackView.getExpandedBubble().getKey()).getEntry());
+        assertEquals(mRow.getEntry().getKey(), mBubbleData.getBubbleInStackWithKey(
+                stackView.getExpandedBubble().getKey()).getKey());
         verify(mBubbleExpandListener).onBubbleExpandChanged(true, mRow.getEntry().getKey());
 
         // Dismiss that one
         mBubbleController.removeBubble(
                 mBubbleData.getBubbleInStackWithKey(
-                        stackView.getExpandedBubble().getKey()).getEntry().getKey(),
+                        stackView.getExpandedBubble().getKey()).getKey(),
                 BubbleController.DISMISS_USER_GESTURE);
 
         // Make sure state changes and collapse happens
@@ -767,6 +769,8 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -785,6 +789,8 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(0);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
         assertTrue(mBubbleData.hasBubbleInStackWithKey(groupedBubble.getEntry().getKey()));
 
@@ -807,6 +813,8 @@
         ExpandableNotificationRow groupSummary = mNotificationTestHelper.createGroup(2);
         ExpandableNotificationRow groupedBubble = mNotificationTestHelper.createBubbleInGroup();
         mEntryListener.onEntryAdded(groupedBubble.getEntry());
+        when(mNotificationEntryManager.getPendingOrActiveNotif(groupedBubble.getEntry().getKey()))
+                .thenReturn(groupedBubble.getEntry());
         groupSummary.addChildNotification(groupedBubble);
 
         // WHEN the summary is dismissed
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
index 1d02b8d..9b8fd11 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubblePersistentRepositoryTest.kt
@@ -32,7 +32,7 @@
 
     private val bubbles = listOf(
             BubbleEntity(0, "com.example.messenger", "shortcut-1", "key-1", 120, 0),
-            BubbleEntity(10, "com.example.chat", "alice and bob", "key-2", 0, 16537428),
+            BubbleEntity(10, "com.example.chat", "alice and bob", "key-2", 0, 16537428, "title"),
             BubbleEntity(0, "com.example.messenger", "shortcut-2", "key-3", 120, 0)
     )
     private lateinit var repository: BubblePersistentRepository
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
index f9d611c..76c5833 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleVolatileRepositoryTest.kt
@@ -37,9 +37,10 @@
     private val user0 = UserHandle.of(0)
     private val user10 = UserHandle.of(10)
 
-    private val bubble1 = BubbleEntity(0, PKG_MESSENGER, "shortcut-1", "k1", 120, 0)
-    private val bubble2 = BubbleEntity(10, PKG_CHAT, "alice and bob", "k2", 0, 16537428)
-    private val bubble3 = BubbleEntity(0, PKG_MESSENGER, "shortcut-2", "k3", 120, 0)
+    private val bubble1 = BubbleEntity(0, "com.example.messenger", "shortcut-1", "key-1", 120, 0)
+    private val bubble2 = BubbleEntity(10, "com.example.chat", "alice and bob",
+            "key-2", 0, 16537428, "title")
+    private val bubble3 = BubbleEntity(0, "com.example.messenger", "shortcut-2", "key-3", 120, 0)
 
     private val bubbles = listOf(bubble1, bubble2, bubble3)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
index 4946787..81687c7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/storage/BubbleXmlHelperTest.kt
@@ -31,17 +31,17 @@
 class BubbleXmlHelperTest : SysuiTestCase() {
 
     private val bubbles = listOf(
-        BubbleEntity(0, "com.example.messenger", "shortcut-1", "k1", 120, 0),
-        BubbleEntity(10, "com.example.chat", "alice and bob", "k2", 0, 16537428),
-        BubbleEntity(0, "com.example.messenger", "shortcut-2", "k3", 120, 0)
+            BubbleEntity(0, "com.example.messenger", "shortcut-1", "k1", 120, 0),
+            BubbleEntity(10, "com.example.chat", "alice and bob", "k2", 0, 16537428, "title"),
+            BubbleEntity(0, "com.example.messenger", "shortcut-2", "k3", 120, 0)
     )
 
     @Test
     fun testWriteXml() {
         val expectedEntries = """
-            <bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
-            <bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" />
-            <bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
+<bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
+<bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" t="title" />
+<bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
         """.trimIndent()
         ByteArrayOutputStream().use {
             writeXml(it, bubbles)
@@ -54,12 +54,12 @@
     @Test
     fun testReadXml() {
         val src = """
-            <?xml version='1.0' encoding='utf-8' standalone='yes' ?>
-            <bs>
-            <bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
-            <bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" />
-            <bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
-            </bs>
+<?xml version='1.0' encoding='utf-8' standalone='yes' ?>
+<bs>
+<bb uid="0" pkg="com.example.messenger" sid="shortcut-1" key="k1" h="120" hid="0" />
+<bb uid="10" pkg="com.example.chat" sid="alice and bob" key="k2" h="0" hid="16537428" t="title" />
+<bb uid="0" pkg="com.example.messenger" sid="shortcut-2" key="k3" h="120" hid="0" />
+</bs>
         """.trimIndent()
         val actual = readXml(ByteArrayInputStream(src.toByteArray(Charsets.UTF_8)))
         assertEquals("failed parsing bubbles from xml\n$src", bubbles, actual)