New Pipeline: Remote Input 1/4: Extract legacy pipeline logic within NotificationRemoteInputManager

Fixes: 204127880
Bug: 203938360
Test: atest NotificationRemoteInputManagerTest SmartReplyControllerTest
Merged-In: I5424d3525601fda2b217c3d13130769106d4f50d
Change-Id: I5424d3525601fda2b217c3d13130769106d4f50d

packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java

@@ -15,8 +15,6 @@
  */
 package com.android.systemui.statusbar;
 
-import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
 import android.app.KeyguardManager;
@@ -48,6 +46,9 @@ import android.widget.RemoteViews;
 import android.widget.RemoteViews.InteractionHandler;
 import android.widget.TextView;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.internal.statusbar.NotificationVisibility;
@@ -55,10 +56,12 @@ import com.android.systemui.Dumpable;
 import com.android.systemui.R;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.dagger.StatusBarDependenciesModule;
 import com.android.systemui.statusbar.notification.NotificationEntryListener;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
+import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry.EditedSuggestionInfo;
 import com.android.systemui.statusbar.notification.logging.NotificationLogger;
@@ -93,27 +96,7 @@ public class NotificationRemoteInputManager implements Dumpable {
     private static final boolean DEBUG = false;
     private static final String TAG = "NotifRemoteInputManager";
 
-    /**
+    private RemoteInputListener mRemoteInputListener;
-     * How long to wait before auto-dismissing a notification that was kept for remote input, and
-     * has now sent a remote input. We auto-dismiss, because the app may not see a reason to cancel
-     * these given that they technically don't exist anymore. We wait a bit in case the app issues
-     * an update.
-     */
-    private static final int REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY = 200;
-
-    /**
-     * 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<NotificationEntry> mEntriesKeptForRemoteInputActive =
-            new ArraySet<>();
 
     // Dependencies:
     private final NotificationLockscreenUserManager mLockscreenUserManager;
@@ -125,6 +108,7 @@ public class NotificationRemoteInputManager implements Dumpable {
     private final Lazy<Optional<StatusBar>> mStatusBarOptionalLazy;
 
     protected final Context mContext;
+    protected final FeatureFlags mFeatureFlags;
     private final UserManager mUserManager;
     private final KeyguardManager mKeyguardManager;
     private final StatusBarStateController mStatusBarStateController;
@@ -132,11 +116,8 @@ public class NotificationRemoteInputManager implements Dumpable {
     private final NotificationClickNotifier mClickNotifier;
 
     protected RemoteInputController mRemoteInputController;
-    protected NotificationLifetimeExtender.NotificationSafeToRemoveCallback
-            mNotificationLifetimeFinishedCallback;
     protected IStatusBarService mBarService;
     protected Callback mCallback;
-    protected final ArrayList<NotificationLifetimeExtender> mLifetimeExtenders = new ArrayList<>();
 
     private final List<RemoteInputController.Callback> mControllerCallbacks = new ArrayList<>();
 
@@ -226,6 +207,7 @@ public class NotificationRemoteInputManager implements Dumpable {
                 ViewGroup actionGroup = (ViewGroup) parent;
                 buttonIndex = actionGroup.indexOfChild(view);
             }
+            // FIXME: get this for the new pipeline!
             final int count = mEntryManager.getActiveNotificationsCount();
             final int rank = entry.getRanking().getRank();
 
@@ -283,9 +265,11 @@ public class NotificationRemoteInputManager implements Dumpable {
      */
     public NotificationRemoteInputManager(
             Context context,
+            FeatureFlags featureFlags,
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
+            NotifPipeline notifPipeline,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             @Main Handler mainHandler,
@@ -294,6 +278,7 @@ public class NotificationRemoteInputManager implements Dumpable {
             ActionClickLogger logger,
             DumpManager dumpManager) {
         mContext = context;
+        mFeatureFlags = featureFlags;
         mLockscreenUserManager = lockscreenUserManager;
         mSmartReplyController = smartReplyController;
         mEntryManager = notificationEntryManager;
@@ -303,7 +288,10 @@ public class NotificationRemoteInputManager implements Dumpable {
         mBarService = IStatusBarService.Stub.asInterface(
                 ServiceManager.getService(Context.STATUS_BAR_SERVICE));
         mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
-        addLifetimeExtenders();
+        if (!featureFlags.isNewNotifPipelineRenderingEnabled()) {
+            mRemoteInputListener = createLegacyRemoteInputLifetimeExtender(mainHandler,
+                    notificationEntryManager, smartReplyController);
+        }
         mKeyguardManager = context.getSystemService(KeyguardManager.class);
         mStatusBarStateController = statusBarStateController;
         mRemoteInputUriController = remoteInputUriController;
@@ -335,10 +323,20 @@ public class NotificationRemoteInputManager implements Dumpable {
         });
     }
 
+    @NonNull
+    @VisibleForTesting
+    protected LegacyRemoteInputLifetimeExtender createLegacyRemoteInputLifetimeExtender(
+            Handler mainHandler,
+            NotificationEntryManager notificationEntryManager,
+            SmartReplyController smartReplyController) {
+        return new LegacyRemoteInputLifetimeExtender();
+    }
+
     /** Initializes this component with the provided dependencies. */
     public void setUpWithCallback(Callback callback, RemoteInputController.Delegate delegate) {
         mCallback = callback;
         mRemoteInputController = new RemoteInputController(delegate, mRemoteInputUriController);
+        mRemoteInputListener.setRemoteInputController(mRemoteInputController);
         // Register all stored callbacks from before the Controller was initialized.
         for (RemoteInputController.Callback cb : mControllerCallbacks) {
             mRemoteInputController.addCallback(cb);
@@ -347,19 +345,8 @@ public class NotificationRemoteInputManager implements Dumpable {
         mRemoteInputController.addCallback(new RemoteInputController.Callback() {
             @Override
             public void onRemoteInputSent(NotificationEntry entry) {
-                if (FORCE_REMOTE_INPUT_HISTORY
+                if (mRemoteInputListener != null) {
-                        && isNotificationKeptForRemoteInputHistory(entry.getKey())) {
+                    mRemoteInputListener.onRemoteInputSent(entry);
-                    mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-                } 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.
-                    mMainHandler.postDelayed(() -> {
-                        if (mEntriesKeptForRemoteInputActive.remove(entry)) {
-                            mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-                        }
-                    }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
                 }
                 try {
                     mBarService.onNotificationDirectReplied(entry.getSbn().getKey());
@@ -381,12 +368,13 @@ public class NotificationRemoteInputManager implements Dumpable {
                 }
             }
         });
-        mSmartReplyController.setCallback((entry, reply) -> {
+        if (!mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
-            StatusBarNotification newSbn =
+            // FIXME: Don't forget to implement this in the coordinator!
-                    rebuildNotificationWithRemoteInputInserted(entry, reply, true /* showSpinner */,
+            mSmartReplyController.setCallback((entry, reply) -> {
-                            null /* mimeType */, null /* uri */);
+                StatusBarNotification newSbn = rebuildNotificationForSendingSmartReply(entry, reply);
-            mEntryManager.updateNotification(newSbn, null /* ranking */);
+                mEntryManager.updateNotification(newSbn, null /* ranking */);
-        });
+            });
+        }
     }
 
     public void addControllerCallback(RemoteInputController.Callback callback) {
@@ -574,51 +562,39 @@ public class NotificationRemoteInputManager implements Dumpable {
         if (v == null) {
             return null;
         }
-        return (RemoteInputView) v.findViewWithTag(RemoteInputView.VIEW_TAG);
+        return v.findViewWithTag(RemoteInputView.VIEW_TAG);
-    }
-
-    /**
-     * 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;
+        // OLD pipeline code ONLY; can assume implementation
+        return ((LegacyRemoteInputLifetimeExtender) mRemoteInputListener).mLifetimeExtenders;
     }
 
     @VisibleForTesting
     void onPerformRemoveNotification(NotificationEntry entry, final String key) {
-        if (mKeysKeptForRemoteInputHistory.contains(key)) {
+        // OLD pipeline code ONLY; can assume implementation
-            mKeysKeptForRemoteInputHistory.remove(key);
+        ((LegacyRemoteInputLifetimeExtender) mRemoteInputListener)
-        }
+                .mKeysKeptForRemoteInputHistory.remove(key);
         if (isRemoteInputActive(entry)) {
             entry.mRemoteEditImeVisible = false;
             mRemoteInputController.removeRemoteInput(entry, null);
         }
     }
 
+    /** Informs the remote input system that the panel has collapsed */
     public void onPanelCollapsed() {
-        for (int i = 0; i < mEntriesKeptForRemoteInputActive.size(); i++) {
+        if (mRemoteInputListener != null) {
-            NotificationEntry entry = mEntriesKeptForRemoteInputActive.valueAt(i);
+            mRemoteInputListener.onPanelCollapsed();
-            if (mRemoteInputController != null) {
-                mRemoteInputController.removeRemoteInput(entry, null);
-            }
-            if (mNotificationLifetimeFinishedCallback != null) {
-                mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-            }
         }
-        mEntriesKeptForRemoteInputActive.clear();
     }
 
+    /** Returns whether the given notification is lifetime extended because of remote input */
     public boolean isNotificationKeptForRemoteInputHistory(String key) {
-        return mKeysKeptForRemoteInputHistory.contains(key);
+        return mRemoteInputListener != null
+                && mRemoteInputListener.isNotificationKeptForRemoteInputHistory(key);
     }
 
+    /** Returns whether the notification should be lifetime extended for remote input history */
     public boolean shouldKeepForRemoteInputHistory(NotificationEntry entry) {
         if (!FORCE_REMOTE_INPUT_HISTORY) {
             return false;
@@ -636,16 +612,12 @@ public class NotificationRemoteInputManager implements Dumpable {
         if (entry == null) {
             return;
         }
-        final String key = entry.getKey();
+        if (mRemoteInputListener != null) {
-        if (isNotificationKeptForRemoteInputHistory(key)) {
+            mRemoteInputListener.releaseNotificationIfKeptForRemoteInputHistory(entry);
-            mMainHandler.postDelayed(() -> {
-                if (isNotificationKeptForRemoteInputHistory(key)) {
-                    mNotificationLifetimeFinishedCallback.onSafeToRemove(key);
-                }
-            }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
         }
     }
 
+    /** Returns whether the notification should be lifetime extended for smart reply history */
     public boolean shouldKeepForSmartReplyHistory(NotificationEntry entry) {
         if (!FORCE_REMOTE_INPUT_HISTORY) {
             return false;
@@ -661,13 +633,36 @@ public class NotificationRemoteInputManager implements Dumpable {
         }
     }
 
-    @VisibleForTesting
+    // FIXME: Move to a helper class and test separately
-    StatusBarNotification rebuildNotificationForCanceledSmartReplies(
+    public StatusBarNotification rebuildNotificationForSendingSmartReply(NotificationEntry entry,
+            CharSequence reply) {
+        return rebuildNotificationWithRemoteInputInserted(entry, reply,
+                true /* showSpinner */,
+                null /* mimeType */, null /* uri */);
+    }
+
+    // FIXME: Move to a helper class and test separately
+    public StatusBarNotification rebuildNotificationForCanceledSmartReplies(
             NotificationEntry entry) {
         return rebuildNotificationWithRemoteInputInserted(entry, null /* remoteInputTest */,
                 false /* showSpinner */, null /* mimeType */, null /* uri */);
     }
 
+    // FIXME: Move to a helper class and test separately
+    public StatusBarNotification rebuildNotificationForBasicExtension(NotificationEntry entry) {
+        CharSequence remoteInputText = entry.remoteInputText;
+        if (TextUtils.isEmpty(remoteInputText)) {
+            remoteInputText = entry.remoteInputTextWhenReset;
+        }
+        String remoteInputMimeType = entry.remoteInputMimeType;
+        Uri remoteInputUri = entry.remoteInputUri;
+        StatusBarNotification newSbn = rebuildNotificationWithRemoteInputInserted(entry,
+                remoteInputText, false /* showSpinner */, remoteInputMimeType,
+                remoteInputUri);
+        return newSbn;
+    }
+
+    // FIXME: Move to a helper class and test separately
     @VisibleForTesting
     StatusBarNotification rebuildNotificationWithRemoteInputInserted(NotificationEntry entry,
             CharSequence remoteInputText, boolean showSpinner, String mimeType, Uri uri) {
@@ -714,11 +709,9 @@ public class NotificationRemoteInputManager implements Dumpable {
 
     @Override
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
-        pw.println("NotificationRemoteInputManager state:");
+        if (mRemoteInputListener instanceof Dumpable) {
-        pw.print("  mKeysKeptForRemoteInputHistory: ");
+            ((Dumpable) mRemoteInputListener).dump(fd, pw, args);
-        pw.println(mKeysKeptForRemoteInputHistory);
+        }
-        pw.print("  mEntriesKeptForRemoteInputActive: ");
-        pw.println(mEntriesKeptForRemoteInputActive);
     }
 
     public void bindRow(ExpandableNotificationRow row) {
@@ -734,11 +727,6 @@ public class NotificationRemoteInputManager implements Dumpable {
         return mInteractionHandler;
     }
 
-    @VisibleForTesting
-    public Set<NotificationEntry> getEntriesKeptForRemoteInputActive() {
-        return mEntriesKeptForRemoteInputActive;
-    }
-
     public boolean isRemoteInputActive() {
         return mRemoteInputController != null && mRemoteInputController.isRemoteInputActive();
     }
@@ -757,131 +745,6 @@ public class NotificationRemoteInputManager implements Dumpable {
         }
     }
 
-    /**
-     * 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 NotificationEntry entry) {
-            return shouldKeepForRemoteInputHistory(entry);
-        }
-
-        @Override
-        public void setShouldManageLifetime(NotificationEntry entry,
-                boolean shouldExtend) {
-            if (shouldExtend) {
-                CharSequence remoteInputText = entry.remoteInputText;
-                if (TextUtils.isEmpty(remoteInputText)) {
-                    remoteInputText = entry.remoteInputTextWhenReset;
-                }
-                String remoteInputMimeType = entry.remoteInputMimeType;
-                Uri remoteInputUri = entry.remoteInputUri;
-                StatusBarNotification newSbn = rebuildNotificationWithRemoteInputInserted(entry,
-                        remoteInputText, false /* showSpinner */, remoteInputMimeType,
-                        remoteInputUri);
-                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.isRemoved()) {
-                    return;
-                }
-
-                if (Log.isLoggable(TAG, Log.DEBUG)) {
-                    Log.d(TAG, "Keeping notification around after sending remote input "
-                            + entry.getKey());
-                }
-
-                mKeysKeptForRemoteInputHistory.add(entry.getKey());
-            } else {
-                mKeysKeptForRemoteInputHistory.remove(entry.getKey());
-            }
-        }
-    }
-
-    /**
-     * 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 NotificationEntry entry) {
-            return shouldKeepForSmartReplyHistory(entry);
-        }
-
-        @Override
-        public void setShouldManageLifetime(NotificationEntry entry,
-                boolean shouldExtend) {
-            if (shouldExtend) {
-                StatusBarNotification newSbn = rebuildNotificationForCanceledSmartReplies(entry);
-
-                if (newSbn == null) {
-                    return;
-                }
-
-                mEntryManager.updateNotification(newSbn, null);
-
-                if (entry.isRemoved()) {
-                    return;
-                }
-
-                if (Log.isLoggable(TAG, Log.DEBUG)) {
-                    Log.d(TAG, "Keeping notification around after sending smart reply "
-                            + entry.getKey());
-                }
-
-                mKeysKeptForRemoteInputHistory.add(entry.getKey());
-            } else {
-                mKeysKeptForRemoteInputHistory.remove(entry.getKey());
-                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 NotificationEntry entry) {
-            return isRemoteInputActive(entry);
-        }
-
-        @Override
-        public void setShouldManageLifetime(NotificationEntry entry,
-                boolean shouldExtend) {
-            if (shouldExtend) {
-                if (Log.isLoggable(TAG, Log.DEBUG)) {
-                    Log.d(TAG, "Keeping notification around while remote input active "
-                            + entry.getKey());
-                }
-                mEntriesKeptForRemoteInputActive.add(entry);
-            } else {
-                mEntriesKeptForRemoteInputActive.remove(entry);
-            }
-        }
-    }
-
     /**
      * Callback for various remote input related events, or for providing information that
      * NotificationRemoteInputManager needs to know to decide what to do.
@@ -975,4 +838,250 @@ public class NotificationRemoteInputManager implements Dumpable {
          */
         boolean showBouncerIfNecessary();
     }
+
+    public interface RemoteInputListener {
+        void onRemoteInputSent(NotificationEntry entry);
+
+        void onPanelCollapsed();
+
+        boolean isNotificationKeptForRemoteInputHistory(String key);
+
+        void releaseNotificationIfKeptForRemoteInputHistory(@NonNull NotificationEntry entry);
+
+        void setRemoteInputController(@NonNull RemoteInputController remoteInputController);
+    }
+
+    @VisibleForTesting
+    protected class LegacyRemoteInputLifetimeExtender implements RemoteInputListener, Dumpable {
+
+        /**
+         * How long to wait before auto-dismissing a notification that was kept for remote input,
+         * and has now sent a remote input. We auto-dismiss, because the app may not see a reason to
+         * cancel these given that they technically don't exist anymore. We wait a bit in case the
+         * app issues an update.
+         */
+        private static final int REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY = 200;
+
+        /**
+         * 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<NotificationEntry> mEntriesKeptForRemoteInputActive =
+                new ArraySet<>();
+
+        protected NotificationLifetimeExtender.NotificationSafeToRemoveCallback
+                mNotificationLifetimeFinishedCallback;
+
+        protected final ArrayList<NotificationLifetimeExtender> mLifetimeExtenders =
+                new ArrayList<>();
+        private RemoteInputController mRemoteInputController;
+
+        LegacyRemoteInputLifetimeExtender() {
+            addLifetimeExtenders();
+        }
+
+        /**
+         * 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());
+        }
+
+        @Override
+        public void setRemoteInputController(@NonNull RemoteInputController remoteInputController) {
+            mRemoteInputController= remoteInputController;
+        }
+
+        @Override
+        public void onRemoteInputSent(NotificationEntry entry) {
+            if (FORCE_REMOTE_INPUT_HISTORY
+                    && isNotificationKeptForRemoteInputHistory(entry.getKey())) {
+                mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
+            } 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.
+                mMainHandler.postDelayed(() -> {
+                    if (mEntriesKeptForRemoteInputActive.remove(entry)) {
+                        mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
+                    }
+                }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
+            }
+        }
+
+        @Override
+        public void onPanelCollapsed() {
+            for (int i = 0; i < mEntriesKeptForRemoteInputActive.size(); i++) {
+                NotificationEntry entry = mEntriesKeptForRemoteInputActive.valueAt(i);
+                if (mRemoteInputController != null) {
+                    mRemoteInputController.removeRemoteInput(entry, null);
+                }
+                if (mNotificationLifetimeFinishedCallback != null) {
+                    mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
+                }
+            }
+            mEntriesKeptForRemoteInputActive.clear();
+        }
+
+        @Override
+        public boolean isNotificationKeptForRemoteInputHistory(String key) {
+            return mKeysKeptForRemoteInputHistory.contains(key);
+        }
+
+        @Override
+        public void releaseNotificationIfKeptForRemoteInputHistory(
+                @NonNull NotificationEntry entry) {
+            final String key = entry.getKey();
+            if (isNotificationKeptForRemoteInputHistory(key)) {
+                mMainHandler.postDelayed(() -> {
+                    if (isNotificationKeptForRemoteInputHistory(key)) {
+                        mNotificationLifetimeFinishedCallback.onSafeToRemove(key);
+                    }
+                }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
+            }
+        }
+
+        @VisibleForTesting
+        public Set<NotificationEntry> getEntriesKeptForRemoteInputActive() {
+            return mEntriesKeptForRemoteInputActive;
+        }
+
+        @Override
+        public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw,
+                @NonNull String[] args) {
+            pw.println("LegacyRemoteInputLifetimeExtender:");
+            pw.print("  mKeysKeptForRemoteInputHistory: ");
+            pw.println(mKeysKeptForRemoteInputHistory);
+            pw.print("  mEntriesKeptForRemoteInputActive: ");
+            pw.println(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 NotificationEntry entry) {
+                return shouldKeepForRemoteInputHistory(entry);
+            }
+
+            @Override
+            public void setShouldManageLifetime(NotificationEntry entry,
+                    boolean shouldExtend) {
+                if (shouldExtend) {
+                    StatusBarNotification newSbn = rebuildNotificationForBasicExtension(entry);
+                    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.isRemoved()) {
+                        return;
+                    }
+
+                    if (Log.isLoggable(TAG, Log.DEBUG)) {
+                        Log.d(TAG, "Keeping notification around after sending remote input "
+                                + entry.getKey());
+                    }
+
+                    mKeysKeptForRemoteInputHistory.add(entry.getKey());
+                } else {
+                    mKeysKeptForRemoteInputHistory.remove(entry.getKey());
+                }
+            }
+        }
+
+        /**
+         * 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 NotificationEntry entry) {
+                return shouldKeepForSmartReplyHistory(entry);
+            }
+
+            @Override
+            public void setShouldManageLifetime(NotificationEntry entry,
+                    boolean shouldExtend) {
+                if (shouldExtend) {
+                    StatusBarNotification newSbn = rebuildNotificationForCanceledSmartReplies(entry);
+
+                    if (newSbn == null) {
+                        return;
+                    }
+
+                    mEntryManager.updateNotification(newSbn, null);
+
+                    if (entry.isRemoved()) {
+                        return;
+                    }
+
+                    if (Log.isLoggable(TAG, Log.DEBUG)) {
+                        Log.d(TAG, "Keeping notification around after sending smart reply "
+                                + entry.getKey());
+                    }
+
+                    mKeysKeptForRemoteInputHistory.add(entry.getKey());
+                } else {
+                    mKeysKeptForRemoteInputHistory.remove(entry.getKey());
+                    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 NotificationEntry entry) {
+                return isRemoteInputActive(entry);
+            }
+
+            @Override
+            public void setShouldManageLifetime(NotificationEntry entry,
+                    boolean shouldExtend) {
+                if (shouldExtend) {
+                    if (Log.isLoggable(TAG, Log.DEBUG)) {
+                        Log.d(TAG, "Keeping notification around while remote input active "
+                                + entry.getKey());
+                    }
+                    mEntriesKeptForRemoteInputActive.add(entry);
+                } else {
+                    mEntriesKeptForRemoteInputActive.remove(entry);
+                }
+            }
+        }
+    }
 }

packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputController.java

@@ -299,6 +299,9 @@ public class RemoteInputController {
         default void onRemoteInputSent(NotificationEntry entry) {}
     }
 
+    /**
+     * This is a delegate which implements some view controller pieces of the remote input process
+     */
     public interface Delegate {
         /**
          * Activate remote input if necessary.

packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java

@@ -96,9 +96,11 @@ public interface StatusBarDependenciesModule {
     @Provides
     static NotificationRemoteInputManager provideNotificationRemoteInputManager(
             Context context,
+            FeatureFlags featureFlags,
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
+            NotifPipeline notifPipeline,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             Handler mainHandler,
@@ -108,9 +110,11 @@ public interface StatusBarDependenciesModule {
             DumpManager dumpManager) {
         return new NotificationRemoteInputManager(
                 context,
+                featureFlags,
                 lockscreenUserManager,
                 smartReplyController,
                 notificationEntryManager,
+                notifPipeline,
                 statusBarOptionalLazy,
                 statusBarStateController,
                 mainHandler,

packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java

@@ -22,14 +22,16 @@ import android.service.notification.StatusBarNotification;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
-import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputActiveExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.LegacyRemoteInputLifetimeExtender.RemoteInputActiveExtender;
-import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputHistoryExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.LegacyRemoteInputLifetimeExtender.RemoteInputHistoryExtender;
-import com.android.systemui.statusbar.NotificationRemoteInputManager.SmartReplyHistoryExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.LegacyRemoteInputLifetimeExtender.SmartReplyHistoryExtender;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
@@ -76,12 +78,15 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
     private RemoteInputHistoryExtender mRemoteInputHistoryExtender;
     private SmartReplyHistoryExtender mSmartReplyHistoryExtender;
     private RemoteInputActiveExtender mRemoteInputActiveExtender;
+    private TestableNotificationRemoteInputManager.FakeLegacyRemoteInputLifetimeExtender
+            mLegacyRemoteInputLifetimeExtender;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
         mRemoteInputManager = new TestableNotificationRemoteInputManager(mContext,
+                mock(FeatureFlags.class),
                 mLockscreenUserManager, mSmartReplyController, mEntryManager,
                 () -> Optional.of(mock(StatusBar.class)),
                 mStateController,
@@ -151,18 +156,19 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
     public void testNotificationWithRemoteInputActiveIsRemovedOnCollapse() {
         mRemoteInputActiveExtender.setShouldManageLifetime(mEntry, true /* shouldManage */);
 
-        assertEquals(mRemoteInputManager.getEntriesKeptForRemoteInputActive(),
+        assertEquals(mLegacyRemoteInputLifetimeExtender.getEntriesKeptForRemoteInputActive(),
                 Sets.newArraySet(mEntry));
 
         mRemoteInputManager.onPanelCollapsed();
 
-        assertTrue(mRemoteInputManager.getEntriesKeptForRemoteInputActive().isEmpty());
+        assertTrue(
+                mLegacyRemoteInputLifetimeExtender.getEntriesKeptForRemoteInputActive().isEmpty());
     }
 
     @Test
     public void testRebuildWithRemoteInput_noExistingInput_image() {
         Uri uri = mock(Uri.class);
-        String mimeType  = "image/jpeg";
+        String mimeType = "image/jpeg";
         String text = "image inserted";
         StatusBarNotification newSbn =
                 mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
@@ -229,7 +235,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
     public void testRebuildWithRemoteInput_withExistingInput_image() {
         // Setup a notification entry with 1 remote input.
         Uri uri = mock(Uri.class);
-        String mimeType  = "image/jpeg";
+        String mimeType = "image/jpeg";
         String text = "image inserted";
         StatusBarNotification newSbn =
                 mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
@@ -266,6 +272,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
 
         TestableNotificationRemoteInputManager(
                 Context context,
+                FeatureFlags featureFlags,
                 NotificationLockscreenUserManager lockscreenUserManager,
                 SmartReplyController smartReplyController,
                 NotificationEntryManager notificationEntryManager,
@@ -278,6 +285,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
                 DumpManager dumpManager) {
             super(
                     context,
+                    featureFlags,
                     lockscreenUserManager,
                     smartReplyController,
                     notificationEntryManager,
@@ -297,14 +305,28 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
             mRemoteInputController = controller;
         }
 
+        @NonNull
         @Override
-        protected void addLifetimeExtenders() {
+        protected LegacyRemoteInputLifetimeExtender createLegacyRemoteInputLifetimeExtender(
-            mRemoteInputActiveExtender = new RemoteInputActiveExtender();
+                Handler mainHandler,
-            mRemoteInputHistoryExtender = new RemoteInputHistoryExtender();
+                NotificationEntryManager notificationEntryManager,
-            mSmartReplyHistoryExtender = new SmartReplyHistoryExtender();
+                SmartReplyController smartReplyController) {
-            mLifetimeExtenders.add(mRemoteInputHistoryExtender);
+            mLegacyRemoteInputLifetimeExtender = new FakeLegacyRemoteInputLifetimeExtender();
-            mLifetimeExtenders.add(mSmartReplyHistoryExtender);
+            return mLegacyRemoteInputLifetimeExtender;
-            mLifetimeExtenders.add(mRemoteInputActiveExtender);
         }
+
+        class FakeLegacyRemoteInputLifetimeExtender extends LegacyRemoteInputLifetimeExtender {
+
+            @Override
+            protected void addLifetimeExtenders() {
+                mRemoteInputActiveExtender = new RemoteInputActiveExtender();
+                mRemoteInputHistoryExtender = new RemoteInputHistoryExtender();
+                mSmartReplyHistoryExtender = new SmartReplyHistoryExtender();
+                mLifetimeExtenders.add(mRemoteInputHistoryExtender);
+                mLifetimeExtenders.add(mSmartReplyHistoryExtender);
+                mLifetimeExtenders.add(mRemoteInputActiveExtender);
+            }
+        }
+
     }
 }

packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java

@@ -39,8 +39,10 @@ import com.android.internal.statusbar.IStatusBarService;
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
+import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
 import com.android.systemui.statusbar.phone.StatusBar;
@@ -95,8 +97,11 @@ public class SmartReplyControllerTest extends SysuiTestCase {
                 mSmartReplyController);
 
         mRemoteInputManager = new NotificationRemoteInputManager(mContext,
+                mock(FeatureFlags.class),
                 mock(NotificationLockscreenUserManager.class), mSmartReplyController,
-                mNotificationEntryManager, () -> Optional.of(mock(StatusBar.class)),
+                mNotificationEntryManager,
+                mock(NotifPipeline.class),
+                () -> Optional.of(mock(StatusBar.class)),
                 mStatusBarStateController,
                 Handler.createAsync(Looper.myLooper()),
                 mRemoteInputUriController,