Refactor managers that defer notification removal

This CL introduces NotificationLifetimeExtender, an interface for the
various notification managers to implement to say that a notification
should stay around even after being removed.  This provides greater
consistency/standardization as opposed to having the same
implementation done differently and scattered throughout different
classes.

Also implements small refactors to try to reduce coupling where possible
(in particular reduce references to NotificationEntryManager) and move
logic to more appropriate locations (in particular move a lot of remote
input logic from NotificationEntryManager to
NotificationRemoteInputManager.)

Test: manual (see below); runtest systemui
* heads up + cancel
* guts active + cancel
* ambient pulse + cancel (UI here is janky even on master)
* have a remote input active + cancel
* have remote history, send, see that the notification stays
Change-Id: I24d345a1f2d8751827e367d1432918b3db7fa5f3
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/AlertingNotificationManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/AlertingNotificationManager.java
index c017104..35e9d55 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/AlertingNotificationManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/AlertingNotificationManager.java
@@ -36,10 +36,20 @@
  * remove notifications that appear on screen for a period of time and dismiss themselves at the
  * appropriate time.  These include heads up notifications and ambient pulses.
  */
-public abstract class AlertingNotificationManager {
+public abstract class AlertingNotificationManager implements NotificationLifetimeExtender {
     private static final String TAG = "AlertNotifManager";
     protected final Clock mClock = new Clock();
     protected final ArrayMap<String, AlertEntry> mAlertEntries = new ArrayMap<>();
+
+    /**
+     * This is the list of entries that have already been removed from the
+     * NotificationManagerService side, but we keep it to prevent the UI from looking weird and
+     * will remove when possible. See {@link NotificationLifetimeExtender}
+     */
+    protected final ArraySet<NotificationData.Entry> mExtendedLifetimeAlertEntries =
+            new ArraySet<>();
+
+    protected NotificationSafeToRemoveCallback mNotificationLifetimeFinishedCallback;
     protected int mMinimumDisplayTime;
     protected int mAutoDismissNotificationDecay;
     @VisibleForTesting
@@ -74,7 +84,7 @@
         if (alertEntry == null) {
             return true;
         }
-        if (releaseImmediately || alertEntry.wasShownLongEnough()) {
+        if (releaseImmediately || canRemoveImmediately(key)) {
             removeAlertEntry(key);
         } else {
             alertEntry.removeAsSoonAsPossible();
@@ -191,6 +201,12 @@
         onAlertEntryRemoved(alertEntry);
         entry.row.sendAccessibilityEvent(AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED);
         alertEntry.reset();
+        if (mExtendedLifetimeAlertEntries.contains(entry)) {
+            if (mNotificationLifetimeFinishedCallback != null) {
+                mNotificationLifetimeFinishedCallback.onSafeToRemove(key);
+            }
+            mExtendedLifetimeAlertEntries.remove(entry);
+        }
     }
 
     /**
@@ -207,6 +223,40 @@
         return new AlertEntry();
     }
 
+    /**
+     * Whether or not the alert can be removed currently.  If it hasn't been on screen long enough
+     * it should not be removed unless forced
+     * @param key the key to check if removable
+     * @return true if the alert entry can be removed
+     */
+    protected boolean canRemoveImmediately(String key) {
+        AlertEntry alertEntry = mAlertEntries.get(key);
+        return alertEntry == null || alertEntry.wasShownLongEnough();
+    }
+
+    ///////////////////////////////////////////////////////////////////////////////////////////////
+    // NotificationLifetimeExtender Methods
+
+    @Override
+    public void setCallback(NotificationSafeToRemoveCallback callback) {
+        mNotificationLifetimeFinishedCallback = callback;
+    }
+
+    @Override
+    public boolean shouldExtendLifetime(NotificationData.Entry entry) {
+        return !canRemoveImmediately(entry.key);
+    }
+
+    @Override
+    public void setShouldExtendLifetime(NotificationData.Entry entry, boolean shouldExtend) {
+        if (shouldExtend) {
+            mExtendedLifetimeAlertEntries.add(entry);
+        } else {
+            mExtendedLifetimeAlertEntries.remove(entry);
+        }
+    }
+    ///////////////////////////////////////////////////////////////////////////////////////////////
+
     protected class AlertEntry implements Comparable<AlertEntry> {
         @Nullable public NotificationData.Entry mEntry;
         public long mPostTime;
@@ -214,11 +264,11 @@
 
         @Nullable protected Runnable mRemoveAlertRunnable;
 
-        public void setEntry(@Nullable final NotificationData.Entry entry) {
+        public void setEntry(@NonNull final NotificationData.Entry entry) {
             setEntry(entry, () -> removeAlertEntry(entry.key));
         }
 
-        public void setEntry(@Nullable final NotificationData.Entry entry,
+        public void setEntry(@NonNull final NotificationData.Entry entry,
                 @Nullable Runnable removeAlertRunnable) {
             mEntry = entry;
             mRemoveAlertRunnable = removeAlertRunnable;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLifetimeExtender.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLifetimeExtender.java
new file mode 100644
index 0000000..42e380f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLifetimeExtender.java
@@ -0,0 +1,53 @@
+package com.android.systemui.statusbar;
+
+import com.android.systemui.statusbar.notification.NotificationData;
+
+import androidx.annotation.NonNull;
+
+/**
+ * Interface for anything that may need to keep notifications managed even after
+ * {@link NotificationListener} removes it.  The lifetime extender is in charge of performing the
+ * callback when the notification is then safe to remove.
+ */
+public interface NotificationLifetimeExtender {
+
+    /**
+     * Set the handler to callback to when the notification is safe to remove.
+     *
+     * @param callback the handler to callback
+     */
+    void setCallback(@NonNull NotificationSafeToRemoveCallback callback);
+
+    /**
+     * Determines whether or not the extender needs the notification kept after removal.
+     *
+     * @param entry the entry containing the notification to check
+     * @return true if the notification lifetime should be extended
+     */
+    boolean shouldExtendLifetime(@NonNull NotificationData.Entry entry);
+
+    /**
+     * Sets whether or not the lifetime should be extended.  In practice, if shouldExtend is
+     * true, this is where the extender starts managing the entry internally and is now
+     * responsible for calling {@link NotificationSafeToRemoveCallback#onSafeToRemove(String)} when
+     * the entry is safe to remove.  If shouldExtend is false, the extender no longer needs to
+     * worry about it (either because we will be removing it anyway or the entry is no longer
+     * removed due to an update).
+     *
+     * @param entry the entry to mark as having an extended lifetime
+     * @param shouldExtend true if the extender should manage the entry now, false otherwise
+     */
+    void setShouldExtendLifetime(@NonNull NotificationData.Entry entry, boolean shouldExtend);
+
+    /**
+     * The callback for when the notification is now safe to remove (i.e. its lifetime has ended).
+     */
+    interface NotificationSafeToRemoveCallback {
+        /**
+         * Called when the lifetime extender determines it's safe to remove.
+         *
+         * @param key key of the entry that is now safe to remove
+         */
+        void onSafeToRemove(String key);
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
index 9b375df..cfa09bc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationListener.java
@@ -76,7 +76,6 @@
             mPresenter.getHandler().post(() -> {
                 processForRemoteInput(sbn.getNotification(), mContext);
                 String key = sbn.getKey();
-                mEntryManager.removeKeyKeptForRemoteInput(key);
                 boolean isUpdate =
                         mEntryManager.getNotificationData().get(key) != null;
                 // In case we don't allow child notifications, we ignore children of
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index 929713c..1a3e812 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
@@ -17,8 +17,10 @@
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN_OR_SPLIT_SCREEN_SECONDARY;
 
+import android.annotation.NonNull;
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
+import android.app.Notification;
 import android.app.PendingIntent;
 import android.app.RemoteInput;
 import android.content.Context;
@@ -29,6 +31,7 @@
 import android.os.SystemProperties;
 import android.os.UserManager;
 import android.service.notification.StatusBarNotification;
+import android.text.TextUtils;
 import android.util.ArraySet;
 import android.util.Log;
 import android.view.MotionEvent;
@@ -50,6 +53,7 @@
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
+import java.util.ArrayList;
 import java.util.Set;
 
 /**
@@ -61,10 +65,10 @@
 public class NotificationRemoteInputManager implements Dumpable {
     public static final boolean ENABLE_REMOTE_INPUT =
             SystemProperties.getBoolean("debug.enable_remote_input", true);
-    public static final boolean FORCE_REMOTE_INPUT_HISTORY =
+    public static boolean FORCE_REMOTE_INPUT_HISTORY =
             SystemProperties.getBoolean("debug.force_remoteinput_history", true);
     private static final boolean DEBUG = false;
-    private static final String TAG = "NotificationRemoteInputManager";
+    private static final String TAG = "NotifRemoteInputManager";
 
     /**
      * How long to wait before auto-dismissing a notification that was kept for remote input, and
@@ -74,12 +78,25 @@
      */
     private static final int REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY = 200;
 
-    protected final ArraySet<NotificationData.Entry> mRemoteInputEntriesToRemoveOnCollapse =
+    /**
+     * Notifications that are already removed but are kept around because we want to show the
+     * remote input history. See {@link RemoteInputHistoryExtender} and
+     * {@link SmartReplyHistoryExtender}.
+     */
+    protected final ArraySet<String> mKeysKeptForRemoteInputHistory = new ArraySet<>();
+
+    /**
+     * Notifications that are already removed but are kept around because the remote input is
+     * actively being used (i.e. user is typing in it).  See {@link RemoteInputActiveExtender}.
+     */
+    protected final ArraySet<NotificationData.Entry> mEntriesKeptForRemoteInputActive =
             new ArraySet<>();
 
     // Dependencies:
     protected final NotificationLockscreenUserManager mLockscreenUserManager =
             Dependency.get(NotificationLockscreenUserManager.class);
+    protected final SmartReplyController mSmartReplyController =
+            Dependency.get(SmartReplyController.class);
 
     protected final Context mContext;
     private final UserManager mUserManager;
@@ -87,8 +104,11 @@
     protected RemoteInputController mRemoteInputController;
     protected NotificationPresenter mPresenter;
     protected NotificationEntryManager mEntryManager;
+    protected NotificationLifetimeExtender.NotificationSafeToRemoveCallback
+            mNotificationLifetimeFinishedCallback;
     protected IStatusBarService mBarService;
     protected Callback mCallback;
+    protected final ArrayList<NotificationLifetimeExtender> mLifetimeExtenders = new ArrayList<>();
 
     private final RemoteViews.OnClickHandler mOnClickHandler = new RemoteViews.OnClickHandler() {
 
@@ -276,6 +296,7 @@
         mBarService = IStatusBarService.Stub.asInterface(
                 ServiceManager.getService(Context.STATUS_BAR_SERVICE));
         mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
+        addLifetimeExtenders();
     }
 
     public void setUpWithPresenter(NotificationPresenter presenter,
@@ -290,16 +311,16 @@
             @Override
             public void onRemoteInputSent(NotificationData.Entry entry) {
                 if (FORCE_REMOTE_INPUT_HISTORY
-                        && mEntryManager.isNotificationKeptForRemoteInput(entry.key)) {
-                    mEntryManager.removeNotification(entry.key, null);
-                } else if (mRemoteInputEntriesToRemoveOnCollapse.contains(entry)) {
+                        && isNotificationKeptForRemoteInputHistory(entry.key)) {
+                    mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.key);
+                } else if (mEntriesKeptForRemoteInputActive.contains(entry)) {
                     // We're currently holding onto this notification, but from the apps point of
                     // view it is already canceled, so we'll need to cancel it on the apps behalf
                     // after sending - unless the app posts an update in the mean time, so wait a
                     // bit.
                     mPresenter.getHandler().postDelayed(() -> {
-                        if (mRemoteInputEntriesToRemoveOnCollapse.remove(entry)) {
-                            mEntryManager.removeNotification(entry.key, null);
+                        if (mEntriesKeptForRemoteInputActive.remove(entry)) {
+                            mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.key);
                         }
                     }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
                 }
@@ -310,45 +331,74 @@
                 }
             }
         });
+        mSmartReplyController.setCallback((entry, reply) -> {
+            StatusBarNotification newSbn =
+                    rebuildNotificationWithRemoteInput(entry, reply, true /* showSpinner */);
+            mEntryManager.updateNotification(newSbn, null /* ranking */);
+        });
+    }
 
+    /**
+     * Adds all the notification lifetime extenders. Each extender represents a reason for the
+     * NotificationRemoteInputManager to keep a notification lifetime extended.
+     */
+    protected void addLifetimeExtenders() {
+        mLifetimeExtenders.add(new RemoteInputHistoryExtender());
+        mLifetimeExtenders.add(new SmartReplyHistoryExtender());
+        mLifetimeExtenders.add(new RemoteInputActiveExtender());
+    }
+
+    public ArrayList<NotificationLifetimeExtender> getLifetimeExtenders() {
+        return mLifetimeExtenders;
     }
 
     public RemoteInputController getController() {
         return mRemoteInputController;
     }
 
-    public void onUpdateNotification(NotificationData.Entry entry) {
-        mRemoteInputEntriesToRemoveOnCollapse.remove(entry);
-    }
-
-    /**
-     * Returns true if NotificationRemoteInputManager wants to keep this notification around.
-     *
-     * @param entry notification being removed
-     */
-    public boolean onRemoveNotification(NotificationData.Entry entry) {
-        if (entry != null && mRemoteInputController.isRemoteInputActive(entry)
-                && (entry.row != null && !entry.row.isDismissed())) {
-            mRemoteInputEntriesToRemoveOnCollapse.add(entry);
-            return true;
-        }
-        return false;
-    }
-
     public void onPerformRemoveNotification(StatusBarNotification n,
             NotificationData.Entry entry) {
+        if (mKeysKeptForRemoteInputHistory.contains(n.getKey())) {
+            mKeysKeptForRemoteInputHistory.remove(n.getKey());
+        }
         if (mRemoteInputController.isRemoteInputActive(entry)) {
             mRemoteInputController.removeRemoteInput(entry, null);
         }
     }
 
-    public void removeRemoteInputEntriesKeptUntilCollapsed() {
-        for (int i = 0; i < mRemoteInputEntriesToRemoveOnCollapse.size(); i++) {
-            NotificationData.Entry entry = mRemoteInputEntriesToRemoveOnCollapse.valueAt(i);
+    public void onPanelCollapsed() {
+        for (int i = 0; i < mEntriesKeptForRemoteInputActive.size(); i++) {
+            NotificationData.Entry entry = mEntriesKeptForRemoteInputActive.valueAt(i);
             mRemoteInputController.removeRemoteInput(entry, null);
-            mEntryManager.removeNotification(entry.key, mEntryManager.getLatestRankingMap());
+            if (mNotificationLifetimeFinishedCallback != null) {
+                mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.key);
+            }
         }
-        mRemoteInputEntriesToRemoveOnCollapse.clear();
+        mEntriesKeptForRemoteInputActive.clear();
+    }
+
+    public boolean isNotificationKeptForRemoteInputHistory(String key) {
+        return mKeysKeptForRemoteInputHistory.contains(key);
+    }
+
+    public boolean shouldKeepForRemoteInputHistory(NotificationData.Entry entry) {
+        if (entry.row == null || entry.row.isDismissed()) {
+            return false;
+        }
+        if (!FORCE_REMOTE_INPUT_HISTORY) {
+            return false;
+        }
+        return (mRemoteInputController.isSpinning(entry.key) || entry.hasJustSentRemoteInput());
+    }
+
+    public boolean shouldKeepForSmartReplyHistory(NotificationData.Entry entry) {
+        if (entry.row == null || entry.row.isDismissed()) {
+            return false;
+        }
+        if (!FORCE_REMOTE_INPUT_HISTORY) {
+            return false;
+        }
+        return mSmartReplyController.isSendingSmartReply(entry.key);
     }
 
     public void checkRemoteInputOutside(MotionEvent event) {
@@ -359,11 +409,63 @@
         }
     }
 
+    @VisibleForTesting
+    StatusBarNotification rebuildNotificationForCanceledSmartReplies(
+            NotificationData.Entry entry) {
+        return rebuildNotificationWithRemoteInput(entry, null /* remoteInputTest */,
+                false /* showSpinner */);
+    }
+
+    @VisibleForTesting
+    StatusBarNotification rebuildNotificationWithRemoteInput(NotificationData.Entry entry,
+            CharSequence remoteInputText, boolean showSpinner) {
+        StatusBarNotification sbn = entry.notification;
+
+        Notification.Builder b = Notification.Builder
+                .recoverBuilder(mContext, sbn.getNotification().clone());
+        if (remoteInputText != null) {
+            CharSequence[] oldHistory = sbn.getNotification().extras
+                    .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+            CharSequence[] newHistory;
+            if (oldHistory == null) {
+                newHistory = new CharSequence[1];
+            } else {
+                newHistory = new CharSequence[oldHistory.length + 1];
+                System.arraycopy(oldHistory, 0, newHistory, 1, oldHistory.length);
+            }
+            newHistory[0] = String.valueOf(remoteInputText);
+            b.setRemoteInputHistory(newHistory);
+        }
+        b.setShowRemoteInputSpinner(showSpinner);
+        b.setHideSmartReplies(true);
+
+        Notification newNotification = b.build();
+
+        // Undo any compatibility view inflation
+        newNotification.contentView = sbn.getNotification().contentView;
+        newNotification.bigContentView = sbn.getNotification().bigContentView;
+        newNotification.headsUpContentView = sbn.getNotification().headsUpContentView;
+
+        return new StatusBarNotification(
+                sbn.getPackageName(),
+                sbn.getOpPkg(),
+                sbn.getId(),
+                sbn.getTag(),
+                sbn.getUid(),
+                sbn.getInitialPid(),
+                newNotification,
+                sbn.getUser(),
+                sbn.getOverrideGroupKey(),
+                sbn.getPostTime());
+    }
+
     @Override
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("NotificationRemoteInputManager state:");
-        pw.print("  mRemoteInputEntriesToRemoveOnCollapse: ");
-        pw.println(mRemoteInputEntriesToRemoveOnCollapse);
+        pw.print("  mKeysKeptForRemoteInputHistory: ");
+        pw.println(mKeysKeptForRemoteInputHistory);
+        pw.print("  mEntriesKeptForRemoteInputActive: ");
+        pw.println(mEntriesKeptForRemoteInputActive);
     }
 
     public void bindRow(ExpandableNotificationRow row) {
@@ -372,8 +474,133 @@
     }
 
     @VisibleForTesting
-    public Set<NotificationData.Entry> getRemoteInputEntriesToRemoveOnCollapse() {
-        return mRemoteInputEntriesToRemoveOnCollapse;
+    public Set<NotificationData.Entry> getEntriesKeptForRemoteInputActive() {
+        return mEntriesKeptForRemoteInputActive;
+    }
+
+    /**
+     * NotificationRemoteInputManager has multiple reasons to keep notification lifetime extended
+     * so we implement multiple NotificationLifetimeExtenders
+     */
+    protected abstract class RemoteInputExtender implements NotificationLifetimeExtender {
+        @Override
+        public void setCallback(NotificationSafeToRemoveCallback callback) {
+            if (mNotificationLifetimeFinishedCallback == null) {
+                mNotificationLifetimeFinishedCallback = callback;
+            }
+        }
+    }
+
+    /**
+     * Notification is kept alive as it was cancelled in response to a remote input interaction.
+     * This allows us to show what you replied and allows you to continue typing into it.
+     */
+    protected class RemoteInputHistoryExtender extends RemoteInputExtender {
+        @Override
+        public boolean shouldExtendLifetime(@NonNull NotificationData.Entry entry) {
+            return shouldKeepForRemoteInputHistory(entry);
+        }
+
+        @Override
+        public void setShouldExtendLifetime(NotificationData.Entry entry,
+                boolean shouldExtend) {
+            if (shouldExtend) {
+                CharSequence remoteInputText = entry.remoteInputText;
+                if (TextUtils.isEmpty(remoteInputText)) {
+                    remoteInputText = entry.remoteInputTextWhenReset;
+                }
+                StatusBarNotification newSbn = rebuildNotificationWithRemoteInput(entry,
+                        remoteInputText, false /* showSpinner */);
+                entry.onRemoteInputInserted();
+
+                if (newSbn == null) {
+                    return;
+                }
+
+                mEntryManager.updateNotification(newSbn, null);
+
+                // Ensure the entry hasn't already been removed. This can happen if there is an
+                // inflation exception while updating the remote history
+                if (entry.row == null || entry.row.isRemoved()) {
+                    return;
+                }
+
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Keeping notification around after sending remote input "
+                            + entry.key);
+                }
+
+                mKeysKeptForRemoteInputHistory.add(entry.key);
+            } else {
+                mKeysKeptForRemoteInputHistory.remove(entry.key);
+            }
+        }
+    }
+
+    /**
+     * Notification is kept alive for smart reply history.  Similar to REMOTE_INPUT_HISTORY but with
+     * {@link SmartReplyController} specific logic
+     */
+    protected class SmartReplyHistoryExtender extends RemoteInputExtender {
+        @Override
+        public boolean shouldExtendLifetime(@NonNull NotificationData.Entry entry) {
+            return shouldKeepForSmartReplyHistory(entry);
+        }
+
+        @Override
+        public void setShouldExtendLifetime(NotificationData.Entry entry,
+                boolean shouldExtend) {
+            if (shouldExtend) {
+                StatusBarNotification newSbn = rebuildNotificationForCanceledSmartReplies(entry);
+
+                if (newSbn == null) {
+                    return;
+                }
+
+                mEntryManager.updateNotification(newSbn, null);
+
+                if (entry.row == null || entry.row.isRemoved()) {
+                    return;
+                }
+
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Keeping notification around after sending smart reply "
+                            + entry.key);
+                }
+
+                mKeysKeptForRemoteInputHistory.add(entry.key);
+            } else {
+                mKeysKeptForRemoteInputHistory.remove(entry.key);
+                mSmartReplyController.stopSending(entry);
+            }
+        }
+    }
+
+    /**
+     * Notification is kept alive because the user is still using the remote input
+     */
+    protected class RemoteInputActiveExtender extends RemoteInputExtender {
+        @Override
+        public boolean shouldExtendLifetime(@NonNull NotificationData.Entry entry) {
+            if (entry.row == null || entry.row.isDismissed()) {
+                return false;
+            }
+            return mRemoteInputController.isRemoteInputActive(entry);
+        }
+
+        @Override
+        public void setShouldExtendLifetime(NotificationData.Entry entry,
+                boolean shouldExtend) {
+            if (shouldExtend) {
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Keeping notification around while remote input active "
+                            + entry.key);
+                }
+                mEntriesKeptForRemoteInputActive.add(entry);
+            } else {
+                mEntriesKeptForRemoteInputActive.remove(entry);
+            }
+        }
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
index e43c9e5..fb888dd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
@@ -33,20 +33,19 @@
 public class SmartReplyController {
     private IStatusBarService mBarService;
     private Set<String> mSendingKeys = new ArraySet<>();
+    private Callback mCallback;
 
     public SmartReplyController() {
         mBarService = Dependency.get(IStatusBarService.class);
     }
 
-    public void smartReplySent(NotificationData.Entry entry, int replyIndex, CharSequence reply) {
-        NotificationEntryManager notificationEntryManager
-                = Dependency.get(NotificationEntryManager.class);
-        StatusBarNotification newSbn =
-                notificationEntryManager.rebuildNotificationWithRemoteInput(entry, reply,
-                        true /* showSpinner */);
-        notificationEntryManager.updateNotification(newSbn, null /* ranking */);
-        mSendingKeys.add(entry.key);
+    public void setCallback(Callback callback) {
+        mCallback = callback;
+    }
 
+    public void smartReplySent(NotificationData.Entry entry, int replyIndex, CharSequence reply) {
+        mCallback.onSmartReplySent(entry, reply);
+        mSendingKeys.add(entry.key);
         try {
             mBarService.onNotificationSmartReplySent(entry.notification.getKey(),
                     replyIndex);
@@ -77,4 +76,17 @@
             mSendingKeys.remove(entry.notification.getKey());
         }
     }
+
+    /**
+     * Callback for any class that needs to do something in response to a smart reply being sent.
+     */
+    public interface Callback {
+        /**
+         * A smart reply has just been sent for a notification
+         *
+         * @param entry the entry for the notification
+         * @param reply the reply that was sent
+         */
+        void onSmartReplySent(NotificationData.Entry entry, CharSequence reply);
+    }
 }
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 b655a6b..906bbb9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
@@ -37,7 +37,6 @@
 import android.service.notification.NotificationListenerService;
 import android.service.notification.NotificationStats;
 import android.service.notification.StatusBarNotification;
-import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.EventLog;
@@ -58,6 +57,7 @@
 import com.android.systemui.R;
 import com.android.systemui.UiOffloadThread;
 import com.android.systemui.recents.misc.SystemServicesProxy;
+import com.android.systemui.statusbar.NotificationLifetimeExtender;
 import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationLockscreenUserManager;
 import com.android.systemui.statusbar.NotificationMediaManager;
@@ -65,7 +65,6 @@
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationUiAdjustment;
 import com.android.systemui.statusbar.NotificationUpdateHandler;
-import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.notification.row.NotificationInflater;
 import com.android.systemui.statusbar.notification.row.RowInflaterTask;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
@@ -100,8 +99,6 @@
     protected final Context mContext;
     protected final HashMap<String, NotificationData.Entry> mPendingNotifications = new HashMap<>();
     protected final NotificationClicker mNotificationClicker = new NotificationClicker();
-    protected final ArraySet<NotificationData.Entry> mHeadsUpEntriesToRemoveOnSwitch =
-            new ArraySet<>();
 
     // Dependencies:
     protected final NotificationLockscreenUserManager mLockscreenUserManager =
@@ -124,8 +121,6 @@
             Dependency.get(ForegroundServiceController.class);
     protected final NotificationListener mNotificationListener =
             Dependency.get(NotificationListener.class);
-    private final SmartReplyController mSmartReplyController =
-            Dependency.get(SmartReplyController.class);
 
     protected IStatusBarService mBarService;
     protected NotificationPresenter mPresenter;
@@ -139,13 +134,9 @@
     protected boolean mUseHeadsUp = false;
     protected boolean mDisableNotificationAlerts;
     protected NotificationListContainer mListContainer;
+    protected final ArrayList<NotificationLifetimeExtender> mNotificationLifetimeExtenders
+            = new ArrayList<>();
     private ExpandableNotificationRow.OnAppOpsClickListener mOnAppOpsClickListener;
-    /**
-     * Notifications with keys in this set are not actually around anymore. We kept them around
-     * when they were canceled in response to a remote input interaction. This allows us to show
-     * what you replied and allows you to continue typing into it.
-     */
-    private final ArraySet<String> mKeysKeptForRemoteInput = new ArraySet<>();
 
 
     private final class NotificationClicker implements View.OnClickListener {
@@ -198,14 +189,6 @@
                 }
             };
 
-    public NotificationListenerService.RankingMap getLatestRankingMap() {
-        return mLatestRankingMap;
-    }
-
-    public void setLatestRankingMap(NotificationListenerService.RankingMap latestRankingMap) {
-        mLatestRankingMap = latestRankingMap;
-    }
-
     public void setDisableNotificationAlerts(boolean disableNotificationAlerts) {
         mDisableNotificationAlerts = disableNotificationAlerts;
         mHeadsUpObserver.onChange(true);
@@ -215,18 +198,6 @@
         mDeviceProvisionedController.removeCallback(mDeviceProvisionedListener);
     }
 
-    public void onHeadsUpStateChanged(NotificationData.Entry entry, boolean isHeadsUp) {
-        if (!isHeadsUp && mHeadsUpEntriesToRemoveOnSwitch.contains(entry)) {
-            removeNotification(entry.key, getLatestRankingMap());
-            mHeadsUpEntriesToRemoveOnSwitch.remove(entry);
-            if (mHeadsUpEntriesToRemoveOnSwitch.isEmpty()) {
-                setLatestRankingMap(null);
-            }
-        } else {
-            updateNotificationRanking(null);
-        }
-    }
-
     @Override
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("NotificationEntryManager state:");
@@ -240,8 +211,6 @@
         }
         pw.print("  mUseHeadsUp=");
         pw.println(mUseHeadsUp);
-        pw.print("  mKeysKeptForRemoteInput: ");
-        pw.println(mKeysKeptForRemoteInput);
     }
 
     public NotificationEntryManager(Context context) {
@@ -294,6 +263,14 @@
                     mHeadsUpObserver);
         }
 
