Ensure the notification is removed when bubble is removed & fix cancel all

Bubbles has some particular notification behaviour needs:

1) Whenever there is a bubble, the notif is kept around but may be hidden
   from the shade
2) When the bubble is dismissed, if the notif is no longer in the
   shade => really remove the notification
3) When the bubble is dismissed, if there is still a notification in the
   shade => notif stays around and is normal
4) Clear all should only hide the bubble notification, not remove the
   bubble
5) Apps canceling a notification that has a bubble will cancel the bubble

This CL does this by:

* Including the removal reason removeNotification path
* Adding a NotificationRemoveInterceptor that gets the option of overriding
  the removal
* BubbleController has this new interceptor and uses the removal reason
  to determine what should happen to the bubble & if the notification
  should be allowed to be removed
* When the bubble is dismissed, if the notif is no longer in the shade,
  then actually remove that notification

Test: atest BubbleControllerTest NotificationEntryManagerTest
Bug: 130347307
Bug: 130687293
Change-Id: I4459864a2ee5522076117f84ae37022bdfe4ee5d

13 files changed, 330 insertions, 57 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index d0713636c540..9ecc6f448f17 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -16,6 +16,10 @@
 
 package com.android.systemui.bubbles;
 
+import static android.service.notification.NotificationListenerService.REASON_APP_CANCEL;
+import static android.service.notification.NotificationListenerService.REASON_APP_CANCEL_ALL;
+import static android.service.notification.NotificationListenerService.REASON_CANCEL;
+import static android.service.notification.NotificationListenerService.REASON_CANCEL_ALL;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.INVALID_DISPLAY;
 import static android.view.View.INVISIBLE;
@@ -53,13 +57,13 @@ import androidx.annotation.MainThread;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.statusbar.IStatusBarService;
-import com.android.internal.statusbar.NotificationVisibility;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.WindowManagerWrapper;
+import com.android.systemui.statusbar.NotificationRemoveInterceptor;
 import com.android.systemui.statusbar.notification.NotificationEntryListener;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.NotificationInterruptionStateProvider;
@@ -210,6 +214,7 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
 
         mNotificationEntryManager = Dependency.get(NotificationEntryManager.class);
         mNotificationEntryManager.addNotificationEntryListener(mEntryListener);
+        mNotificationEntryManager.setNotificationRemoveInterceptor(mRemoveInterceptor);
 
         mStatusBarWindowController = statusBarWindowController;
         mStatusBarStateListener = new StatusBarStateListener();
@@ -389,6 +394,46 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
     }
 
     @SuppressWarnings("FieldCanBeLocal")
+    private final NotificationRemoveInterceptor mRemoveInterceptor =
+            new NotificationRemoveInterceptor() {
+            @Override
+            public boolean onNotificationRemoveRequested(String key, int reason) {
+                if (!mBubbleData.hasBubbleWithKey(key)) {
+                    return false;
+                }
+                NotificationEntry entry = mBubbleData.getBubbleWithKey(key).entry;
+
+                final boolean isClearAll = reason == REASON_CANCEL_ALL;
+                final boolean isUserDimiss = reason == REASON_CANCEL;
+                final boolean isAppCancel = reason == REASON_APP_CANCEL
+                        || reason == REASON_APP_CANCEL_ALL;
+
+                // 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)
+                        || isClearAll || isUserDimiss;
+
+                // The bubble notification sticks around in the data as long as the bubble is
+                // not dismissed and the app hasn't cancelled the notification.
+                boolean bubbleExtended = entry.isBubble() && !entry.isBubbleDismissed()
+                        && userRemovedNotif;
+                if (bubbleExtended) {
+                    entry.setShowInShadeWhenBubble(false);
+                    if (mStackView != null) {
+                        mStackView.updateDotVisibility(entry.key);
+                    }
+                    mNotificationEntryManager.updateNotifications();
+                    return true;
+                } else if (!userRemovedNotif && !entry.isBubbleDismissed()) {
+                    // This wasn't a user removal so we should remove the bubble as well
+                    mBubbleData.notificationEntryRemoved(entry, DISMISS_NOTIF_CANCEL);
+                    return false;
+                }
+                return false;
+            }
+        };
+
+    @SuppressWarnings("FieldCanBeLocal")
     private final NotificationEntryListener mEntryListener = new NotificationEntryListener() {
         @Override
         public void onPendingEntryAdded(NotificationEntry entry) {
@@ -396,7 +441,6 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
                 return;
             }
             if (mNotificationInterruptionStateProvider.shouldBubbleUp(entry)) {
-                // TODO: handle group summaries?
                 updateShowInShadeForSuppressNotification(entry);
             }
         }