+        mNotificationLifetimeExtenders.add(mHeadsUpManager);
+        mNotificationLifetimeExtenders.add(mGutsManager);
+        mNotificationLifetimeExtenders.addAll(mRemoteInputManager.getLifetimeExtenders());
+
+        for (NotificationLifetimeExtender extender : mNotificationLifetimeExtenders) {
+            extender.setCallback(key -> removeNotification(key, mLatestRankingMap));
+        }
+
         mDeviceProvisionedController.addCallback(mDeviceProvisionedListener);
 
         mHeadsUpObserver.onChange(true); // set up
@@ -397,11 +374,6 @@
                 true);
         NotificationData.Entry entry = mNotificationData.get(n.getKey());
 
-        if (FORCE_REMOTE_INPUT_HISTORY
-                && mKeysKeptForRemoteInput.contains(n.getKey())) {
-            mKeysKeptForRemoteInput.remove(n.getKey());
-        }
-
         mRemoteInputManager.onPerformRemoveNotification(n, entry);
         final String pkg = n.getPackageName();
         final String tag = n.getTag();
@@ -433,7 +405,7 @@
      * WARNING: this will call back into us.  Don't hold any locks.
      */
     void handleNotificationError(StatusBarNotification n, String message) {
-        removeNotification(n.getKey(), null);
+        removeNotificationInternal(n.getKey(), null, true /* forceRemove */);
         try {
             mBarService.onNotificationError(n.getPackageName(), n.getTag(), n.getId(), n.getUid(),
                     n.getInitialPid(), message, n.getUserId());
@@ -487,7 +459,11 @@
 
     @Override
     public void removeNotification(String key, NotificationListenerService.RankingMap ranking) {
-        boolean deferRemoval = false;
+        removeNotificationInternal(key, ranking, false /* forceRemove */);
+    }
+
+    private void removeNotificationInternal(String key,
+            @Nullable NotificationListenerService.RankingMap ranking, boolean forceRemove) {
         abortExistingInflation(key);
         if (mHeadsUpManager.contains(key)) {
             // A cancel() in response to a remote input shouldn't be delayed, as it makes the
@@ -497,154 +473,53 @@
             boolean ignoreEarliestRemovalTime = mRemoteInputManager.getController().isSpinning(key)
                     && !FORCE_REMOTE_INPUT_HISTORY
                     || !mVisualStabilityManager.isReorderingAllowed();
-            deferRemoval = !mHeadsUpManager.removeNotification(key,  ignoreEarliestRemovalTime);
+
+            // Attempt to remove notification.
+            mHeadsUpManager.removeNotification(key, ignoreEarliestRemovalTime);
         }
-        mMediaManager.onNotificationRemoved(key);
 
         NotificationData.Entry entry = mNotificationData.get(key);
-        if (FORCE_REMOTE_INPUT_HISTORY
-                && shouldKeepForRemoteInput(entry)
-                && entry.row != null && !entry.row.isDismissed()) {
-            CharSequence remoteInputText = entry.remoteInputText;
-            if (TextUtils.isEmpty(remoteInputText)) {
-                remoteInputText = entry.remoteInputTextWhenReset;
-            }
-            StatusBarNotification newSbn = rebuildNotificationWithRemoteInput(entry,
-                    remoteInputText, false /* showSpinner */);
-            boolean updated = false;
-            entry.onRemoteInputInserted();
-            try {
-                updateNotificationInternal(newSbn, null);
-                updated = true;
-            } catch (InflationException e) {
-                deferRemoval = false;
-            }
-            if (updated) {
-                Log.w(TAG, "Keeping notification around after sending remote input "+ entry.key);
-                addKeyKeptForRemoteInput(entry.key);
-                return;
-            }
-        }
 
-        if (FORCE_REMOTE_INPUT_HISTORY
-                && shouldKeepForSmartReply(entry)
-                && entry.row != null && !entry.row.isDismissed()) {
-            // Turn off the spinner and hide buttons when an app cancels the notification.
-            StatusBarNotification newSbn = rebuildNotificationForCanceledSmartReplies(entry);
-            boolean updated = false;
-            try {
-                updateNotificationInternal(newSbn, null);
-                updated = true;
-            } catch (InflationException e) {
-                // Ignore just don't keep the notification around.
-            }
-            // Treat the reply as longer sending.
-            mSmartReplyController.stopSending(entry);
-            if (updated) {
-                Log.w(TAG, "Keeping notification around after sending smart reply " + entry.key);
-                addKeyKeptForRemoteInput(entry.key);
-                return;
-            }
-        }
-
-        // Actually removing notification so smart reply controller can forget about it.
-        mSmartReplyController.stopSending(entry);
-
-        if (deferRemoval) {
-            mLatestRankingMap = ranking;
-            mHeadsUpEntriesToRemoveOnSwitch.add(mHeadsUpManager.getEntry(key));
+        if (entry == null) {
+            mCallback.onNotificationRemoved(key, null /* old */);
             return;
         }
 
-        if (mRemoteInputManager.onRemoveNotification(entry)) {
-            mLatestRankingMap = ranking;
-            return;
+        // If a manager needs to keep the notification around for whatever reason, we return early
+        // and keep the notification
+        if (!forceRemove) {
+            for (NotificationLifetimeExtender extender : mNotificationLifetimeExtenders) {
+                if (extender.shouldExtendLifetime(entry)) {
+                    mLatestRankingMap = ranking;
+                    extender.setShouldExtendLifetime(entry, true /* shouldExtend */);
+                    return;
+                }
+            }
         }
 
-        if (entry != null && mGutsManager.getExposedGuts() != null
-                && mGutsManager.getExposedGuts() == entry.row.getGuts()
-                && entry.row.getGuts() != null && !entry.row.getGuts().isLeavebehind()) {
-            Log.w(TAG, "Keeping notification because it's showing guts. " + key);
-            mLatestRankingMap = ranking;
-            mGutsManager.setKeyToRemoveOnGutsClosed(key);
-            return;
+        // At this point, we are guaranteed the notification will be removed
+
+        // Ensure any managers keeping the lifetime extended stop managing the entry
+        for (NotificationLifetimeExtender extender: mNotificationLifetimeExtenders) {
+            extender.setShouldExtendLifetime(entry, false /* shouldExtend */);
         }
 
-        if (entry != null) {
-            mForegroundServiceController.removeNotification(entry.notification);
-        }
+        mMediaManager.onNotificationRemoved(key);
+        mForegroundServiceController.removeNotification(entry.notification);
 
-        if (entry != null && entry.row != null) {
+        if (entry.row != null) {
             entry.row.setRemoved();
             mListContainer.cleanUpViewState(entry.row);
         }
+
         // Let's remove the children if this was a summary
         handleGroupSummaryRemoved(key);
+
         StatusBarNotification old = removeNotificationViews(key, ranking);
 
         mCallback.onNotificationRemoved(key, old);
     }
 
-    public StatusBarNotification rebuildNotificationWithRemoteInput(NotificationData.Entry entry,
-            CharSequence remoteInputText, boolean showSpinner) {
-        StatusBarNotification sbn = entry.notification;
-
-        Notification.Builder b = Notification.Builder
-                .recoverBuilder(mContext, sbn.getNotification().clone());
-        if (remoteInputText != null) {
-            CharSequence[] oldHistory = sbn.getNotification().extras
-                    .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
-            CharSequence[] newHistory;
-            if (oldHistory == null) {
-                newHistory = new CharSequence[1];
-            } else {
-                newHistory = new CharSequence[oldHistory.length + 1];
-                System.arraycopy(oldHistory, 0, newHistory, 1, oldHistory.length);
-            }
-            newHistory[0] = String.valueOf(remoteInputText);
-            b.setRemoteInputHistory(newHistory);
-        }
-        b.setShowRemoteInputSpinner(showSpinner);
-        b.setHideSmartReplies(true);
-
-        Notification newNotification = b.build();
-
-        // Undo any compatibility view inflation
-        newNotification.contentView = sbn.getNotification().contentView;
-        newNotification.bigContentView = sbn.getNotification().bigContentView;
-        newNotification.headsUpContentView = sbn.getNotification().headsUpContentView;
-
-        StatusBarNotification newSbn = new StatusBarNotification(sbn.getPackageName(),
-                sbn.getOpPkg(),
-                sbn.getId(), sbn.getTag(), sbn.getUid(), sbn.getInitialPid(),
-                newNotification, sbn.getUser(), sbn.getOverrideGroupKey(), sbn.getPostTime());
-        return newSbn;
-    }
-
-    @VisibleForTesting
-    StatusBarNotification rebuildNotificationForCanceledSmartReplies(
-            NotificationData.Entry entry) {
-        return rebuildNotificationWithRemoteInput(entry, null /* remoteInputTest */,
-                false /* showSpinner */);
-    }
-
-    private boolean shouldKeepForSmartReply(NotificationData.Entry entry) {
-        return entry != null && mSmartReplyController.isSendingSmartReply(entry.key);
-    }
-
-    private boolean shouldKeepForRemoteInput(NotificationData.Entry entry) {
-        if (entry == null) {
-            return false;
-        }
-        if (mRemoteInputManager.getController().isSpinning(entry.key)) {
-            return true;
-        }
-        if (entry.hasJustSentRemoteInput()) {
-            return true;
-        }
-        return false;
-    }
-
     private StatusBarNotification removeNotificationViews(String key,
             NotificationListenerService.RankingMap ranking) {
         NotificationData.Entry entry = mNotificationData.remove(key, ranking);
@@ -683,9 +558,9 @@
                 NotificationData.Entry childEntry = row.getEntry();
                 boolean isForeground = (row.getStatusBarNotification().getNotification().flags
                         & Notification.FLAG_FOREGROUND_SERVICE) != 0;
-                boolean keepForReply = FORCE_REMOTE_INPUT_HISTORY
-                        && (shouldKeepForRemoteInput(childEntry)
-                                || shouldKeepForSmartReply(childEntry));
+                boolean keepForReply =
+                        mRemoteInputManager.shouldKeepForRemoteInputHistory(childEntry)
+                        || mRemoteInputManager.shouldKeepForSmartReplyHistory(childEntry);
                 if (isForeground || keepForReply) {
                     // the child is a foreground service notification which we can't remove or it's
                     // a child we're keeping around for reply!
@@ -868,13 +743,11 @@
         if (entry == null) {
             return;
         }
-        mHeadsUpEntriesToRemoveOnSwitch.remove(entry);
-        mRemoteInputManager.onUpdateNotification(entry);
-        mSmartReplyController.stopSending(entry);
 
-        if (key.equals(mGutsManager.getKeyToRemoveOnGutsClosed())) {
-            mGutsManager.setKeyToRemoveOnGutsClosed(null);
-            Log.w(TAG, "Notification that was kept for guts was updated. " + key);
+        // Notification is updated so it is essentially re-added and thus alive again.  Don't need
+        // to keep it's lifetime extended.
+        for (NotificationLifetimeExtender extender : mNotificationLifetimeExtenders) {
+            extender.setShouldExtendLifetime(entry, false /* shouldExtend */);
         }
 
         Notification n = notification.getNotification();
@@ -1080,20 +953,6 @@
         return mHeadsUpManager.contains(key);
     }
 
-    public boolean isNotificationKeptForRemoteInput(String key) {
-        return mKeysKeptForRemoteInput.contains(key);
-    }
-
-    public void removeKeyKeptForRemoteInput(String key) {
-        mKeysKeptForRemoteInput.remove(key);
-    }
-
-    public void addKeyKeptForRemoteInput(String key) {
-        if (FORCE_REMOTE_INPUT_HISTORY) {
-            mKeysKeptForRemoteInput.add(key);
-        }
-    }
-
     /**
      * Callback for NotificationEntryManager.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationGutsManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationGutsManager.java
index e635976..a096baa 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationGutsManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationGutsManager.java
@@ -38,27 +38,27 @@
 import android.view.View;
 import android.view.accessibility.AccessibilityManager;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto;
 import com.android.systemui.Dependency;
 import com.android.systemui.Dumpable;
 import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
+import com.android.systemui.statusbar.NotificationLifetimeExtender;
+import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.NotificationLockscreenUserManager;
 import com.android.systemui.statusbar.NotificationPresenter;
-import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.stack.NotificationListContainer;
 import com.android.systemui.statusbar.phone.StatusBar;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 
-import androidx.annotation.VisibleForTesting;
-
 /**
  * Handles various NotificationGuts related tasks, such as binding guts to a row, opening and
  * closing guts, and keeping track of the currently exposed notification guts.
  */
-public class NotificationGutsManager implements Dumpable {
+public class NotificationGutsManager implements Dumpable, NotificationLifetimeExtender {
     private static final String TAG = "NotificationGutsManager";
 
     // Must match constant in Settings. Used to highlight preferences when linking to Settings.
@@ -75,12 +75,13 @@
     // which notification is currently being longpress-examined by the user
     private NotificationGuts mNotificationGutsExposed;
     private NotificationMenuRowPlugin.MenuItem mGutsMenuItem;
-    protected NotificationPresenter mPresenter;
-    protected NotificationEntryManager mEntryManager;
+    private NotificationPresenter mPresenter;
+    private NotificationSafeToRemoveCallback mNotificationLifetimeFinishedCallback;
     private NotificationListContainer mListContainer;
     private NotificationInfo.CheckSaveListener mCheckSaveListener;
     private OnSettingsClickListener mOnSettingsClickListener;
-    private String mKeyToRemoveOnGutsClosed;
+    @VisibleForTesting
+    protected String mKeyToRemoveOnGutsClosed;
 
     public NotificationGutsManager(Context context) {
         mContext = context;
@@ -91,24 +92,15 @@
     }
 
     public void setUpWithPresenter(NotificationPresenter presenter,
-            NotificationEntryManager entryManager, NotificationListContainer listContainer,
+            NotificationListContainer listContainer,
             NotificationInfo.CheckSaveListener checkSaveListener,
             OnSettingsClickListener onSettingsClickListener) {
         mPresenter = presenter;
-        mEntryManager = entryManager;
         mListContainer = listContainer;
         mCheckSaveListener = checkSaveListener;
         mOnSettingsClickListener = onSettingsClickListener;
     }
 
-    public String getKeyToRemoveOnGutsClosed() {
-        return mKeyToRemoveOnGutsClosed;
-    }
-
-    public void setKeyToRemoveOnGutsClosed(String keyToRemoveOnGutsClosed) {
-        mKeyToRemoveOnGutsClosed = keyToRemoveOnGutsClosed;
-    }
-
     public void onDensityOrFontScaleChanged(ExpandableNotificationRow row) {
         setExposedGuts(row.getGuts());
         bindGuts(row);
@@ -171,7 +163,9 @@
             String key = sbn.getKey();
             if (key.equals(mKeyToRemoveOnGutsClosed)) {
                 mKeyToRemoveOnGutsClosed = null;
-                mEntryManager.removeNotification(key, mEntryManager.getLatestRankingMap());
+                if (mNotificationLifetimeFinishedCallback != null) {
+                    mNotificationLifetimeFinishedCallback.onSafeToRemove(key);
+                }
             }
         });
 
@@ -410,6 +404,37 @@
     }
 
     @Override
+    public void setCallback(NotificationSafeToRemoveCallback callback) {
+        mNotificationLifetimeFinishedCallback = callback;
+    }
+
+    @Override
+    public boolean shouldExtendLifetime(NotificationData.Entry entry) {
+        return entry != null
+                &&(mNotificationGutsExposed != null
+                    && entry.row.getGuts() != null
+                    && mNotificationGutsExposed == entry.row.getGuts()
+                    && !mNotificationGutsExposed.isLeavebehind());
+    }
+
+    @Override
+    public void setShouldExtendLifetime(NotificationData.Entry entry, boolean shouldExtend) {
+        if (shouldExtend) {
+            mKeyToRemoveOnGutsClosed = entry.key;
+            if (Log.isLoggable(TAG, Log.DEBUG)) {
+                Log.d(TAG, "Keeping notification because it's showing guts. " + entry.key);
+            }
+        } else {
+            if (mKeyToRemoveOnGutsClosed != null && mKeyToRemoveOnGutsClosed.equals(entry.key)) {
+                mKeyToRemoveOnGutsClosed = null;
+                if (Log.isLoggable(TAG, Log.DEBUG)) {
+                    Log.d(TAG, "Notification that was kept for guts was updated. " + entry.key);
+                }
+            }
+        }
+    }
+
+    @Override
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("NotificationGutsManager state:");
         pw.print("  mKeyToRemoveOnGutsClosed: ");
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
index 6150c2f..4a05989 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
@@ -258,18 +258,6 @@
         return mTrackingHeadsUp;
     }
 
-    /**
-     * React to the removal of the notification in the heads up.
-     *
-     * @return true if the notification was removed and false if it still needs to be kept around
-     * for a bit since it wasn't shown long enough
-     */
-    @Override
-    public boolean removeNotification(@NonNull String key, boolean releaseImmediately) {
-        return super.removeNotification(key, canRemoveImmediately(key)
-                || releaseImmediately);
-    }
-
     @Override
     public void snooze() {
         super.snooze();
@@ -405,7 +393,8 @@
         return (HeadsUpEntryPhone) getTopHeadsUpEntry();
     }
 
-    private boolean canRemoveImmediately(@NonNull String key) {
+    @Override
+    protected boolean canRemoveImmediately(@NonNull String key) {
         if (mSwipedOutKeys.contains(key)) {
             // We always instantly dismiss views being manually swiped out.
             mSwipedOutKeys.remove(key);
@@ -414,7 +403,8 @@
 
         HeadsUpEntryPhone headsUpEntry = getHeadsUpEntryPhone(key);
         HeadsUpEntryPhone topEntry = getTopHeadsUpEntryPhone();
-        return headsUpEntry != topEntry || headsUpEntry.wasShownLongEnough();
+
+        return headsUpEntry == null || headsUpEntry != topEntry || super.canRemoveImmediately(key);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index 22a727c..9beaa10 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -800,7 +800,7 @@
                 this,
                 mNotificationPanel,
                 notifListContainer);
-        mGutsManager.setUpWithPresenter(this, mEntryManager, notifListContainer, mCheckSaveListener,
+        mGutsManager.setUpWithPresenter(this, notifListContainer, mCheckSaveListener,
                 key -> {
                     try {
                         mBarService.onNotificationSettingsViewed(key);
@@ -1811,7 +1811,7 @@
                         mStatusBarWindowController.setHeadsUpShowing(false);
                         mHeadsUpManager.setHeadsUpGoingAway(false);
                     }
-                    mRemoteInputManager.removeRemoteInputEntriesKeptUntilCollapsed();
+                    mRemoteInputManager.onPanelCollapsed();
                 });
             }
         }
@@ -1828,7 +1828,7 @@
 
     @Override
     public void onHeadsUpStateChanged(Entry entry, boolean isHeadsUp) {
-        mEntryManager.onHeadsUpStateChanged(entry, isHeadsUp);
+        mEntryManager.updateNotificationRanking(null /* rankingMap */);
 
         if (isHeadsUp) {
             mDozeServiceHost.fireNotificationHeadsUp();
@@ -1858,7 +1858,7 @@
         }
 
         if (!isExpanded) {
-            mRemoteInputManager.removeRemoteInputEntriesKeptUntilCollapsed();
+            mRemoteInputManager.onPanelCollapsed();
         }
     }
 
@@ -3751,7 +3751,7 @@
             clearNotificationEffects();
         }
         if (newState == StatusBarState.KEYGUARD) {
-            mRemoteInputManager.removeRemoteInputEntriesKeptUntilCollapsed();
+            mRemoteInputManager.onPanelCollapsed();
             maybeEscalateHeadsUp();
         }
     }
@@ -4862,7 +4862,8 @@
                     removeNotification(parentToCancelFinal);
                 }
                 if (shouldAutoCancel(sbn)
-                        || mEntryManager.isNotificationKeptForRemoteInput(notificationKey)) {
+                        || mRemoteInputManager.isNotificationKeptForRemoteInputHistory(
+                                notificationKey)) {
                     // Automatically remove all notifications that we may have kept around longer
                     removeNotification(sbn);
                 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/AlertingNotificationManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/AlertingNotificationManagerTest.java
index f04a115..32c972c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/AlertingNotificationManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/AlertingNotificationManagerTest.java
@@ -91,7 +91,7 @@
         return new TestableAlertingNotificationManager();
     }
 
-    private StatusBarNotification createNewNotification(int id) {
+    protected StatusBarNotification createNewNotification(int id) {
         Notification.Builder n = new Notification.Builder(mContext, "")
                 .setSmallIcon(R.drawable.ic_person)
                 .setContentTitle("Title")
@@ -154,7 +154,7 @@
     public void testRemoveNotification_forceRemove() {
         mAlertingNotificationManager.showNotification(mEntry);
 
-        //Remove forcibly with releaseImmediately = true.
+        // Remove forcibly with releaseImmediately = true.
         mAlertingNotificationManager.removeNotification(mEntry.key, true /* releaseImmediately */);
 
         assertFalse(mAlertingNotificationManager.contains(mEntry.key));
@@ -173,4 +173,30 @@
 
         assertEquals(0, mAlertingNotificationManager.getAllEntries().count());
     }
+
+    @Test
+    public void testShouldExtendLifetime_notShownLongEnough() {
+        mAlertingNotificationManager.showNotification(mEntry);
+
+        // The entry has just been added so the lifetime should be extended
+        assertTrue(mAlertingNotificationManager.shouldExtendLifetime(mEntry));
+    }
+
+    @Test
+    public void testSetShouldExtendLifetime_setShouldExtend() {
+        mAlertingNotificationManager.showNotification(mEntry);
+
+        mAlertingNotificationManager.setShouldExtendLifetime(mEntry, true /* shouldExtend */);
+
+        assertTrue(mAlertingNotificationManager.mExtendedLifetimeAlertEntries.contains(mEntry));
+    }
+
+    @Test
+    public void testSetShouldExtendLifetime_setShouldNotExtend() {
+        mAlertingNotificationManager.mExtendedLifetimeAlertEntries.add(mEntry);
+
+        mAlertingNotificationManager.setShouldExtendLifetime(mEntry, false /* shouldExtend */);
+
+        assertFalse(mAlertingNotificationManager.mExtendedLifetimeAlertEntries.contains(mEntry));
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
index 8129b01..09c1931 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NonPhoneDependencyTest.java
@@ -86,8 +86,8 @@
 
         entryManager.setUpWithPresenter(mPresenter, mListContainer, mEntryManagerCallback,
                 mHeadsUpManager);
-        gutsManager.setUpWithPresenter(mPresenter, entryManager, mListContainer,
-                mCheckSaveListener, mOnClickListener);
+        gutsManager.setUpWithPresenter(mPresenter, mListContainer, mCheckSaveListener,
+                mOnClickListener);
         notificationLogger.setUpWithEntryManager(entryManager, mListContainer);
         mediaManager.setUpWithPresenter(mPresenter, entryManager);
         remoteInputManager.setUpWithPresenter(mPresenter, entryManager, mRemoteInputManagerCallback,
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 3cafaf4..7b0c0a0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
@@ -85,13 +85,6 @@
     }
 
     @Test
-    public void testPostNotificationRemovesKeyKeptForRemoteInput() {
-        mListener.onNotificationPosted(mSbn, mRanking);
-        TestableLooper.get(this).processAllMessages();
-        verify(mEntryManager).removeKeyKeptForRemoteInput(mSbn.getKey());
-    }
-
-    @Test
     public void testNotificationUpdateCallsUpdateNotification() {
         when(mNotificationData.get(mSbn.getKey())).thenReturn(new NotificationData.Entry(mSbn));
         mListener.onNotificationPosted(mSbn, mRanking);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
index afe2cf6..b2493b3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
@@ -1,8 +1,10 @@
 package com.android.systemui.statusbar;
 
-import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertTrue;
+import static junit.framework.Assert.assertFalse;
 
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -10,6 +12,7 @@
 import android.content.Context;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.SystemClock;
 import android.os.UserHandle;
 import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
@@ -22,8 +25,14 @@
 import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputActiveExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputHistoryExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.SmartReplyHistoryExtender;
+
 import com.google.android.collect.Sets;
 
+import junit.framework.Assert;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -41,6 +50,7 @@
     @Mock private RemoteInputController.Delegate mDelegate;
     @Mock private NotificationRemoteInputManager.Callback mCallback;
     @Mock private RemoteInputController mController;
+    @Mock private SmartReplyController mSmartReplyController;
     @Mock private NotificationListenerService.RankingMap mRanking;
     @Mock private ExpandableNotificationRow mRow;
 
@@ -51,6 +61,9 @@
     private TestableNotificationRemoteInputManager mRemoteInputManager;
     private StatusBarNotification mSbn;
     private NotificationData.Entry mEntry;
+    private RemoteInputHistoryExtender mRemoteInputHistoryExtender;
+    private SmartReplyHistoryExtender mSmartReplyHistoryExtender;
+    private RemoteInputActiveExtender mRemoteInputActiveExtender;
 
     @Before
     public void setUp() {
@@ -58,9 +71,9 @@
         mDependency.injectTestDependency(NotificationEntryManager.class, mEntryManager);
         mDependency.injectTestDependency(NotificationLockscreenUserManager.class,
                 mLockscreenUserManager);
+        mDependency.injectTestDependency(SmartReplyController.class, mSmartReplyController);
 
         when(mPresenter.getHandler()).thenReturn(Handler.createAsync(Looper.myLooper()));
-        when(mEntryManager.getLatestRankingMap()).thenReturn(mRanking);
 
         mRemoteInputManager = new TestableNotificationRemoteInputManager(mContext);
         mSbn = new StatusBarNotification(TEST_PACKAGE_NAME, TEST_PACKAGE_NAME, 0, null, TEST_UID,
@@ -70,20 +83,10 @@
 
         mRemoteInputManager.setUpWithPresenterForTest(mPresenter, mEntryManager, mCallback,
                 mDelegate, mController);
-    }
-
-    @Test
-    public void testOnRemoveNotificationNotKept() {
-        assertFalse(mRemoteInputManager.onRemoveNotification(mEntry));
-        assertTrue(mRemoteInputManager.getRemoteInputEntriesToRemoveOnCollapse().isEmpty());
-    }
-
-    @Test
-    public void testOnRemoveNotificationKept() {
-        when(mController.isRemoteInputActive(mEntry)).thenReturn(true);
-        assertTrue(mRemoteInputManager.onRemoveNotification(mEntry));
-        assertTrue(mRemoteInputManager.getRemoteInputEntriesToRemoveOnCollapse().equals(
-                Sets.newArraySet(mEntry)));
+        for (NotificationLifetimeExtender extender : mRemoteInputManager.getLifetimeExtenders()) {
+            extender.setCallback(
+                    mock(NotificationLifetimeExtender.NotificationSafeToRemoveCallback.class));
+        }
     }
 
     @Test
@@ -95,15 +98,104 @@
     }
 
     @Test
-    public void testRemoveRemoteInputEntriesKeptUntilCollapsed() {
-        mRemoteInputManager.getRemoteInputEntriesToRemoveOnCollapse().add(mEntry);
-        mRemoteInputManager.removeRemoteInputEntriesKeptUntilCollapsed();
+    public void testShouldExtendLifetime_remoteInputActive() {
+        when(mController.isRemoteInputActive(mEntry)).thenReturn(true);
 
-        assertTrue(mRemoteInputManager.getRemoteInputEntriesToRemoveOnCollapse().isEmpty());
-        verify(mController).removeRemoteInput(mEntry, null);
-        verify(mEntryManager).removeNotification(mEntry.key, mRanking);
+        assertTrue(mRemoteInputActiveExtender.shouldExtendLifetime(mEntry));
     }
 
+    @Test
+    public void testShouldExtendLifetime_isSpinning() {
+        NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
+        when(mController.isSpinning(mEntry.key)).thenReturn(true);
+
+        assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
+    }
+
+    @Test
+    public void testShouldExtendLifetime_recentRemoteInput() {
+        NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
+        mEntry.lastRemoteInputSent = SystemClock.elapsedRealtime();
+
+        assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
+    }
+
+    @Test
+    public void testShouldExtendLifetime_smartReplySending() {
+        NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
+        when(mSmartReplyController.isSendingSmartReply(mEntry.key)).thenReturn(true);
+
+        assertTrue(mSmartReplyHistoryExtender.shouldExtendLifetime(mEntry));
+    }
+
+    @Test
+    public void testNotificationWithRemoteInputActiveIsRemovedOnCollapse() {
+        mRemoteInputActiveExtender.setShouldExtendLifetime(mEntry, true);
+
+        assertEquals(mRemoteInputManager.getEntriesKeptForRemoteInputActive(),
+                Sets.newArraySet(mEntry));
+
+        mRemoteInputManager.onPanelCollapsed();
+
+        assertTrue(mRemoteInputManager.getEntriesKeptForRemoteInputActive().isEmpty());
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInputNoSpinner() {
+        StatusBarNotification newSbn =
+                mRemoteInputManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        assertEquals(1, messages.length);
+        assertEquals("A Reply", messages[0]);
+        assertFalse(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+        assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInputWithSpinner() {
+        StatusBarNotification newSbn =
+                mRemoteInputManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", true);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        assertEquals(1, messages.length);
+        assertEquals("A Reply", messages[0]);
+        assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+        assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_withExistingInput() {
+        // Setup a notification entry with 1 remote input.
+        StatusBarNotification newSbn =
+                mRemoteInputManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
+        NotificationData.Entry entry = new NotificationData.Entry(newSbn);
+
+        // Try rebuilding to add another reply.
+        newSbn = mRemoteInputManager.rebuildNotificationWithRemoteInput(entry, "Reply 2", true);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        assertEquals(2, messages.length);
+        assertEquals("Reply 2", messages[0]);
+        assertEquals("A Reply", messages[1]);
+    }
+
+    @Test
+    public void testRebuildNotificationForCanceledSmartReplies() {
+        // Try rebuilding to remove spinner and hide buttons.
+        StatusBarNotification newSbn =
+                mRemoteInputManager.rebuildNotificationForCanceledSmartReplies(mEntry);
+        assertFalse(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+        assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
+    }
+
+
     private class TestableNotificationRemoteInputManager extends NotificationRemoteInputManager {
 
         public TestableNotificationRemoteInputManager(Context context) {
@@ -118,5 +210,15 @@
             super.setUpWithPresenter(presenter, entryManager, callback, delegate);
             mRemoteInputController = controller;
         }
+
+        @Override
+        protected void addLifetimeExtenders() {
+            mRemoteInputActiveExtender = new RemoteInputActiveExtender();
+            mRemoteInputHistoryExtender = new RemoteInputHistoryExtender();
+            mSmartReplyHistoryExtender = new SmartReplyHistoryExtender();
+            mLifetimeExtenders.add(mRemoteInputHistoryExtender);
+            mLifetimeExtenders.add(mSmartReplyHistoryExtender);
+            mLifetimeExtenders.add(mRemoteInputActiveExtender);
+        }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
index ada5785..17daaac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
@@ -23,8 +23,11 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.ActivityManager;
 import android.app.Notification;
 import android.os.RemoteException;
+import android.os.UserHandle;
+import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
@@ -35,6 +38,7 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -46,97 +50,88 @@
 @TestableLooper.RunWithLooper
 @SmallTest
 public class SmartReplyControllerTest extends SysuiTestCase {
-    private static final String TEST_NOTIFICATION_KEY = "akey";
+    private static final String TEST_PACKAGE_NAME = "test";
+    private static final int TEST_UID = 0;
     private static final String TEST_CHOICE_TEXT = "A Reply";
     private static final int TEST_CHOICE_INDEX = 2;
     private static final int TEST_CHOICE_COUNT = 4;
 
     private Notification mNotification;
     private NotificationData.Entry mEntry;
+    private SmartReplyController mSmartReplyController;
+    private NotificationRemoteInputManager mRemoteInputManager;
 
-    @Mock
-    private NotificationEntryManager mNotificationEntryManager;
-    @Mock
-    private IStatusBarService mIStatusBarService;
+    @Mock private NotificationPresenter mPresenter;
+    @Mock private RemoteInputController.Delegate mDelegate;
+    @Mock private NotificationRemoteInputManager.Callback mCallback;
+    @Mock private StatusBarNotification mSbn;
+    @Mock private NotificationEntryManager mNotificationEntryManager;
+    @Mock private IStatusBarService mIStatusBarService;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-
         mDependency.injectTestDependency(NotificationEntryManager.class,
                 mNotificationEntryManager);
         mDependency.injectTestDependency(IStatusBarService.class, mIStatusBarService);
 
+        mSmartReplyController = new SmartReplyController();
+        mDependency.injectTestDependency(SmartReplyController.class,
+                mSmartReplyController);
+
+        mRemoteInputManager = new NotificationRemoteInputManager(mContext);
+        mRemoteInputManager.setUpWithPresenter(mPresenter, mNotificationEntryManager, mCallback,
+                mDelegate);
         mNotification = new Notification.Builder(mContext, "")
                 .setSmallIcon(R.drawable.ic_person)
                 .setContentTitle("Title")
                 .setContentText("Text").build();
-        StatusBarNotification sbn = mock(StatusBarNotification.class);
-        when(sbn.getNotification()).thenReturn(mNotification);
-        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
-        mEntry = new NotificationData.Entry(sbn);
+
+        mSbn = new StatusBarNotification(TEST_PACKAGE_NAME, TEST_PACKAGE_NAME, 0, null, TEST_UID,
+                0, mNotification, new UserHandle(ActivityManager.getCurrentUser()), null, 0);
+        mEntry = new NotificationData.Entry(mSbn);
     }
 
     @Test
     public void testSendSmartReply_updatesRemoteInput() {
-        StatusBarNotification sbn = mock(StatusBarNotification.class);
-        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
-        when(mNotificationEntryManager.rebuildNotificationWithRemoteInput(
-                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
-                eq(TEST_CHOICE_TEXT), eq(true))).thenReturn(sbn);
-
-        SmartReplyController controller = new SmartReplyController();
-        controller.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
 
         // Sending smart reply should make calls to NotificationEntryManager
         // to update the notification with reply and spinner.
-        verify(mNotificationEntryManager).rebuildNotificationWithRemoteInput(
-                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
-                eq(TEST_CHOICE_TEXT), eq(true));
         verify(mNotificationEntryManager).updateNotification(
-                argThat(sbn2 -> sbn2.getKey().equals(TEST_NOTIFICATION_KEY)), isNull());
+                argThat(sbn -> sbn.getKey().equals(mSbn.getKey())), isNull());
     }
 
     @Test
     public void testSendSmartReply_logsToStatusBar() throws RemoteException {
-        StatusBarNotification sbn = mock(StatusBarNotification.class);
-        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
-        when(mNotificationEntryManager.rebuildNotificationWithRemoteInput(
-                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
-                eq(TEST_CHOICE_TEXT), eq(true))).thenReturn(sbn);
-
-        SmartReplyController controller = new SmartReplyController();
-        controller.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
 
         // Check we log the result to the status bar service.
-        verify(mIStatusBarService).onNotificationSmartReplySent(TEST_NOTIFICATION_KEY,
+        verify(mIStatusBarService).onNotificationSmartReplySent(mSbn.getKey(),
                 TEST_CHOICE_INDEX);
     }
 
     @Test
     public void testShowSmartReply_logsToStatusBar() throws RemoteException {
-        SmartReplyController controller = new SmartReplyController();
-        controller.smartRepliesAdded(mEntry, TEST_CHOICE_COUNT);
+        mSmartReplyController.smartRepliesAdded(mEntry, TEST_CHOICE_COUNT);
 
         // Check we log the result to the status bar service.
-        verify(mIStatusBarService).onNotificationSmartRepliesAdded(TEST_NOTIFICATION_KEY,
+        verify(mIStatusBarService).onNotificationSmartRepliesAdded(mSbn.getKey(),
                 TEST_CHOICE_COUNT);
     }
 
     @Test
     public void testSendSmartReply_reportsSending() {
-        SmartReplyController controller = new SmartReplyController();
-        controller.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
 
-        assertTrue(controller.isSendingSmartReply(TEST_NOTIFICATION_KEY));
+        assertTrue(mSmartReplyController.isSendingSmartReply(mSbn.getKey()));
     }
 
     @Test
     public void testSendingSmartReply_afterRemove_shouldReturnFalse() {
-        SmartReplyController controller = new SmartReplyController();
-        controller.isSendingSmartReply(TEST_NOTIFICATION_KEY);
-        controller.stopSending(mEntry);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.stopSending(mEntry);
 
-        assertFalse(controller.isSendingSmartReply(TEST_NOTIFICATION_KEY));
+        assertFalse(mSmartReplyController.isSendingSmartReply(mSbn.getKey()));
     }
 }
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 6543bdb..dacf59c 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
@@ -57,6 +57,7 @@
 import com.android.systemui.ForegroundServiceController;
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.statusbar.NotificationLifetimeExtender;
 import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationLockscreenUserManager;
 import com.android.systemui.statusbar.NotificationMediaManager;
@@ -84,7 +85,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -143,6 +143,10 @@
         public CountDownLatch getCountDownLatch() {
             return mCountDownLatch;
         }
+
+        public ArrayList<NotificationLifetimeExtender> getLifetimeExtenders() {
+            return mNotificationLifetimeExtenders;
+        }
     }
 
     private void setUserSentiment(String key, int sentiment) {
@@ -279,7 +283,6 @@
         verify(mBarService, never()).onNotificationError(any(), any(), anyInt(), anyInt(), anyInt(),
                 any(), anyInt());
 
-        verify(mRemoteInputManager).onUpdateNotification(mEntry);
         verify(mPresenter).updateNotificationViews();
         verify(mForegroundServiceController).updateNotification(eq(mSbn), anyInt());
         verify(mCallback).onNotificationUpdated(mSbn);
@@ -301,8 +304,6 @@
                 any(), anyInt());
 
         verify(mMediaManager).onNotificationRemoved(mSbn.getKey());
-        verify(mRemoteInputManager).onRemoveNotification(mEntry);
-        verify(mSmartReplyController).stopSending(mEntry);
         verify(mForegroundServiceController).removeNotification(mSbn);
         verify(mListContainer).cleanUpViewState(mRow);
         verify(mPresenter).updateNotificationViews();
@@ -313,17 +314,23 @@
     }
 
     @Test
-    public void testRemoveNotification_blockedBySendingSmartReply() throws Exception {
+    public void testRemoveNotification_blockedByLifetimeExtender() {
         com.android.systemui.util.Assert.isNotMainThread();
 
+        NotificationLifetimeExtender extender = mock(NotificationLifetimeExtender.class);
+        when(extender.shouldExtendLifetime(mEntry)).thenReturn(true);
+
+        ArrayList<NotificationLifetimeExtender> extenders = mEntryManager.getLifetimeExtenders();
+        extenders.clear();
+        extenders.add(extender);
+
         mEntry.row = mRow;
         mEntryManager.getNotificationData().add(mEntry);
-        when(mSmartReplyController.isSendingSmartReply(mEntry.key)).thenReturn(true);
 
         mEntryManager.removeNotification(mSbn.getKey(), mRankingMap);
 
         assertNotNull(mEntryManager.getNotificationData().get(mSbn.getKey()));
-        assertTrue(mEntryManager.isNotificationKeptForRemoteInput(mEntry.key));
+        verify(extender).setShouldExtendLifetime(mEntry, true);
     }
 
     @Test
@@ -411,61 +418,6 @@
     }
 
     @Test
-    public void testRebuildWithRemoteInput_noExistingInputNoSpinner() {
-        StatusBarNotification newSbn =
-                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
-        CharSequence[] messages = newSbn.getNotification().extras
-                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
-        Assert.assertEquals(1, messages.length);
-        Assert.assertEquals("A Reply", messages[0]);
-        Assert.assertFalse(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
-        Assert.assertTrue(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
-    }
-
-    @Test
-    public void testRebuildWithRemoteInput_noExistingInputWithSpinner() {
-        StatusBarNotification newSbn =
-                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", true);
-        CharSequence[] messages = newSbn.getNotification().extras
-                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
-        Assert.assertEquals(1, messages.length);
-        Assert.assertEquals("A Reply", messages[0]);
-        Assert.assertTrue(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
-        Assert.assertTrue(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
-    }
-
-    @Test
-    public void testRebuildWithRemoteInput_withExistingInput() {
-        // Setup a notification entry with 1 remote input.
-        StatusBarNotification newSbn =
-                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
-        NotificationData.Entry entry = new NotificationData.Entry(newSbn);
-
-        // Try rebuilding to add another reply.
-        newSbn = mEntryManager.rebuildNotificationWithRemoteInput(entry, "Reply 2", true);
-        CharSequence[] messages = newSbn.getNotification().extras
-                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
-        Assert.assertEquals(2, messages.length);
-        Assert.assertEquals("Reply 2", messages[0]);
-        Assert.assertEquals("A Reply", messages[1]);
-    }
-
-    @Test
-    public void testRebuildNotificationForCanceledSmartReplies() {
-        // Try rebuilding to remove spinner and hide buttons.
-        StatusBarNotification newSbn =
-                mEntryManager.rebuildNotificationForCanceledSmartReplies(mEntry);
-        Assert.assertFalse(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
-        Assert.assertTrue(newSbn.getNotification().extras
-                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
-    }
-
-    @Test
     public void testUpdateNotificationRanking() {
         when(mPresenter.isDeviceProvisioned()).thenReturn(true);
         when(mPresenter.isNotificationForCurrentProfiles(any())).thenReturn(true);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
index 676cb61..6656fdd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsManagerTest.java
@@ -22,6 +22,8 @@
 
 import static android.service.notification.NotificationListenerService.Ranking.USER_SENTIMENT_NEGATIVE;
 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.junit.Assert.fail;
@@ -30,6 +32,7 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.spy;
@@ -54,6 +57,7 @@
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
+import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.NotificationPresenter;
 import com.android.systemui.statusbar.NotificationTestHelper;
@@ -99,7 +103,7 @@
         mHelper = new NotificationTestHelper(mContext);
 
         mGutsManager = new NotificationGutsManager(mContext);
-        mGutsManager.setUpWithPresenter(mPresenter, mEntryManager, mStackScroller,
+        mGutsManager.setUpWithPresenter(mPresenter, mStackScroller,
                 mCheckSaveListener, mOnSettingsClickListener);
     }
 
@@ -346,6 +350,35 @@
                 eq(true) /* isUserSentimentNegative */);
     }
 
+    @Test
+    public void testShouldExtendLifetime() {
+        NotificationGuts guts = new NotificationGuts(mContext);
+        ExpandableNotificationRow row = spy(createTestNotificationRow());
+        doReturn(guts).when(row).getGuts();
+        NotificationData.Entry entry = row.getEntry();
+        entry.row = row;
+        mGutsManager.setExposedGuts(guts);
+
+        assertTrue(mGutsManager.shouldExtendLifetime(entry));
+    }
+
+    @Test
+    public void testSetShouldExtendLifetime_setShouldExtend() {
+        NotificationData.Entry entry = createTestNotificationRow().getEntry();
+        mGutsManager.setShouldExtendLifetime(entry, true /* shouldExtend */);
+
+        assertTrue(entry.key.equals(mGutsManager.mKeyToRemoveOnGutsClosed));
+    }
+
+    @Test
+    public void testSetShouldExtendLifetime_setShouldNotExtend() {
+        NotificationData.Entry entry = createTestNotificationRow().getEntry();
+        mGutsManager.mKeyToRemoveOnGutsClosed = entry.key;
+        mGutsManager.setShouldExtendLifetime(entry, false /* shouldExtend */);
+
+        assertNull(mGutsManager.mKeyToRemoveOnGutsClosed);
+    }
+
     ////////////////////////////////////////////////////////////////////////////////////////////////
     // Utility methods:
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
index c19188c..ce0bd58 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
@@ -112,7 +112,8 @@
 
         mEntryManager = new TestableNotificationEntryManager(mSystemServicesProxy, mPowerManager,
                 mContext);
-        mEntryManager.setUpForTest(mock(NotificationPresenter.class), null, null, null, mNotificationData);
+        mEntryManager.setUpForTest(mock(NotificationPresenter.class), null, null, mHeadsUpManager,
+                mNotificationData);
         mDependency.injectTestDependency(NotificationEntryManager.class, mEntryManager);
 
         NotificationShelf notificationShelf = spy(new NotificationShelf(getContext(), null));
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhoneTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhoneTest.java
index bdf7cd3..a81d17f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhoneTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhoneTest.java
@@ -23,6 +23,7 @@
 
 import com.android.systemui.statusbar.AlertingNotificationManager;
 import com.android.systemui.statusbar.AlertingNotificationManagerTest;
+import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.VisualStabilityManager;
 
 import org.junit.Before;
@@ -83,4 +84,26 @@
 
         assertFalse(mHeadsUpManager.contains(mEntry.key));
     }
+
+    @Test
+    public void testShouldExtendLifetime_swipedOut() {
+        mHeadsUpManager.showNotification(mEntry);
+        mHeadsUpManager.addSwipedOutNotification(mEntry.key);
+
+        // Notification is swiped so its lifetime should not be extended even if it hasn't been
+        // shown long enough
+        assertFalse(mHeadsUpManager.shouldExtendLifetime(mEntry));
+    }
+
+    @Test
+    public void testShouldExtendLifetime_notTopEntry() {
+        NotificationData.Entry laterEntry = new NotificationData.Entry(createNewNotification(1));
+        laterEntry.row = mRow;
+        mHeadsUpManager.showNotification(mEntry);
+        mHeadsUpManager.showNotification(laterEntry);
+
+        // Notification is "behind" a higher priority notification so we have no reason to keep
+        // its lifetime extended
+        assertFalse(mHeadsUpManager.shouldExtendLifetime(mEntry));
+    }
 }