@@ -426,23 +470,6 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
                 updateBubble(entry);
             }
         }
-
-        @Override
-        public void onEntryRemoved(NotificationEntry entry,
-                @Nullable NotificationVisibility visibility,
-                boolean removedByUser) {
-            if (!areBubblesEnabled(mContext)) {
-                return;
-            }
-            entry.setShowInShadeWhenBubble(false);
-            if (mStackView != null) {
-                mStackView.updateDotVisibility(entry.key);
-            }
-            if (!removedByUser) {
-                // This was a cancel so we should remove the bubble
-                removeBubble(entry.key, DISMISS_NOTIF_CANCEL);
-            }
-        }
     };
 
     @SuppressWarnings("FieldCanBeLocal")
@@ -455,13 +482,15 @@ public class BubbleController implements ConfigurationController.ConfigurationLi
         }
 
         @Override
-        public void onBubbleRemoved(Bubble bubble, int reason) {
+        public void onBubbleRemoved(Bubble bubble, @DismissReason int reason) {
             if (mStackView != null) {
                 mStackView.removeBubble(bubble);
             }
-            if (!bubble.entry.showInShadeWhenBubble()) {
-                // The notification is gone & bubble is gone, time to actually remove it
-                mNotificationEntryManager.performRemoveNotification(bubble.entry.notification);
+            if (!mBubbleData.hasBubbleWithKey(bubble.getKey())
+                    && !bubble.entry.showInShadeWhenBubble()) {
+                // The bubble is gone & the notification is gone, time to actually remove it
+                mNotificationEntryManager.performRemoveNotification(bubble.entry.notification,
+                        0 /* reason */);
             } else {
                 // The notification is still in the shade but we've removed the bubble so
                 // lets make sure NoMan knows it's not a bubble anymore
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
index aeaceb0c11e5..d59a5e8593cc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
@@ -104,7 +104,7 @@ public class NotificationListener extends NotificationListenerWithPlugins {
 
                     // Remove existing notification to avoid stale data.
                     if (isUpdate) {
-                        mEntryManager.removeNotification(key, rankingMap);
+                        mEntryManager.removeNotification(key, rankingMap, 0 /* reason */);
                     } else {
                         mEntryManager.getNotificationData()
                                 .updateRanking(rankingMap);
@@ -121,18 +121,23 @@ public class NotificationListener extends NotificationListenerWithPlugins {
     }
 
     @Override
-    public void onNotificationRemoved(StatusBarNotification sbn,
-            final RankingMap rankingMap) {
-        if (DEBUG) Log.d(TAG, "onNotificationRemoved: " + sbn);
+    public void onNotificationRemoved(StatusBarNotification sbn, RankingMap rankingMap,
+            int reason) {
+        if (DEBUG) Log.d(TAG, "onNotificationRemoved: " + sbn + " reason: " + reason);
         if (sbn != null && !onPluginNotificationRemoved(sbn, rankingMap)) {
             final String key = sbn.getKey();
             Dependency.get(Dependency.MAIN_HANDLER).post(() -> {
-                mEntryManager.removeNotification(key, rankingMap);
+                mEntryManager.removeNotification(key, rankingMap, reason);
             });
         }
     }
 
     @Override
+    public void onNotificationRemoved(StatusBarNotification sbn, RankingMap rankingMap) {
+        onNotificationRemoved(sbn, rankingMap, 0 /* reason */);
+    }
+
+    @Override
     public void onNotificationRankingUpdate(final RankingMap rankingMap) {
         if (DEBUG) Log.d(TAG, "onRankingUpdate");
         if (rankingMap != null) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoveInterceptor.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoveInterceptor.java
new file mode 100644
index 000000000000..930116ee4ae7
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoveInterceptor.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar;
+
+import android.service.notification.NotificationListenerService;
+
+/**
+ * Interface for anything that may need to prevent notifications from being removed. This is
+ * similar to a {@link NotificationLifetimeExtender} in the sense that it extends the life of
+ * a notification by preventing the removal, however, unlike the extender, the remove interceptor
+ * gets first pick at intercepting any type of removal -- the life time extender is unable to
+ * extend the life of a user dismissed or force removed notification.
+ */
+public interface NotificationRemoveInterceptor {
+
+    /**
+     * Called when a notification has been removed.
+     *
+     * @param key the entry key of the notification being removed.
+     * @param removeReason why the notification is being removed, e.g.
+     * {@link NotificationListenerService#REASON_CANCEL} or 0 if unknown.
+     *
+     * @return true if the removal should be ignored, false otherwise.
+     */
+    boolean onNotificationRemoveRequested(String key, int removeReason);
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationUpdateHandler.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationUpdateHandler.java
index 0044194e886c..1ac81982505e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationUpdateHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationUpdateHandler.java
@@ -37,8 +37,10 @@ public interface NotificationUpdateHandler {
      *
      * @param key Key identifying the notification to remove
      * @param ranking RankingMap to update with
+     * @param reason why the notification is being removed, e.g.
+     * {@link NotificationListenerService#REASON_CANCEL}.
      */
-    void removeNotification(String key, NotificationListenerService.RankingMap ranking);
+    void removeNotification(String key, NotificationListenerService.RankingMap ranking, int reason);
 
     /**
      * Update a given notification and the current notification ranking map.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
index 7d224fb13e2d..d926f88d9f60 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
@@ -15,6 +15,9 @@
  */
 package com.android.systemui.statusbar.notification;
 
+import static android.service.notification.NotificationListenerService.REASON_CANCEL;
+import static android.service.notification.NotificationListenerService.REASON_ERROR;
+
 import android.annotation.Nullable;
 import android.app.Notification;
 import android.content.Context;
@@ -30,6 +33,7 @@ import com.android.systemui.Dumpable;
 import com.android.systemui.statusbar.NotificationLifetimeExtender;
 import com.android.systemui.statusbar.NotificationPresenter;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
+import com.android.systemui.statusbar.NotificationRemoveInterceptor;
 import com.android.systemui.statusbar.NotificationUiAdjustment;
 import com.android.systemui.statusbar.NotificationUpdateHandler;
 import com.android.systemui.statusbar.notification.collection.NotificationData;
@@ -82,6 +86,7 @@ public class NotificationEntryManager implements
     final ArrayList<NotificationLifetimeExtender> mNotificationLifetimeExtenders
             = new ArrayList<>();
     private final List<NotificationEntryListener> mNotificationEntryListeners = new ArrayList<>();
+    private NotificationRemoveInterceptor mRemoveInterceptor;
 
     @Override
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
@@ -115,6 +120,11 @@ public class NotificationEntryManager implements
         mNotificationEntryListeners.add(listener);
     }
 
+    /** Sets the {@link NotificationRemoveInterceptor}. */
+    public void setNotificationRemoveInterceptor(NotificationRemoveInterceptor interceptor) {
+        mRemoveInterceptor = interceptor;
+    }
+
     /**
      * Our dependencies can have cyclic references, so some need to be lazy
      */
@@ -146,7 +156,7 @@ public class NotificationEntryManager implements
     /** Adds a {@link NotificationLifetimeExtender}. */
     public void addNotificationLifetimeExtender(NotificationLifetimeExtender extender) {
         mNotificationLifetimeExtenders.add(extender);
-        extender.setCallback(key -> removeNotification(key, mLatestRankingMap));
+        extender.setCallback(key -> removeNotification(key, mLatestRankingMap, 0));
     }
 
     public NotificationData getNotificationData() {
@@ -158,10 +168,18 @@ public class NotificationEntryManager implements
         updateNotifications();
     }
 
-    public void performRemoveNotification(StatusBarNotification n) {
+    /**
+     * Requests a notification to be removed.
+     *
+     * @param n the notification to remove.
+     * @param reason why it is being removed e.g. {@link NotificationListenerService#REASON_CANCEL},
+     *               or 0 if unknown.
+     */
+    public void performRemoveNotification(StatusBarNotification n, int reason) {
         final NotificationVisibility nv = obtainVisibility(n.getKey());
         removeNotificationInternal(
-                n.getKey(), null, nv, false /* forceRemove */, true /* removedByUser */);
+                n.getKey(), null, nv, false /* forceRemove */, true /* removedByUser */,
+                reason);
     }
 
     private NotificationVisibility obtainVisibility(String key) {
@@ -193,7 +211,8 @@ public class NotificationEntryManager implements
     @Override
     public void handleInflationException(StatusBarNotification n, Exception e) {
         removeNotificationInternal(
-                n.getKey(), null, null, true /* forceRemove */, false /* removedByUser */);
+                n.getKey(), null, null, true /* forceRemove */, false /* removedByUser */,
+                REASON_ERROR);
         for (NotificationEntryListener listener : mNotificationEntryListeners) {
             listener.onInflationError(n, e);
         }
@@ -228,9 +247,10 @@ public class NotificationEntryManager implements
     }
 
     @Override
-    public void removeNotification(String key, NotificationListenerService.RankingMap ranking) {
+    public void removeNotification(String key, NotificationListenerService.RankingMap ranking,
+            int reason) {
         removeNotificationInternal(key, ranking, obtainVisibility(key), false /* forceRemove */,
-                false /* removedByUser */);
+                false /* removedByUser */, reason);
     }
 
     private void removeNotificationInternal(
@@ -238,7 +258,15 @@ public class NotificationEntryManager implements
             @Nullable NotificationListenerService.RankingMap ranking,
             @Nullable NotificationVisibility visibility,
             boolean forceRemove,
-            boolean removedByUser) {
+            boolean removedByUser,
+            int reason) {
+
+        if (mRemoveInterceptor != null
+                && mRemoveInterceptor.onNotificationRemoveRequested(key, reason)) {
+            // Remove intercepted; skip
+            return;
+        }
+
         final NotificationEntry entry = mNotificationData.get(key);
 
         abortExistingInflation(key);
@@ -342,7 +370,8 @@ public class NotificationEntryManager implements
 
         Dependency.get(LeakDetector.class).trackInstance(entry);
         // Construct the expanded view.
-        requireBinder().inflateViews(entry, () -> performRemoveNotification(notification));
+        requireBinder().inflateViews(entry, () -> performRemoveNotification(notification,
+                REASON_CANCEL));
 
         abortExistingInflation(key);
 
@@ -383,7 +412,8 @@ public class NotificationEntryManager implements
             listener.onPreEntryUpdated(entry);
         }
 
-        requireBinder().inflateViews(entry, () -> performRemoveNotification(notification));
+        requireBinder().inflateViews(entry, () -> performRemoveNotification(notification,
+                REASON_CANCEL));
         updateNotifications();
 
         if (DEBUG) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
index d89354bea8bf..a3e18efa1915 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
@@ -247,7 +247,7 @@ public final class NotificationEntry {
      */
     public boolean showInShadeWhenBubble() {
         // We always show it in the shade if non-clearable
-        return !isClearable() || mShowInShadeWhenBubble;
+        return !isRowDismissed() && (!isClearable() || mShowInShadeWhenBubble);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationLogger.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationLogger.java
index c4ecb8259a4f..94f7e651d73d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationLogger.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationLogger.java
@@ -270,6 +270,8 @@ public class NotificationLogger implements StateListener {
         }
     }
 
+    // TODO: This method has side effects, it is NOT just logging that a notification
+    // was cleared, it also actually removes the notification
     private void logNotificationClear(String key, StatusBarNotification notification,
             NotificationVisibility nv) {
         final String pkg = notification.getPackageName();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index ebda5852faea..c6d6ca65928f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -46,6 +46,7 @@ import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.ServiceManager;
 import android.provider.Settings;
+import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
 import android.util.AttributeSet;
 import android.util.DisplayMetrics;
@@ -5586,7 +5587,8 @@ public class NotificationStackScrollLayout extends ViewGroup implements ScrollAd
             setDismissAllInProgress(false);
             for (ExpandableNotificationRow rowToRemove : viewsToRemove) {
                 if (canChildBeDismissed(rowToRemove)) {
-                    mEntryManager.removeNotification(rowToRemove.getEntry().key, null);
+                    mEntryManager.removeNotification(rowToRemove.getEntry().key, null /* ranking */,
+                            NotificationListenerService.REASON_CANCEL_ALL);
                 } else {
                     rowToRemove.resetTranslation();
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
index e4af15cf7dd1..e00d439dc1c7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.statusbar.phone;
 
+import static android.service.notification.NotificationListenerService.REASON_CLICK;
+
 import static com.android.systemui.statusbar.phone.StatusBar.getActivityOptions;
 
 import android.app.ActivityManager;
@@ -483,7 +485,7 @@ public class StatusBarNotificationActivityStarter implements NotificationActivit
         // We have to post it to the UI thread for synchronization
         mMainThreadHandler.post(() -> {
             Runnable removeRunnable =
-                    () -> mEntryManager.performRemoveNotification(notification);
+                    () -> mEntryManager.performRemoveNotification(notification, REASON_CLICK);
             if (mPresenter.isCollapsing()) {
                 // To avoid lags we're only performing the remove
                 // after the shade was collapsed
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 8b0ba9448d2b..ec8dae27809a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bubbles/BubbleControllerTest.java
@@ -18,13 +18,18 @@ package com.android.systemui.bubbles;
 
 import static android.app.Notification.FLAG_BUBBLE;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
+import static android.service.notification.NotificationListenerService.REASON_APP_CANCEL;
+import static android.service.notification.NotificationListenerService.REASON_CANCEL;
+import static android.service.notification.NotificationListenerService.REASON_CANCEL_ALL;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -50,6 +55,7 @@ import androidx.test.filters.SmallTest;
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.NotificationPresenter;
+import com.android.systemui.statusbar.NotificationRemoveInterceptor;
 import com.android.systemui.statusbar.NotificationTestHelper;
 import com.android.systemui.statusbar.notification.NotificationEntryListener;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
@@ -95,10 +101,13 @@ public class BubbleControllerTest extends SysuiTestCase {
     private FrameLayout mStatusBarView;
     @Captor
     private ArgumentCaptor<NotificationEntryListener> mEntryListenerCaptor;
+    @Captor
+    private ArgumentCaptor<NotificationRemoveInterceptor> mRemoveInterceptorCaptor;
 
     private TestableBubbleController mBubbleController;
     private StatusBarWindowController mStatusBarWindowController;
     private NotificationEntryListener mEntryListener;
+    private NotificationRemoveInterceptor mRemoveInterceptor;
 
     private NotificationTestHelper mNotificationTestHelper;
     private ExpandableNotificationRow mRow;
@@ -167,6 +176,10 @@ public class BubbleControllerTest extends SysuiTestCase {
         verify(mNotificationEntryManager, atLeastOnce())
                 .addNotificationEntryListener(mEntryListenerCaptor.capture());
         mEntryListener = mEntryListenerCaptor.getValue();
+        // And the remove interceptor
+        verify(mNotificationEntryManager, atLeastOnce())
+                .setNotificationRemoveInterceptor(mRemoveInterceptorCaptor.capture());
+        mRemoveInterceptor = mRemoveInterceptorCaptor.getValue();
     }
 
     @Test
@@ -199,6 +212,27 @@ public class BubbleControllerTest extends SysuiTestCase {
     }
 
     @Test
+    public void testRemoveBubble_withDismissedNotif() {
+        mEntryListener.onPendingEntryAdded(mRow.getEntry());
+        mBubbleController.updateBubble(mRow.getEntry());
+
+        assertTrue(mBubbleController.hasBubbles());
+        assertTrue(mRow.getEntry().showInShadeWhenBubble());
+
+        // Make it look like dismissed notif
+        mRow.getEntry().setShowInShadeWhenBubble(false);
+
+        // Now remove the bubble
+        mBubbleController.removeBubble(mRow.getEntry().key, BubbleController.DISMISS_USER_GESTURE);
+
+        // Since the notif is dismissed, once the bubble is removed, performRemoveNotification gets
+        // called to really remove the notif
+        verify(mNotificationEntryManager, times(1)).performRemoveNotification(
+                mRow.getEntry().notification, 0);
+        assertFalse(mBubbleController.hasBubbles());
+    }
+
+    @Test
     public void testDismissStack() {
         mBubbleController.updateBubble(mRow.getEntry());
         verify(mNotificationEntryManager, times(1)).updateNotifications();
@@ -455,8 +489,7 @@ public class BubbleControllerTest extends SysuiTestCase {
         mBubbleController.updateBubble(mRow.getEntry());
 
         // Simulate notification cancellation.
-        mEntryListener.onEntryRemoved(mRow.getEntry(), null /* notificationVisibility (unused) */,
-                false /* removedbyUser */);
+        mRemoveInterceptor.onNotificationRemoveRequested(mRow.getEntry().key, REASON_APP_CANCEL);
 
         mBubbleController.expandStackAndSelectBubble(key);
     }
@@ -511,6 +544,83 @@ public class BubbleControllerTest extends SysuiTestCase {
         verify(mDeleteIntent, never()).send();
     }
 
+    @Test
+    public void testRemoveBubble_succeeds_appCancel() {
+        mEntryListener.onPendingEntryAdded(mRow.getEntry());
+        mBubbleController.updateBubble(mRow.getEntry());
+
+        assertTrue(mBubbleController.hasBubbles());
+
+        boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
+                mRow.getEntry().key, REASON_APP_CANCEL);
+
+        // Cancels always remove so no need to intercept
+        assertFalse(intercepted);
+        assertFalse(mBubbleController.hasBubbles());
+    }
+
+    @Test
+    public void removeBubble_fails_clearAll()  {
+        mEntryListener.onPendingEntryAdded(mRow.getEntry());
+        mBubbleController.updateBubble(mRow.getEntry());
+
+        assertTrue(mBubbleController.hasBubbles());
+        assertTrue(mRow.getEntry().showInShadeWhenBubble());
+
+        boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
+                mRow.getEntry().key, REASON_CANCEL_ALL);
+
+        // Intercept!
+        assertTrue(intercepted);
+        // Should update show in shade state
+        assertFalse(mRow.getEntry().showInShadeWhenBubble());
+
+        verify(mNotificationEntryManager, never()).performRemoveNotification(
+                any(), anyInt());
+        assertTrue(mBubbleController.hasBubbles());
+    }
+
+    @Test
+    public void removeBubble_fails_userDismissNotif() {
+        mEntryListener.onPendingEntryAdded(mRow.getEntry());
+        mBubbleController.updateBubble(mRow.getEntry());
+
+        assertTrue(mBubbleController.hasBubbles());
+        assertTrue(mRow.getEntry().showInShadeWhenBubble());
+
+        boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
+                mRow.getEntry().key, REASON_CANCEL);
+
+        // Intercept!
+        assertTrue(intercepted);
+        // Should update show in shade state
+        assertFalse(mRow.getEntry().showInShadeWhenBubble());
+
+        verify(mNotificationEntryManager, never()).performRemoveNotification(
+                any(), anyInt());
+        assertTrue(mBubbleController.hasBubbles());
+    }
+
+    @Test
+    public void removeBubble_succeeds_userDismissBubble_userDimissNotif() {
+        mEntryListener.onPendingEntryAdded(mRow.getEntry());
+        mBubbleController.updateBubble(mRow.getEntry());
+
+        assertTrue(mBubbleController.hasBubbles());
+        assertTrue(mRow.getEntry().showInShadeWhenBubble());
+
+        // Dismiss the bubble
+        mBubbleController.removeBubble(mRow.getEntry().key, BubbleController.DISMISS_USER_GESTURE);
+        assertFalse(mBubbleController.hasBubbles());
+
+        // Dismiss the notification
+        boolean intercepted = mRemoveInterceptor.onNotificationRemoveRequested(
+                mRow.getEntry().key, REASON_CANCEL);
+
+        // It's no longer a bubble so we shouldn't intercept
+        assertFalse(intercepted);
+    }
+
     static class TestableBubbleController extends BubbleController {
         // Let's assume surfaces can be synchronized immediately.
         TestableBubbleController(Context context,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
index 1a1acdf4590c..0800cb9c33f7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
@@ -17,6 +17,8 @@
 package com.android.systemui.statusbar;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -99,7 +101,7 @@ public class NotificationListenerTest extends SysuiTestCase {
     public void testNotificationRemovalCallsRemoveNotification() {
         mListener.onNotificationRemoved(mSbn, mRanking);
         TestableLooper.get(this).processAllMessages();
-        verify(mEntryManager).removeNotification(mSbn.getKey(), mRanking);
+        verify(mEntryManager).removeNotification(eq(mSbn.getKey()), eq(mRanking), anyInt());
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationEntryManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationEntryManagerTest.java
index c8005ddacdb1..70941d3a602b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationEntryManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationEntryManagerTest.java
@@ -16,13 +16,18 @@
 
 package com.android.systemui.statusbar.notification;
 
+import static android.service.notification.NotificationListenerService.REASON_CANCEL;
+
 import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
@@ -51,6 +56,7 @@ import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.logging.MetricsLogger;
+import com.android.internal.statusbar.NotificationVisibility;
 import com.android.systemui.Dependency;
 import com.android.systemui.ForegroundServiceController;
 import com.android.systemui.InitController;
@@ -61,6 +67,7 @@ import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationLockscreenUserManager;
 import com.android.systemui.statusbar.NotificationPresenter;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
+import com.android.systemui.statusbar.NotificationRemoveInterceptor;
 import com.android.systemui.statusbar.RemoteInputController;
 import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.StatusBarIconView;
@@ -105,6 +112,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
     @Mock private ExpandableNotificationRow mRow;
     @Mock private NotificationListContainer mListContainer;
     @Mock private NotificationEntryListener mEntryListener;
+    @Mock private NotificationRemoveInterceptor mRemoveInterceptor;
     @Mock private NotificationRowBinderImpl.BindRowCallback mBindCallback;
     @Mock private HeadsUpManager mHeadsUpManager;
     @Mock private NotificationListenerService.RankingMap mRankingMap;
@@ -234,6 +242,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         Dependency.get(InitController.class).executePostInitTasks();
         mEntryManager.setUpWithPresenter(mPresenter, mListContainer, mHeadsUpManager);
         mEntryManager.addNotificationEntryListener(mEntryListener);
+        mEntryManager.setNotificationRemoveInterceptor(mRemoveInterceptor);
 
         NotificationRowBinderImpl notificationRowBinder =
                 new NotificationRowBinderImpl(mContext, true /* allowLongPress */);
@@ -341,7 +350,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         mEntry.setRow(mRow);
         mEntryManager.getNotificationData().add(mEntry);
 
-        mEntryManager.removeNotification(mSbn.getKey(), mRankingMap);
+        mEntryManager.removeNotification(mSbn.getKey(), mRankingMap, 0 /* reason */);
 
         verify(mEntryListener, never()).onInflationError(any(), any());
 
@@ -357,7 +366,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
     public void testRemoveNotification_onEntryRemoveNotFiredIfEntryDoesntExist() {
         com.android.systemui.util.Assert.isNotMainThread();
 
-        mEntryManager.removeNotification("not_a_real_key", mRankingMap);
+        mEntryManager.removeNotification("not_a_real_key", mRankingMap, 0 /* reason */);
 
         verify(mEntryListener, never()).onEntryRemoved(
                 eq(mEntry), any(), eq(false) /* removedByUser */);
@@ -370,7 +379,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         mEntryManager.setRowBinder(mMockedRowBinder);
 
         mEntryManager.addNotification(mSbn, mRankingMap);
-        mEntryManager.removeNotification(mSbn.getKey(), mRankingMap);
+        mEntryManager.removeNotification(mSbn.getKey(), mRankingMap, 0 /* reason */);
 
         verify(mEntryListener, never()).onEntryRemoved(
                 eq(mEntry), any(), eq(false /* removedByUser */));
@@ -449,7 +458,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         mEntryManager.addNotificationLifetimeExtender(extender);
 
         // WHEN the notification is removed
-        mEntryManager.removeNotification(mEntry.key, mRankingMap);
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
 
         // THEN the extender is asked to manage the lifetime
         verify(extender).setShouldManageLifetime(mEntry, true);
@@ -465,7 +474,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         mEntryManager.getNotificationData().add(mEntry);
         final FakeNotificationLifetimeExtender extender = new FakeNotificationLifetimeExtender();
         mEntryManager.addNotificationLifetimeExtender(extender);
-        mEntryManager.removeNotification(mEntry.key, mRankingMap);
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
         assertTrue(extender.isManaging(mEntry.key));
 
         // WHEN the extender finishes its extension
@@ -485,7 +494,7 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         NotificationLifetimeExtender extender = mock(NotificationLifetimeExtender.class);
         when(extender.shouldExtendLifetime(mEntry)).thenReturn(true);
         mEntryManager.addNotificationLifetimeExtender(extender);
-        mEntryManager.removeNotification(mEntry.key, mRankingMap);
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
 
         // WHEN the notification is updated
         mEntryManager.updateNotification(mEntry.notification, mRankingMap);
@@ -510,13 +519,13 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
         mEntryManager.addNotificationLifetimeExtender(extender2);
 
         // GIVEN a notification was lifetime-extended and extender2 is managing it
-        mEntryManager.removeNotification(mEntry.key, mRankingMap);
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
         verify(extender1, never()).setShouldManageLifetime(mEntry, true);
         verify(extender2).setShouldManageLifetime(mEntry, true);
 
         // WHEN the extender1 changes its mind and wants to extend the lifetime of the notif
         when(extender1.shouldExtendLifetime(mEntry)).thenReturn(true);
-        mEntryManager.removeNotification(mEntry.key, mRankingMap);
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
 
         // THEN extender2 stops managing the notif and extender1 starts managing it
         verify(extender1).setShouldManageLifetime(mEntry, true);
@@ -530,7 +539,45 @@ public class NotificationEntryManagerTest extends SysuiTestCase {
     @Test
     public void testPerformRemoveNotification_removedEntry() {
         mEntryManager.getNotificationData().remove(mSbn.getKey(), null /* ranking */);
-        mEntryManager.performRemoveNotification(mSbn);
+        mEntryManager.performRemoveNotification(mSbn, REASON_CANCEL);
+    }
+
+    @Test
+    public void testRemoveInterceptor_interceptsDontGetRemoved() {
+        // GIVEN an entry manager with a notification
+        mEntryManager.setRowBinder(mMockedRowBinder);
+        mEntryManager.getNotificationData().add(mEntry);
+
+        // GIVEN interceptor that intercepts that entry
+        when(mRemoveInterceptor.onNotificationRemoveRequested(eq(mEntry.key), anyInt()))
+                .thenReturn(true);
+
+        // WHEN the notification is removed
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
+
+        // THEN the interceptor intercepts & the entry is not removed & no listeners are called
+        assertNotNull(mEntryManager.getNotificationData().get(mEntry.key));
+        verify(mEntryListener, never()).onEntryRemoved(eq(mEntry),
+                any(NotificationVisibility.class), anyBoolean());
+    }
+
+    @Test
+    public void testRemoveInterceptor_notInterceptedGetsRemoved() {
+        // GIVEN an entry manager with a notification
+        mEntryManager.setRowBinder(mMockedRowBinder);
+        mEntryManager.getNotificationData().add(mEntry);
+
+        // GIVEN interceptor that doesn't intercept
+        when(mRemoveInterceptor.onNotificationRemoveRequested(eq(mEntry.key), anyInt()))
+                .thenReturn(false);
+
+        // WHEN the notification is removed
+        mEntryManager.removeNotification(mEntry.key, mRankingMap, 0 /* reason */);
+
+        // THEN the interceptor intercepts & the entry is not removed & no listeners are called
+        assertNull(mEntryManager.getNotificationData().get(mEntry.key));
+        verify(mEntryListener, atLeastOnce()).onEntryRemoved(eq(mEntry),
+                any(NotificationVisibility.class), anyBoolean());
     }
 
     private Notification.Action createAction() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
index 51fb47bd049c..06d76ebcff28 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.statusbar.phone;
 
+import static android.service.notification.NotificationListenerService.REASON_CLICK;
+
 import static org.mockito.AdditionalAnswers.answerVoid;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
@@ -224,7 +226,7 @@ public class StatusBarNotificationActivityStarterTest extends SysuiTestCase {
                 eq(sbn.getKey()), any(NotificationVisibility.class));
 
         // Notification is removed due to FLAG_AUTO_CANCEL
-        verify(mEntryManager).performRemoveNotification(eq(sbn));
+        verify(mEntryManager).performRemoveNotification(eq(sbn), eq(REASON_CLICK));
     }
 
     @Test
@@ -253,7 +255,7 @@ public class StatusBarNotificationActivityStarterTest extends SysuiTestCase {
         verifyZeroInteractions(mContentIntent);
 
         // Notification should not be cancelled.
-        verify(mEntryManager, never()).performRemoveNotification(eq(sbn));
+        verify(mEntryManager, never()).performRemoveNotification(eq(sbn), anyInt());
     }
 
     @Test
@@ -283,7 +285,7 @@ public class StatusBarNotificationActivityStarterTest extends SysuiTestCase {
         verifyZeroInteractions(mContentIntent);
 
         // Notification should not be cancelled.
-        verify(mEntryManager, never()).performRemoveNotification(eq(sbn));
+        verify(mEntryManager, never()).performRemoveNotification(eq(sbn), anyInt());
     }
 
     @Test
@@ -315,6 +317,6 @@ public class StatusBarNotificationActivityStarterTest extends SysuiTestCase {
         verifyNoMoreInteractions(mContentIntent);
 
         // Notification should not be cancelled.
-        verify(mEntryManager, never()).performRemoveNotification(eq(sbn));
+        verify(mEntryManager, never()).performRemoveNotification(eq(sbn), anyInt());
     }
 }