Merge changes from topic "b204127880_pipeline_backport_3" into sc-v2-dev

* changes:
  New Pipeline: Remote Input 4/4: Add RemoteInputCoordinator
  New Pipeline: Remote Input 3/4: Extract notification rebuilder methods to a utility
  New Pipeline: Remote Input 2/4: Add ability to internally update notifications
  New Pipeline: Remote Input 1/4: Extract legacy pipeline logic within NotificationRemoteInputManager
  Make SmartReplyController a Dumpable

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

@@ -15,21 +15,16 @@
  */
 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;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.app.RemoteInput;
-import android.app.RemoteInputHistoryItem;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.UserInfo;
-import android.net.Uri;
 import android.os.Handler;
-import android.os.Parcelable;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.SystemClock;
@@ -48,6 +43,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,6 +53,7 @@ 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;
@@ -70,12 +69,10 @@ import com.android.systemui.statusbar.policy.RemoteInputView;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Stream;
 
 import dagger.Lazy;
 
@@ -93,27 +90,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,18 +102,17 @@ 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 RemoteInputNotificationRebuilder mRebuilder;
     private final StatusBarStateController mStatusBarStateController;
     private final RemoteInputUriController mRemoteInputUriController;
     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 +202,7 @@ public class NotificationRemoteInputManager implements Dumpable {
                 ViewGroup actionGroup = (ViewGroup) parent;
                 buttonIndex = actionGroup.indexOfChild(view);
             }
+            // TODO(b/204183781): get this from the current pipeline
             final int count = mEntryManager.getActiveNotificationsCount();
             final int rank = entry.getRanking().getRank();
 
@@ -283,9 +260,11 @@ public class NotificationRemoteInputManager implements Dumpable {
      */
     public NotificationRemoteInputManager(
             Context context,
+            FeatureFlags featureFlags,
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
+            RemoteInputNotificationRebuilder rebuilder,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             @Main Handler mainHandler,
@@ -294,6 +273,7 @@ public class NotificationRemoteInputManager implements Dumpable {
             ActionClickLogger logger,
             DumpManager dumpManager) {
         mContext = context;
+        mFeatureFlags = featureFlags;
         mLockscreenUserManager = lockscreenUserManager;
         mSmartReplyController = smartReplyController;
         mEntryManager = notificationEntryManager;
@@ -303,7 +283,11 @@ public class NotificationRemoteInputManager implements Dumpable {
         mBarService = IStatusBarService.Stub.asInterface(
                 ServiceManager.getService(Context.STATUS_BAR_SERVICE));
         mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
-        addLifetimeExtenders();
+        mRebuilder = rebuilder;
+        if (!featureFlags.isNewNotifPipelineRenderingEnabled()) {
+            mRemoteInputListener = createLegacyRemoteInputLifetimeExtender(mainHandler,
+                    notificationEntryManager, smartReplyController);
+        }
         mKeyguardManager = context.getSystemService(KeyguardManager.class);
         mStatusBarStateController = statusBarStateController;
         mRemoteInputUriController = remoteInputUriController;
@@ -335,10 +319,35 @@ public class NotificationRemoteInputManager implements Dumpable {
         });
     }
 
+    /** Add a listener for various remote input events.  Works with NEW pipeline only. */
+    public void setRemoteInputListener(@NonNull RemoteInputListener remoteInputListener) {
+        if (mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
+            if (mRemoteInputListener != null) {
+                throw new IllegalStateException("mRemoteInputListener is already set");
+            }
+            mRemoteInputListener = remoteInputListener;
+            if (mRemoteInputController != null) {
+                mRemoteInputListener.setRemoteInputController(mRemoteInputController);
+            }
+        }
+    }
+
+    @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);
+        if (mRemoteInputListener != null) {
+            mRemoteInputListener.setRemoteInputController(mRemoteInputController);
+        }
         // Register all stored callbacks from before the Controller was initialized.
         for (RemoteInputController.Callback cb : mControllerCallbacks) {
             mRemoteInputController.addCallback(cb);
@@ -347,19 +356,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 +379,12 @@ public class NotificationRemoteInputManager implements Dumpable {
                 }
             }
         });
-        mSmartReplyController.setCallback((entry, reply) -> {
+        if (!mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
-            StatusBarNotification newSbn =
+            mSmartReplyController.setCallback((entry, reply) -> {
-                    rebuildNotificationWithRemoteInputInserted(entry, reply, true /* showSpinner */,
+                StatusBarNotification newSbn = mRebuilder.rebuildForSendingSmartReply(entry, reply);
-                            null /* mimeType */, null /* uri */);
+                mEntryManager.updateNotification(newSbn, null /* ranking */);
-            mEntryManager.updateNotification(newSbn, null /* ranking */);
+            });
-        });
+        }
     }
 
     public void addControllerCallback(RemoteInputController.Callback callback) {
@@ -574,51 +572,47 @@ 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);
+        cleanUpRemoteInputForUserRemoval(entry);
+    }
+
+    /**
+     * Disable remote input on the entry and remove the remote input view.
+     * This should be called when a user dismisses a notification that won't be lifetime extended.
+     */
+    public void cleanUpRemoteInputForUserRemoval(NotificationEntry entry) {
         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 +630,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,64 +651,11 @@ public class NotificationRemoteInputManager implements Dumpable {
         }
     }
 
-    @VisibleForTesting
-    StatusBarNotification rebuildNotificationForCanceledSmartReplies(
-            NotificationEntry entry) {
-        return rebuildNotificationWithRemoteInputInserted(entry, null /* remoteInputTest */,
-                false /* showSpinner */, null /* mimeType */, null /* uri */);
-    }
-
-    @VisibleForTesting
-    StatusBarNotification rebuildNotificationWithRemoteInputInserted(NotificationEntry entry,
-            CharSequence remoteInputText, boolean showSpinner, String mimeType, Uri uri) {
-        StatusBarNotification sbn = entry.getSbn();
-
-        Notification.Builder b = Notification.Builder
-                .recoverBuilder(mContext, sbn.getNotification().clone());
-        if (remoteInputText != null || uri != null) {
-            RemoteInputHistoryItem newItem = uri != null
-                    ? new RemoteInputHistoryItem(mimeType, uri, remoteInputText)
-                    : new RemoteInputHistoryItem(remoteInputText);
-            Parcelable[] oldHistoryItems = sbn.getNotification().extras
-                    .getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-            RemoteInputHistoryItem[] newHistoryItems = oldHistoryItems != null
-                    ? Stream.concat(
-                                Stream.of(newItem),
-                                Arrays.stream(oldHistoryItems).map(p -> (RemoteInputHistoryItem) p))
-                            .toArray(RemoteInputHistoryItem[]::new)
-                    : new RemoteInputHistoryItem[] { newItem };
-            b.setRemoteInputHistory(newHistoryItems);
-        }
-        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:");
+        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 +671,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 +689,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 +782,256 @@ public class NotificationRemoteInputManager implements Dumpable {
          */
         boolean showBouncerIfNecessary();
     }
+
+    /** An interface for listening to remote input events that relate to notification lifetime */
+    public interface RemoteInputListener {
+        /** Called when remote input pending intent has been sent */
+        void onRemoteInputSent(@NonNull NotificationEntry entry);
+
+        /** Called when the notification shade becomes fully closed */
+        void onPanelCollapsed();
+
+        /** @return whether lifetime of a notification is being extended by the listener */
+        boolean isNotificationKeptForRemoteInputHistory(@NonNull String key);
+
+        /** Called on user interaction to end lifetime extension for history */
+        void releaseNotificationIfKeptForRemoteInputHistory(@NonNull NotificationEntry entry);
+
+        /** Called when the RemoteInputController is attached to the manager */
+        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(@NonNull 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(@NonNull 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 = mRebuilder.rebuildForRemoteInputReply(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 = mRebuilder.rebuildForCanceledSmartReplies(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/RemoteInputNotificationRebuilder.java

@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar;
+
+import android.annotation.NonNull;
+import android.app.Notification;
+import android.app.RemoteInputHistoryItem;
+import android.content.Context;
+import android.net.Uri;
+import android.os.Parcelable;
+import android.service.notification.StatusBarNotification;
+import android.text.TextUtils;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+
+import java.util.Arrays;
+import java.util.stream.Stream;
+
+import javax.inject.Inject;
+
+/**
+ * A helper class which will augment the notifications using arguments and other information
+ * accessible to the entry in order to provide intermediate remote input states.
+ */
+@SysUISingleton
+public class RemoteInputNotificationRebuilder {
+
+    private final Context mContext;
+
+    @Inject
+    RemoteInputNotificationRebuilder(Context context) {
+        mContext = context;
+    }
+
+    /**
+     * When a smart reply is sent off to the app, we insert the text into the remote input history,
+     * and show a spinner to indicate that the app has yet to respond.
+     */
+    @NonNull
+    public StatusBarNotification rebuildForSendingSmartReply(NotificationEntry entry,
+            CharSequence reply) {
+        return rebuildWithRemoteInputInserted(entry, reply,
+                true /* showSpinner */,
+                null /* mimeType */, null /* uri */);
+    }
+
+    /**
+     * When the app cancels a notification in response to a smart reply, we remove the spinner
+     * and leave the previously-added reply.  This is the lifetime-extended appearance of the
+     * notification.
+     */
+    @NonNull
+    public StatusBarNotification rebuildForCanceledSmartReplies(
+            NotificationEntry entry) {
+        return rebuildWithRemoteInputInserted(entry, null /* remoteInputTest */,
+                false /* showSpinner */, null /* mimeType */, null /* uri */);
+    }
+
+    /**
+     * When the app cancels a notification in response to a remote input reply, we update the
+     * notification with the reply text and/or attachment. This is the lifetime-extended
+     * appearance of the notification.
+     */
+    @NonNull
+    public StatusBarNotification rebuildForRemoteInputReply(NotificationEntry entry) {
+        CharSequence remoteInputText = entry.remoteInputText;
+        if (TextUtils.isEmpty(remoteInputText)) {
+            remoteInputText = entry.remoteInputTextWhenReset;
+        }
+        String remoteInputMimeType = entry.remoteInputMimeType;
+        Uri remoteInputUri = entry.remoteInputUri;
+        StatusBarNotification newSbn = rebuildWithRemoteInputInserted(entry,
+                remoteInputText, false /* showSpinner */, remoteInputMimeType,
+                remoteInputUri);
+        return newSbn;
+    }
+
+    /** Inner method for generating the SBN */
+    @VisibleForTesting
+    @NonNull
+    StatusBarNotification rebuildWithRemoteInputInserted(NotificationEntry entry,
+            CharSequence remoteInputText, boolean showSpinner, String mimeType, Uri uri) {
+        StatusBarNotification sbn = entry.getSbn();
+
+        Notification.Builder b = Notification.Builder
+                .recoverBuilder(mContext, sbn.getNotification().clone());
+        if (remoteInputText != null || uri != null) {
+            RemoteInputHistoryItem newItem = uri != null
+                    ? new RemoteInputHistoryItem(mimeType, uri, remoteInputText)
+                    : new RemoteInputHistoryItem(remoteInputText);
+            Parcelable[] oldHistoryItems = sbn.getNotification().extras
+                    .getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+            RemoteInputHistoryItem[] newHistoryItems = oldHistoryItems != null
+                    ? Stream.concat(
+                    Stream.of(newItem),
+                    Arrays.stream(oldHistoryItems).map(p -> (RemoteInputHistoryItem) p))
+                    .toArray(RemoteInputHistoryItem[]::new)
+                    : new RemoteInputHistoryItem[] { newItem };
+            b.setRemoteInputHistory(newHistoryItems);
+        }
+        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());
+    }
+
+
+}

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

@@ -19,35 +19,44 @@ import android.app.Notification;
 import android.os.RemoteException;
 import android.util.ArraySet;
 
+import androidx.annotation.NonNull;
+
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.internal.statusbar.NotificationVisibility;
+import com.android.systemui.Dumpable;
+import com.android.systemui.dump.DumpManager;
 import com.android.systemui.statusbar.dagger.StatusBarModule;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.logging.NotificationLogger;
 
+import java.io.FileDescriptor;
+import java.io.PrintWriter;
 import java.util.Set;
 
 /**
  * Handles when smart replies are added to a notification
  * and clicked upon.
  */
-public class SmartReplyController {
+public class SmartReplyController implements Dumpable {
     private final IStatusBarService mBarService;
     private final NotificationEntryManager mEntryManager;
     private final NotificationClickNotifier mClickNotifier;
-    private Set<String> mSendingKeys = new ArraySet<>();
+    private final Set<String> mSendingKeys = new ArraySet<>();
     private Callback mCallback;
 
     /**
      * Injected constructor. See {@link StatusBarModule}.
      */
-    public SmartReplyController(NotificationEntryManager entryManager,
+    public SmartReplyController(
+            DumpManager dumpManager,
+            NotificationEntryManager entryManager,
             IStatusBarService statusBarService,
             NotificationClickNotifier clickNotifier) {
         mBarService = statusBarService;
         mEntryManager = entryManager;
         mClickNotifier = clickNotifier;
+        dumpManager.registerDumpable(this);
     }
 
     public void setCallback(Callback callback) {
@@ -75,6 +84,7 @@ public class SmartReplyController {
     public void smartActionClicked(
             NotificationEntry entry, int actionIndex, Notification.Action action,
             boolean generatedByAssistant) {
+        // TODO(b/204183781): get this from the current pipeline
         final int count = mEntryManager.getActiveNotificationsCount();
         final int rank = entry.getRanking().getRank();
         NotificationVisibility.NotificationLocation location =
@@ -112,6 +122,14 @@ public class SmartReplyController {
         }
     }
 
+    @Override
+    public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw, @NonNull String[] args) {
+        pw.println("mSendingKeys: " + mSendingKeys.size());
+        for (String key : mSendingKeys) {
+            pw.println(" * " + key);
+        }
+    }
+
     /**
      * Callback for any class that needs to do something in response to a smart reply being sent.
      */

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

@@ -43,6 +43,7 @@ import com.android.systemui.statusbar.NotificationMediaManager;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
 import com.android.systemui.statusbar.NotificationViewHierarchyManager;
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder;
 import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.StatusBarStateControllerImpl;
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
@@ -96,9 +97,11 @@ public interface StatusBarDependenciesModule {
     @Provides
     static NotificationRemoteInputManager provideNotificationRemoteInputManager(
             Context context,
+            FeatureFlags featureFlags,
             NotificationLockscreenUserManager lockscreenUserManager,
             SmartReplyController smartReplyController,
             NotificationEntryManager notificationEntryManager,
+            RemoteInputNotificationRebuilder rebuilder,
             Lazy<Optional<StatusBar>> statusBarOptionalLazy,
             StatusBarStateController statusBarStateController,
             Handler mainHandler,
@@ -108,9 +111,11 @@ public interface StatusBarDependenciesModule {
             DumpManager dumpManager) {
         return new NotificationRemoteInputManager(
                 context,
+                featureFlags,
                 lockscreenUserManager,
                 smartReplyController,
                 notificationEntryManager,
+                rebuilder,
                 statusBarOptionalLazy,
                 statusBarStateController,
                 mainHandler,
@@ -166,10 +171,11 @@ public interface StatusBarDependenciesModule {
     @SysUISingleton
     @Provides
     static SmartReplyController provideSmartReplyController(
+            DumpManager dumpManager,
             NotificationEntryManager entryManager,
             IStatusBarService statusBarService,
             NotificationClickNotifier clickNotifier) {
-        return new SmartReplyController(entryManager, statusBarService, clickNotifier);
+        return new SmartReplyController(dumpManager, entryManager, statusBarService, clickNotifier);
     }
 
 

packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java

@@ -689,8 +689,9 @@ public class NotificationEntryManager implements
         for (NotificationEntryListener listener : mNotificationEntryListeners) {
             listener.onPreEntryUpdated(entry);
         }
+        final boolean fromSystem = ranking != null;
         for (NotifCollectionListener listener : mNotifCollectionListeners) {
-            listener.onEntryUpdated(entry);
+            listener.onEntryUpdated(entry, fromSystem);
         }
 
         if (!mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollection.java

@@ -47,6 +47,7 @@ import android.annotation.MainThread;
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.app.Notification;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.os.Trace;
 import android.os.UserHandle;
@@ -62,6 +63,7 @@ import androidx.annotation.NonNull;
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.systemui.Dumpable;
 import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.dump.LogBufferEulogizer;
 import com.android.systemui.flags.FeatureFlags;
@@ -76,6 +78,7 @@ import com.android.systemui.statusbar.notification.collection.notifcollection.En
 import com.android.systemui.statusbar.notification.collection.notifcollection.EntryRemovedEvent;
 import com.android.systemui.statusbar.notification.collection.notifcollection.EntryUpdatedEvent;
 import com.android.systemui.statusbar.notification.collection.notifcollection.InitEntryEvent;
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionLogger;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifDismissInterceptor;
@@ -131,6 +134,7 @@ public class NotifCollection implements Dumpable {
     private final SystemClock mClock;
     private final FeatureFlags mFeatureFlags;
     private final NotifCollectionLogger mLogger;
+    private final Handler mMainHandler;
     private final LogBufferEulogizer mEulogizer;
 
     private final Map<String, NotificationEntry> mNotificationSet = new ArrayMap<>();
@@ -154,6 +158,7 @@ public class NotifCollection implements Dumpable {
             SystemClock clock,
             FeatureFlags featureFlags,
             NotifCollectionLogger logger,
+            @Main Handler mainHandler,
             LogBufferEulogizer logBufferEulogizer,
             DumpManager dumpManager) {
         Assert.isMainThread();
@@ -161,6 +166,7 @@ public class NotifCollection implements Dumpable {
         mClock = clock;
         mFeatureFlags = featureFlags;
         mLogger = logger;
+        mMainHandler = mainHandler;
         mEulogizer = logBufferEulogizer;
 
         dumpManager.registerDumpable(TAG, this);
@@ -442,7 +448,7 @@ public class NotifCollection implements Dumpable {
             mEventQueue.add(new BindEntryEvent(entry, sbn));
 
             mLogger.logNotifUpdated(sbn.getKey());
-            mEventQueue.add(new EntryUpdatedEvent(entry));
+            mEventQueue.add(new EntryUpdatedEvent(entry, true /* fromSystem */));
         }
     }
 
@@ -791,6 +797,51 @@ public class NotifCollection implements Dumpable {
 
     private static final String TAG = "NotifCollection";
 
+    /**
+     * Get an object which can be used to update a notification (internally to the pipeline)
+     * in response to a user action.
+     *
+     * @param name the name of the component that will update notifiations
+     * @return an updater
+     */
+    public InternalNotifUpdater getInternalNotifUpdater(String name) {
+        return (sbn, reason) -> mMainHandler.post(
+                () -> updateNotificationInternally(sbn, name, reason));
+    }
+
+    /**
+     * Provide an updated StatusBarNotification for an existing entry.  If no entry exists for the
+     * given notification key, this method does nothing.
+     *
+     * @param sbn the updated notification
+     * @param name the component which is updating the notification
+     * @param reason the reason the notification is being updated
+     */
+    private void updateNotificationInternally(StatusBarNotification sbn, String name,
+            String reason) {
+        Assert.isMainThread();
+        checkForReentrantCall();
+
+        // Make sure we have the notification to update
+        NotificationEntry entry = mNotificationSet.get(sbn.getKey());
+        if (entry == null) {
+            mLogger.logNotifInternalUpdateFailed(sbn.getKey(), name, reason);
+            return;
+        }
+        mLogger.logNotifInternalUpdate(sbn.getKey(), name, reason);
+
+        // First do the pieces of postNotification which are not about assuming the notification
+        // was sent by the app
+        entry.setSbn(sbn);
+        mEventQueue.add(new BindEntryEvent(entry, sbn));
+
+        mLogger.logNotifUpdated(sbn.getKey());
+        mEventQueue.add(new EntryUpdatedEvent(entry, false /* fromSystem */));
+
+        // Skip the applyRanking step and go straight to dispatching the events
+        dispatchEventsAndRebuildList();
+    }
+
     @IntDef(prefix = { "REASON_" }, value = {
             REASON_NOT_CANCELED,
             REASON_UNKNOWN,

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifPipeline.java

@@ -16,6 +16,8 @@
 
 package com.android.systemui.statusbar.notification.collection;
 
+import android.os.Handler;
+
 import androidx.annotation.Nullable;
 
 import com.android.systemui.dagger.SysUISingleton;
@@ -30,6 +32,7 @@ import com.android.systemui.statusbar.notification.collection.listbuilder.plugga
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifStabilityManager;
 import com.android.systemui.statusbar.notification.collection.notifcollection.CommonNotifCollection;
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifDismissInterceptor;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender;
@@ -222,6 +225,17 @@ public class NotifPipeline implements CommonNotifCollection {
         mShadeListBuilder.addPreRenderInvalidator(invalidator);
     }
 
+    /**
+     * Get an object which can be used to update a notification (internally to the pipeline)
+     * in response to a user action.
+     *
+     * @param name the name of the component that will update notifiations
+     * @return an updater
+     */
+    public InternalNotifUpdater getInternalNotifUpdater(String name) {
+        return mNotifCollection.getInternalNotifUpdater(name);
+    }
+
     /**
      * Returns a read-only view in to the current shade list, i.e. the list of notifications that
      * are currently present in the shade. If this method is called during pipeline execution it

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt

@@ -48,6 +48,7 @@ class NotifCoordinatorsImpl @Inject constructor(
     conversationCoordinator: ConversationCoordinator,
     preparationCoordinator: PreparationCoordinator,
     mediaCoordinator: MediaCoordinator,
+    remoteInputCoordinator: RemoteInputCoordinator,
     shadeEventCoordinator: ShadeEventCoordinator,
     smartspaceDedupingCoordinator: SmartspaceDedupingCoordinator,
     viewConfigCoordinator: ViewConfigCoordinator,
@@ -72,6 +73,7 @@ class NotifCoordinatorsImpl @Inject constructor(
         mCoordinators.add(bubbleCoordinator)
         mCoordinators.add(conversationCoordinator)
         mCoordinators.add(mediaCoordinator)
+        mCoordinators.add(remoteInputCoordinator)
         mCoordinators.add(shadeEventCoordinator)
         mCoordinators.add(viewConfigCoordinator)
         mCoordinators.add(visualStabilityCoordinator)

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt

@@ -0,0 +1,225 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.systemui.statusbar.notification.collection.coordinator
+
+import android.os.Handler
+import android.service.notification.NotificationListenerService.REASON_CANCEL
+import android.service.notification.NotificationListenerService.REASON_CLICK
+import android.util.Log
+import androidx.annotation.VisibleForTesting
+import com.android.systemui.Dumpable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputListener
+import com.android.systemui.statusbar.RemoteInputController
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder
+import com.android.systemui.statusbar.SmartReplyController
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.notifcollection.SelfTrackingLifetimeExtender
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender
+import java.io.FileDescriptor
+import java.io.PrintWriter
+import javax.inject.Inject
+
+private const val TAG = "RemoteInputCoordinator"
+
+/**
+ * How long to wait before auto-dismissing a notification that was kept for active remote input, and
+ * has now sent a remote input. We auto-dismiss, because the app may not cannot cancel
+ * these given that they technically don't exist anymore. We wait a bit in case the app issues
+ * an update, and to also give the other lifetime extenders a beat to decide they want it.
+ */
+private const val REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY: Long = 500
+
+/**
+ * How long to wait before releasing a lifetime extension when requested to do so due to a user
+ * interaction (such as tapping another action).
+ * We wait a bit in case the app issues an update in response to the action, but not too long or we
+ * risk appearing unresponsive to the user.
+ */
+private const val REMOTE_INPUT_EXTENDER_RELEASE_DELAY: Long = 200
+
+/** Whether this class should print spammy debug logs */
+private val DEBUG: Boolean by lazy { Log.isLoggable(TAG, Log.DEBUG) }
+
+@SysUISingleton
+class RemoteInputCoordinator @Inject constructor(
+    dumpManager: DumpManager,
+    private val mRebuilder: RemoteInputNotificationRebuilder,
+    private val mNotificationRemoteInputManager: NotificationRemoteInputManager,
+    @Main private val mMainHandler: Handler,
+    private val mSmartReplyController: SmartReplyController
+) : Coordinator, RemoteInputListener, Dumpable {
+
+    @VisibleForTesting val mRemoteInputHistoryExtender = RemoteInputHistoryExtender()
+    @VisibleForTesting val mSmartReplyHistoryExtender = SmartReplyHistoryExtender()
+    @VisibleForTesting val mRemoteInputActiveExtender = RemoteInputActiveExtender()
+    private val mRemoteInputLifetimeExtenders = listOf(
+            mRemoteInputHistoryExtender,
+            mSmartReplyHistoryExtender,
+            mRemoteInputActiveExtender
+    )
+
+    private lateinit var mNotifUpdater: InternalNotifUpdater
+
+    init {
+        dumpManager.registerDumpable(this)
+    }
+
+    fun getLifetimeExtenders(): List<NotifLifetimeExtender> = mRemoteInputLifetimeExtenders
+
+    override fun attach(pipeline: NotifPipeline) {
+        mNotificationRemoteInputManager.setRemoteInputListener(this)
+        mRemoteInputLifetimeExtenders.forEach { pipeline.addNotificationLifetimeExtender(it) }
+        mNotifUpdater = pipeline.getInternalNotifUpdater(TAG)
+        pipeline.addCollectionListener(mCollectionListener)
+    }
+
+    val mCollectionListener = object : NotifCollectionListener {
+        override fun onEntryUpdated(entry: NotificationEntry, fromSystem: Boolean) {
+            if (DEBUG) {
+                Log.d(TAG, "mCollectionListener.onEntryUpdated(entry=${entry.key}," +
+                        " fromSystem=$fromSystem)")
+            }
+            if (fromSystem) {
+                // Mark smart replies as sent whenever a notification is updated by the app,
+                // otherwise the smart replies are never marked as sent.
+                mSmartReplyController.stopSending(entry)
+            }
+        }
+
+        override fun onEntryRemoved(entry: NotificationEntry, reason: Int) {
+            if (DEBUG) Log.d(TAG, "mCollectionListener.onEntryRemoved(entry=${entry.key})")
+            // We're removing the notification, the smart reply controller can forget about it.
+            // TODO(b/145659174): track 'sending' state on the entry to avoid having to clear it.
+            mSmartReplyController.stopSending(entry)
+
+            // When we know the entry will not be lifetime extended, clean up the remote input view
+            // TODO: Share code with NotifCollection.cannotBeLifetimeExtended
+            if (reason == REASON_CANCEL || reason == REASON_CLICK) {
+                mNotificationRemoteInputManager.cleanUpRemoteInputForUserRemoval(entry)
+            }
+        }
+    }
+
+    override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
+        mRemoteInputLifetimeExtenders.forEach { it.dump(fd, pw, args) }
+    }
+
+    override fun onRemoteInputSent(entry: NotificationEntry) {
+        if (DEBUG) Log.d(TAG, "onRemoteInputSent(entry=${entry.key})")
+        // These calls effectively ensure the freshness of the lifetime extensions.
+        // NOTE: This is some trickery! By removing the lifetime extensions when we know they should
+        // be immediately re-upped, we ensure that the side-effects of the lifetime extenders get to
+        // fire again, thus ensuring that we add subsequent replies to the notification.
+        mRemoteInputHistoryExtender.endLifetimeExtension(entry.key)
+        mSmartReplyHistoryExtender.endLifetimeExtension(entry.key)
+
+        // If we're extending for remote input being active, then from the apps point of
+        // view it is already canceled, so we'll need to cancel it on the apps behalf
+        // now that a reply has been sent. However, delay so that the app has time to posts an
+        // update in the mean time, and to give another lifetime extender time to pick it up.
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY)
+    }
+
+    private fun onSmartReplySent(entry: NotificationEntry, reply: CharSequence) {
+        if (DEBUG) Log.d(TAG, "onSmartReplySent(entry=${entry.key})")
+        val newSbn = mRebuilder.rebuildForSendingSmartReply(entry, reply)
+        mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                "Adding smart reply spinner for sent")
+
+        // If we're extending for remote input being active, then from the apps point of
+        // view it is already canceled, so we'll need to cancel it on the apps behalf
+        // now that a reply has been sent. However, delay so that the app has time to posts an
+        // update in the mean time, and to give another lifetime extender time to pick it up.
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_ACTIVE_EXTENDER_AUTO_CANCEL_DELAY)
+    }
+
+    override fun onPanelCollapsed() {
+        mRemoteInputActiveExtender.endAllLifetimeExtensions()
+    }
+
+    override fun isNotificationKeptForRemoteInputHistory(key: String) =
+            mRemoteInputHistoryExtender.isExtending(key) ||
+                    mSmartReplyHistoryExtender.isExtending(key)
+
+    override fun releaseNotificationIfKeptForRemoteInputHistory(entry: NotificationEntry) {
+        if (DEBUG) Log.d(TAG, "releaseNotificationIfKeptForRemoteInputHistory(entry=${entry.key})")
+        mRemoteInputHistoryExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+        mSmartReplyHistoryExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+        mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(entry.key,
+                REMOTE_INPUT_EXTENDER_RELEASE_DELAY)
+    }
+
+    override fun setRemoteInputController(remoteInputController: RemoteInputController) {
+        mSmartReplyController.setCallback(this::onSmartReplySent)
+    }
+
+    @VisibleForTesting
+    inner class RemoteInputHistoryExtender :
+            SelfTrackingLifetimeExtender(TAG, "RemoteInputHistory", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.shouldKeepForRemoteInputHistory(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            val newSbn = mRebuilder.rebuildForRemoteInputReply(entry)
+            entry.onRemoteInputInserted()
+            mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                    "Extending lifetime of notification with remote input")
+            // TODO: Check if the entry was removed due perhaps to an inflation exception?
+        }
+    }
+
+    @VisibleForTesting
+    inner class SmartReplyHistoryExtender :
+            SelfTrackingLifetimeExtender(TAG, "SmartReplyHistory", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.shouldKeepForSmartReplyHistory(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            val newSbn = mRebuilder.rebuildForCanceledSmartReplies(entry)
+            mSmartReplyController.stopSending(entry)
+            mNotifUpdater.onInternalNotificationUpdate(newSbn,
+                    "Extending lifetime of notification with smart reply")
+            // TODO: Check if the entry was removed due perhaps to an inflation exception?
+        }
+
+        override fun onCanceledLifetimeExtension(entry: NotificationEntry) {
+            // TODO(b/145659174): track 'sending' state on the entry to avoid having to clear it.
+            mSmartReplyController.stopSending(entry)
+        }
+    }
+
+    @VisibleForTesting
+    inner class RemoteInputActiveExtender :
+            SelfTrackingLifetimeExtender(TAG, "RemoteInputActive", DEBUG, mMainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean =
+                mNotificationRemoteInputManager.isRemoteInputActive(entry)
+    }
+}

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/InternalNotifUpdater.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.collection.notifcollection;
+
+import android.service.notification.StatusBarNotification;
+
+/**
+ * An object that allows Coordinators to update notifications internally to SystemUI.
+ * This is used when part of the UI involves updating the underlying appearance of a notification
+ * on behalf of an app, such as to add a spinner or remote input history.
+ */
+public interface InternalNotifUpdater {
+    /**
+     * Called when an already-existing notification needs to be updated to a new temporary
+     * appearance.
+     * This update is local to the SystemUI process.
+     * This has no effect if no notification with the given key exists in the pipeline.
+     *
+     * @param sbn a notification to update
+     * @param reason a debug reason for the update
+     */
+    void onInternalNotificationUpdate(StatusBarNotification sbn, String reason);
+}

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionListener.java

@@ -53,6 +53,17 @@ public interface NotifCollectionListener {
     default void onEntryAdded(@NonNull NotificationEntry entry) {
     }
 
+    /**
+     * Called whenever a notification with the same key as an existing notification is posted. By
+     * the time this listener is called, the entry's SBN and Ranking will already have been updated.
+     * This delegates to {@link #onEntryUpdated(NotificationEntry)} by default.
+     * @param fromSystem If true, this update came from the NotificationManagerService.
+     *                   If false, the notification update is an internal change within systemui.
+     */
+    default void onEntryUpdated(@NonNull NotificationEntry entry, boolean fromSystem) {
+        onEntryUpdated(entry);
+    }
+
     /**
      * Called whenever a notification with the same key as an existing notification is posted. By
      * the time this listener is called, the entry's SBN and Ranking will already have been updated.

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionLogger.kt

@@ -121,6 +121,26 @@ class NotifCollectionLogger @Inject constructor(
         })
     }
 
+    fun logNotifInternalUpdate(key: String, name: String, reason: String) {
+        buffer.log(TAG, INFO, {
+            str1 = key
+            str2 = name
+            str3 = reason
+        }, {
+            "UPDATED INTERNALLY $str1 BY $str2 BECAUSE $str3"
+        })
+    }
+
+    fun logNotifInternalUpdateFailed(key: String, name: String, reason: String) {
+        buffer.log(TAG, INFO, {
+            str1 = key
+            str2 = name
+            str3 = reason
+        }, {
+            "FAILED INTERNAL UPDATE $str1 BY $str2 BECAUSE $str3"
+        })
+    }
+
     fun logNoNotificationToRemoveWithKey(key: String) {
         buffer.log(TAG, ERROR, {
             str1 = key

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifEvent.kt

@@ -64,10 +64,11 @@ data class EntryAddedEvent(
 }
 
 data class EntryUpdatedEvent(
-    val entry: NotificationEntry
+    val entry: NotificationEntry,
+    val fromSystem: Boolean
 ) : NotifEvent() {
     override fun dispatchToListener(listener: NotifCollectionListener) {
-        listener.onEntryUpdated(entry)
+        listener.onEntryUpdated(entry, fromSystem)
     }
 }
 

packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt

@@ -0,0 +1,113 @@
+package com.android.systemui.statusbar.notification.collection.notifcollection
+
+import android.os.Handler
+import android.util.ArrayMap
+import android.util.Log
+import com.android.systemui.Dumpable
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import java.io.FileDescriptor
+import java.io.PrintWriter
+
+/**
+ * A helpful class that implements the core contract of the lifetime extender internally,
+ * making it easier for coordinators to interact with them
+ */
+abstract class SelfTrackingLifetimeExtender(
+    private val tag: String,
+    private val name: String,
+    private val debug: Boolean,
+    private val mainHandler: Handler
+) : NotifLifetimeExtender, Dumpable {
+    private lateinit var mCallback: NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+    protected val mEntriesExtended = ArrayMap<String, NotificationEntry>()
+    private var mEnding = false
+
+    /**
+     * When debugging, warn if the call is happening during and "end lifetime extension" call.
+     *
+     * Note: this will warn a lot! The pipeline explicitly re-invokes all lifetime extenders
+     * whenever one ends, giving all of them a chance to re-up their lifetime extension.
+     */
+    private fun warnIfEnding() {
+        if (debug && mEnding) Log.w(tag, "reentrant code while ending a lifetime extension")
+    }
+
+    fun endAllLifetimeExtensions() {
+        // clear the map before iterating over a copy of the items, because the pipeline will
+        // always give us another chance to extend the lifetime again, and we don't want
+        // concurrent modification
+        val entries = mEntriesExtended.values.toList()
+        if (debug) Log.d(tag, "$name.endAllLifetimeExtensions() entries=$entries")
+        mEntriesExtended.clear()
+        warnIfEnding()
+        mEnding = true
+        entries.forEach { mCallback.onEndLifetimeExtension(this, it) }
+        mEnding = false
+    }
+
+    fun endLifetimeExtensionAfterDelay(key: String, delayMillis: Long) {
+        if (debug) {
+            Log.d(tag, "$name.endLifetimeExtensionAfterDelay" +
+                    "(key=$key, delayMillis=$delayMillis)" +
+                    " isExtending=${isExtending(key)}")
+        }
+        if (isExtending(key)) {
+            mainHandler.postDelayed({ endLifetimeExtension(key) }, delayMillis)
+        }
+    }
+
+    fun endLifetimeExtension(key: String) {
+        if (debug) {
+            Log.d(tag, "$name.endLifetimeExtension(key=$key)" +
+                    " isExtending=${isExtending(key)}")
+        }
+        warnIfEnding()
+        mEnding = true
+        mEntriesExtended.remove(key)?.let { removedEntry ->
+            mCallback.onEndLifetimeExtension(this, removedEntry)
+        }
+        mEnding = false
+    }
+
+    fun isExtending(key: String) = mEntriesExtended.contains(key)
+
+    final override fun getName(): String = name
+
+    final override fun shouldExtendLifetime(entry: NotificationEntry, reason: Int): Boolean {
+        val shouldExtend = queryShouldExtendLifetime(entry)
+        if (debug) {
+            Log.d(tag, "$name.shouldExtendLifetime(key=${entry.key}, reason=$reason)" +
+                    " isExtending=${isExtending(entry.key)}" +
+                    " shouldExtend=$shouldExtend")
+        }
+        warnIfEnding()
+        if (shouldExtend && mEntriesExtended.put(entry.key, entry) == null) {
+            onStartedLifetimeExtension(entry)
+        }
+        return shouldExtend
+    }
+
+    final override fun cancelLifetimeExtension(entry: NotificationEntry) {
+        if (debug) {
+            Log.d(tag, "$name.cancelLifetimeExtension(key=${entry.key})" +
+                    " isExtending=${isExtending(entry.key)}")
+        }
+        warnIfEnding()
+        mEntriesExtended.remove(entry.key)
+        onCanceledLifetimeExtension(entry)
+    }
+
+    abstract fun queryShouldExtendLifetime(entry: NotificationEntry): Boolean
+    open fun onStartedLifetimeExtension(entry: NotificationEntry) {}
+    open fun onCanceledLifetimeExtension(entry: NotificationEntry) {}
+
+    final override fun setCallback(callback: NotifLifetimeExtender.OnEndLifetimeExtensionCallback) {
+        mCallback = callback
+    }
+
+    final override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
+        pw.println("LifetimeExtender: $name:")
+        pw.println("  mEntriesExtended: ${mEntriesExtended.size}")
+        mEntriesExtended.forEach { pw.println("  * ${it.key}") }
+    }
+}

packages/SystemUI/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallController.kt

@@ -86,7 +86,7 @@ class OngoingCallController @Inject constructor(
         //
         // TODO(b/183229367): Remove this function override when b/178406514 is fixed.
         override fun onEntryAdded(entry: NotificationEntry) {
-            onEntryUpdated(entry)
+            onEntryUpdated(entry, true)
         }
 
         override fun onEntryUpdated(entry: NotificationEntry) {

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

@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 
 package com.android.systemui.statusbar;
 
@@ -10,26 +25,25 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.Notification;
-import android.app.RemoteInputHistoryItem;
 import android.content.Context;
-import android.net.Uri;
 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;
 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,13 +90,19 @@ 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,
-                mLockscreenUserManager, mSmartReplyController, mEntryManager,
+                mock(FeatureFlags.class),
+                mLockscreenUserManager,
+                mSmartReplyController,
+                mEntryManager,
+                mock(RemoteInputNotificationRebuilder.class),
                 () -> Optional.of(mock(StatusBar.class)),
                 mStateController,
                 Handler.createAsync(Looper.myLooper()),
@@ -120,6 +140,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
     public void testShouldExtendLifetime_remoteInputActive() {
         when(mController.isRemoteInputActive(mEntry)).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.isRemoteInputActive(mEntry));
         assertTrue(mRemoteInputActiveExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -128,6 +149,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         when(mController.isSpinning(mEntry.getKey())).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.shouldKeepForRemoteInputHistory(mEntry));
         assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -136,6 +158,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         mEntry.lastRemoteInputSent = SystemClock.elapsedRealtime();
 
+        assertTrue(mRemoteInputManager.shouldKeepForRemoteInputHistory(mEntry));
         assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -144,6 +167,7 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         when(mSmartReplyController.isSendingSmartReply(mEntry.getKey())).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.shouldKeepForSmartReplyHistory(mEntry));
         assertTrue(mSmartReplyHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -151,124 +175,24 @@ 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 text = "image inserted";
-        StatusBarNotification newSbn =
-                mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
-                        mEntry, text, false, mimeType, uri);
-        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
-                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-        assertEquals(1, messages.length);
-        assertEquals(text, messages[0].getText());
-        assertEquals(mimeType, messages[0].getMimeType());
-        assertEquals(uri, messages[0].getUri());
-    }
-
-    @Test
-    public void testRebuildWithRemoteInput_noExistingInputNoSpinner() {
-        StatusBarNotification newSbn =
-                mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
-                        mEntry, "A Reply", false, null, null);
-        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
-                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-        assertEquals(1, messages.length);
-        assertEquals("A Reply", messages[0].getText());
-        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.rebuildNotificationWithRemoteInputInserted(
-                        mEntry, "A Reply", true, null, null);
-        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
-                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-        assertEquals(1, messages.length);
-        assertEquals("A Reply", messages[0].getText());
-        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.rebuildNotificationWithRemoteInputInserted(
-                        mEntry, "A Reply", false, null, null);
-        NotificationEntry entry = new NotificationEntryBuilder()
-                .setSbn(newSbn)
-                .build();
-
-        // Try rebuilding to add another reply.
-        newSbn = mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
-                entry, "Reply 2", true, null, null);
-        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
-                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-        assertEquals(2, messages.length);
-        assertEquals("Reply 2", messages[0].getText());
-        assertEquals("A Reply", messages[1].getText());
-    }
-
-    @Test
-    public void testRebuildWithRemoteInput_withExistingInput_image() {
-        // Setup a notification entry with 1 remote input.
-        Uri uri = mock(Uri.class);
-        String mimeType  = "image/jpeg";
-        String text = "image inserted";
-        StatusBarNotification newSbn =
-                mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
-                        mEntry, text, false, mimeType, uri);
-        NotificationEntry entry = new NotificationEntryBuilder()
-                .setSbn(newSbn)
-                .build();
-
-        // Try rebuilding to add another reply.
-        newSbn = mRemoteInputManager.rebuildNotificationWithRemoteInputInserted(
-                entry, "Reply 2", true, null, null);
-        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
-                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
-        assertEquals(2, messages.length);
-        assertEquals("Reply 2", messages[0].getText());
-        assertEquals(text, messages[1].getText());
-        assertEquals(mimeType, messages[1].getMimeType());
-        assertEquals(uri, messages[1].getUri());
-    }
-
-    @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 {
 
         TestableNotificationRemoteInputManager(
                 Context context,
+                FeatureFlags featureFlags,
                 NotificationLockscreenUserManager lockscreenUserManager,
                 SmartReplyController smartReplyController,
                 NotificationEntryManager notificationEntryManager,
+                RemoteInputNotificationRebuilder rebuilder,
                 Lazy<Optional<StatusBar>> statusBarOptionalLazy,
                 StatusBarStateController statusBarStateController,
                 Handler mainHandler,
@@ -278,9 +202,11 @@ public class NotificationRemoteInputManagerTest extends SysuiTestCase {
                 DumpManager dumpManager) {
             super(
                     context,
+                    featureFlags,
                     lockscreenUserManager,
                     smartReplyController,
                     notificationEntryManager,
+                    rebuilder,
                     statusBarOptionalLazy,
                     statusBarStateController,
                     mainHandler,
@@ -297,14 +223,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/RemoteInputNotificationRebuilderTest.java

@@ -0,0 +1,174 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertTrue;
+
+import static org.mockito.Mockito.mock;
+
+import android.app.Notification;
+import android.app.RemoteInputHistoryItem;
+import android.net.Uri;
+import android.os.UserHandle;
+import android.service.notification.StatusBarNotification;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
+import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@SmallTest
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+public class RemoteInputNotificationRebuilderTest extends SysuiTestCase {
+    private static final String TEST_PACKAGE_NAME = "test";
+    private static final int TEST_UID = 0;
+    @Mock
+    private ExpandableNotificationRow mRow;
+
+    private RemoteInputNotificationRebuilder mRebuilder;
+    private NotificationEntry mEntry;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mRebuilder = new RemoteInputNotificationRebuilder(mContext);
+        mEntry = new NotificationEntryBuilder()
+                .setPkg(TEST_PACKAGE_NAME)
+                .setOpPkg(TEST_PACKAGE_NAME)
+                .setUid(TEST_UID)
+                .setNotification(new Notification())
+                .setUser(UserHandle.CURRENT)
+                .build();
+        mEntry.setRow(mRow);
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInput_image() {
+        Uri uri = mock(Uri.class);
+        String mimeType = "image/jpeg";
+        String text = "image inserted";
+        StatusBarNotification newSbn =
+                mRebuilder.rebuildWithRemoteInputInserted(
+                        mEntry, text, false, mimeType, uri);
+        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
+                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+        assertEquals(1, messages.length);
+        assertEquals(text, messages[0].getText());
+        assertEquals(mimeType, messages[0].getMimeType());
+        assertEquals(uri, messages[0].getUri());
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInputNoSpinner() {
+        StatusBarNotification newSbn =
+                mRebuilder.rebuildWithRemoteInputInserted(
+                        mEntry, "A Reply", false, null, null);
+        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
+                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+        assertEquals(1, messages.length);
+        assertEquals("A Reply", messages[0].getText());
+        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 =
+                mRebuilder.rebuildWithRemoteInputInserted(
+                        mEntry, "A Reply", true, null, null);
+        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
+                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+        assertEquals(1, messages.length);
+        assertEquals("A Reply", messages[0].getText());
+        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 =
+                mRebuilder.rebuildWithRemoteInputInserted(
+                        mEntry, "A Reply", false, null, null);
+        NotificationEntry entry = new NotificationEntryBuilder()
+                .setSbn(newSbn)
+                .build();
+
+        // Try rebuilding to add another reply.
+        newSbn = mRebuilder.rebuildWithRemoteInputInserted(
+                entry, "Reply 2", true, null, null);
+        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
+                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+        assertEquals(2, messages.length);
+        assertEquals("Reply 2", messages[0].getText());
+        assertEquals("A Reply", messages[1].getText());
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_withExistingInput_image() {
+        // Setup a notification entry with 1 remote input.
+        Uri uri = mock(Uri.class);
+        String mimeType = "image/jpeg";
+        String text = "image inserted";
+        StatusBarNotification newSbn =
+                mRebuilder.rebuildWithRemoteInputInserted(
+                        mEntry, text, false, mimeType, uri);
+        NotificationEntry entry = new NotificationEntryBuilder()
+                .setSbn(newSbn)
+                .build();
+
+        // Try rebuilding to add another reply.
+        newSbn = mRebuilder.rebuildWithRemoteInputInserted(
+                entry, "Reply 2", true, null, null);
+        RemoteInputHistoryItem[] messages = (RemoteInputHistoryItem[]) newSbn.getNotification()
+                .extras.getParcelableArray(Notification.EXTRA_REMOTE_INPUT_HISTORY_ITEMS);
+        assertEquals(2, messages.length);
+        assertEquals("Reply 2", messages[0].getText());
+        assertEquals(text, messages[1].getText());
+        assertEquals(mimeType, messages[1].getMimeType());
+        assertEquals(uri, messages[1].getUri());
+    }
+
+    @Test
+    public void testRebuildNotificationForCanceledSmartReplies() {
+        // Try rebuilding to remove spinner and hide buttons.
+        StatusBarNotification newSbn =
+                mRebuilder.rebuildForCanceledSmartReplies(mEntry);
+        assertFalse(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+        assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_HIDE_SMART_REPLIES, false));
+    }
+}

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

@@ -39,6 +39,7 @@ 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.NotificationEntry;
@@ -86,14 +87,20 @@ public class SmartReplyControllerTest extends SysuiTestCase {
         mDependency.injectTestDependency(NotificationEntryManager.class,
                 mNotificationEntryManager);
 
-        mSmartReplyController = new SmartReplyController(mNotificationEntryManager,
+        mSmartReplyController = new SmartReplyController(
-                mIStatusBarService, mClickNotifier);
+                mock(DumpManager.class),
+                mNotificationEntryManager,
+                mIStatusBarService,
+                mClickNotifier);
         mDependency.injectTestDependency(SmartReplyController.class,
                 mSmartReplyController);
 
         mRemoteInputManager = new NotificationRemoteInputManager(mContext,
+                mock(FeatureFlags.class),
                 mock(NotificationLockscreenUserManager.class), mSmartReplyController,
-                mNotificationEntryManager, () -> Optional.of(mock(StatusBar.class)),
+                mNotificationEntryManager,
+                new RemoteInputNotificationRebuilder(mContext),
+                () -> Optional.of(mock(StatusBar.class)),
                 mStatusBarStateController,
                 Handler.createAsync(Looper.myLooper()),
                 mRemoteInputUriController,

packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionTest.java

@@ -35,6 +35,7 @@ import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
@@ -50,6 +51,7 @@ import static java.util.Objects.requireNonNull;
 
 import android.annotation.Nullable;
 import android.app.Notification;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.service.notification.NotificationListenerService.Ranking;
 import android.service.notification.NotificationListenerService.RankingMap;
@@ -77,6 +79,7 @@ import com.android.systemui.statusbar.notification.collection.coalescer.GroupCoa
 import com.android.systemui.statusbar.notification.collection.coalescer.GroupCoalescer.BatchableNotificationHandler;
 import com.android.systemui.statusbar.notification.collection.notifcollection.CollectionReadyForBuildListener;
 import com.android.systemui.statusbar.notification.collection.notifcollection.DismissedByUserStats;
+import com.android.systemui.statusbar.notification.collection.notifcollection.InternalNotifUpdater;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionLogger;
 import com.android.systemui.statusbar.notification.collection.notifcollection.NotifDismissInterceptor;
@@ -107,6 +110,7 @@ public class NotifCollectionTest extends SysuiTestCase {
     @Mock private FeatureFlags mFeatureFlags;
     @Mock private NotifCollectionLogger mLogger;
     @Mock private LogBufferEulogizer mEulogizer;
+    @Mock private Handler mMainHandler;
 
     @Mock private GroupCoalescer mGroupCoalescer;
     @Spy private RecordingCollectionListener mCollectionListener;
@@ -152,6 +156,7 @@ public class NotifCollectionTest extends SysuiTestCase {
                 mClock,
                 mFeatureFlags,
                 mLogger,
+                mMainHandler,
                 mEulogizer,
                 mock(DumpManager.class));
         mCollection.attach(mGroupCoalescer);
@@ -1322,6 +1327,78 @@ public class NotifCollectionTest extends SysuiTestCase {
         verify(mCollectionListener, never()).onEntryRemoved(any(NotificationEntry.class), anyInt());
     }
 
+    private Runnable getInternalNotifUpdateRunnable(StatusBarNotification sbn) {
+        InternalNotifUpdater updater = mCollection.getInternalNotifUpdater("Test");
+        updater.onInternalNotificationUpdate(sbn, "reason");
+        ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mMainHandler).post(runnableCaptor.capture());
+        return runnableCaptor.getValue();
+    }
+
+    @Test
+    public void testGetInternalNotifUpdaterPostsToMainHandler() {
+        InternalNotifUpdater updater = mCollection.getInternalNotifUpdater("Test");
+        updater.onInternalNotificationUpdate(mock(StatusBarNotification.class), "reason");
+        verify(mMainHandler).post(any());
+    }
+
+    @Test
+    public void testSecondPostCallsUpdateWithTrue() {
+        // GIVEN a pipeline with one notification
+        NotifEvent notifEvent = mNoMan.postNotif(buildNotif(TEST_PACKAGE, 47, "myTag"));
+        NotificationEntry entry = mCollectionListener.getEntry(notifEvent.key);
+
+        // KNOWING that it already called listener methods once
+        verify(mCollectionListener).onEntryAdded(eq(entry));
+        verify(mCollectionListener).onRankingApplied();
+
+        // WHEN we update the notification via the system
+        mNoMan.postNotif(buildNotif(TEST_PACKAGE, 47, "myTag"));
+
+        // THEN entry updated gets called, added does not, and ranking is called again
+        verify(mCollectionListener).onEntryUpdated(eq(entry));
+        verify(mCollectionListener).onEntryUpdated(eq(entry), eq(true));
+        verify(mCollectionListener).onEntryAdded((entry));
+        verify(mCollectionListener, times(2)).onRankingApplied();
+    }
+
+    @Test
+    public void testInternalNotifUpdaterCallsUpdate() {
+        // GIVEN a pipeline with one notification
+        NotifEvent notifEvent = mNoMan.postNotif(buildNotif(TEST_PACKAGE, 47, "myTag"));
+        NotificationEntry entry = mCollectionListener.getEntry(notifEvent.key);
+
+        // KNOWING that it will call listener methods once
+        verify(mCollectionListener).onEntryAdded(eq(entry));
+        verify(mCollectionListener).onRankingApplied();
+
+        // WHEN we update that notification internally
+        StatusBarNotification sbn = notifEvent.sbn;
+        getInternalNotifUpdateRunnable(sbn).run();
+
+        // THEN only entry updated gets called a second time
+        verify(mCollectionListener).onEntryAdded(eq(entry));
+        verify(mCollectionListener).onRankingApplied();
+        verify(mCollectionListener).onEntryUpdated(eq(entry));
+        verify(mCollectionListener).onEntryUpdated(eq(entry), eq(false));
+    }
+
+    @Test
+    public void testInternalNotifUpdaterIgnoresNew() {
+        // GIVEN a pipeline without any notifications
+        StatusBarNotification sbn = buildNotif(TEST_PACKAGE, 47, "myTag").build().getSbn();
+
+        // WHEN we internally update an unknown notification
+        getInternalNotifUpdateRunnable(sbn).run();
+
+        // THEN only entry updated gets called a second time
+        verify(mCollectionListener, never()).onEntryAdded(any());
+        verify(mCollectionListener, never()).onRankingUpdate(any());
+        verify(mCollectionListener, never()).onRankingApplied();
+        verify(mCollectionListener, never()).onEntryUpdated(any());
+        verify(mCollectionListener, never()).onEntryUpdated(any(), anyBoolean());
+    }
+
     private static NotificationEntryBuilder buildNotif(String pkg, int id, String tag) {
         return new NotificationEntryBuilder()
                 .setPkg(pkg)
@@ -1371,6 +1448,11 @@ public class NotifCollectionTest extends SysuiTestCase {
             mLastSeenEntries.put(entry.getKey(), entry);
         }
 
+        @Override
+        public void onEntryUpdated(NotificationEntry entry, boolean fromSystem) {
+            onEntryUpdated(entry);
+        }
+
         @Override
         public void onEntryRemoved(NotificationEntry entry, int reason) {
         }

packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt

@@ -0,0 +1,145 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.systemui.statusbar.notification.collection.coordinator
+
+import android.os.Handler
+import android.service.notification.StatusBarNotification
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager
+import com.android.systemui.statusbar.NotificationRemoteInputManager.RemoteInputListener
+import com.android.systemui.statusbar.RemoteInputNotificationRebuilder
+import com.android.systemui.statusbar.SmartReplyController
+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.notification.collection.notifcollection.InternalNotifUpdater
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifCollectionListener
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations.initMocks
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class RemoteInputCoordinatorTest : SysuiTestCase() {
+    private lateinit var coordinator: RemoteInputCoordinator
+    private lateinit var listener: RemoteInputListener
+    private lateinit var collectionListener: NotifCollectionListener
+
+    private lateinit var entry1: NotificationEntry
+    private lateinit var entry2: NotificationEntry
+
+    @Mock private lateinit var lifetimeExtensionCallback: OnEndLifetimeExtensionCallback
+    @Mock private lateinit var rebuilder: RemoteInputNotificationRebuilder
+    @Mock private lateinit var remoteInputManager: NotificationRemoteInputManager
+    @Mock private lateinit var mainHandler: Handler
+    @Mock private lateinit var smartReplyController: SmartReplyController
+    @Mock private lateinit var pipeline: NotifPipeline
+    @Mock private lateinit var notifUpdater: InternalNotifUpdater
+    @Mock private lateinit var dumpManager: DumpManager
+    @Mock private lateinit var sbn: StatusBarNotification
+
+    @Before
+    fun setUp() {
+        initMocks(this)
+        coordinator = RemoteInputCoordinator(
+                dumpManager,
+                rebuilder,
+                remoteInputManager,
+                mainHandler,
+                smartReplyController
+        )
+        `when`(pipeline.addNotificationLifetimeExtender(any())).thenAnswer {
+            (it.arguments[0] as NotifLifetimeExtender).setCallback(lifetimeExtensionCallback)
+        }
+        `when`(pipeline.getInternalNotifUpdater(any())).thenReturn(notifUpdater)
+        coordinator.attach(pipeline)
+        listener = withArgCaptor {
+            verify(remoteInputManager).setRemoteInputListener(capture())
+        }
+        collectionListener = withArgCaptor {
+            verify(pipeline).addCollectionListener(capture())
+        }
+        entry1 = NotificationEntryBuilder().setId(1).build()
+        entry2 = NotificationEntryBuilder().setId(2).build()
+        `when`(rebuilder.rebuildForCanceledSmartReplies(any())).thenReturn(sbn)
+        `when`(rebuilder.rebuildForRemoteInputReply(any())).thenReturn(sbn)
+        `when`(rebuilder.rebuildForSendingSmartReply(any(), any())).thenReturn(sbn)
+    }
+
+    val remoteInputActiveExtender get() = coordinator.mRemoteInputActiveExtender
+    val remoteInputHistoryExtender get() = coordinator.mRemoteInputHistoryExtender
+    val smartReplyHistoryExtender get() = coordinator.mSmartReplyHistoryExtender
+
+    @Test
+    fun testRemoteInputActive() {
+        `when`(remoteInputManager.isRemoteInputActive(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isFalse()
+    }
+
+    @Test
+    fun testRemoteInputHistory() {
+        `when`(remoteInputManager.shouldKeepForRemoteInputHistory(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testSmartReplyHistory() {
+        `when`(remoteInputManager.shouldKeepForSmartReplyHistory(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(remoteInputHistoryExtender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(smartReplyHistoryExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(listener.isNotificationKeptForRemoteInputHistory(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testNotificationWithRemoteInputActiveIsRemovedOnCollapse() {
+        `when`(remoteInputManager.isRemoteInputActive(entry1)).thenReturn(true)
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+
+        // Nothing should happen on panel collapse before we start extending the lifetime
+        listener.onPanelCollapsed()
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+        verify(lifetimeExtensionCallback, never()).onEndLifetimeExtension(any(), any())
+
+        // Start extending lifetime & validate that the extension is ended
+        assertThat(remoteInputActiveExtender.shouldExtendLifetime(entry1, 0)).isTrue()
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isTrue()
+        listener.onPanelCollapsed()
+        verify(lifetimeExtensionCallback).onEndLifetimeExtension(remoteInputActiveExtender, entry1)
+        assertThat(remoteInputActiveExtender.isExtending(entry1.key)).isFalse()
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt

@@ -0,0 +1,230 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.systemui.statusbar.notification.collection.notifcollection
+
+import android.os.Handler
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper.RunWithLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.statusbar.notification.collection.notifcollection.NotifLifetimeExtender.OnEndLifetimeExtensionCallback
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.Mockito.any
+import org.mockito.Mockito.anyLong
+import org.mockito.Mockito.eq
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations.initMocks
+import java.util.function.Consumer
+import java.util.function.Predicate
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class SelfTrackingLifetimeExtenderTest : SysuiTestCase() {
+    private lateinit var extender: TestableSelfTrackingLifetimeExtender
+
+    private lateinit var entry1: NotificationEntry
+    private lateinit var entry2: NotificationEntry
+
+    @Mock
+    private lateinit var callback: OnEndLifetimeExtensionCallback
+    @Mock
+    private lateinit var mainHandler: Handler
+    @Mock
+    private lateinit var shouldExtend: Predicate<NotificationEntry>
+    @Mock
+    private lateinit var onStarted: Consumer<NotificationEntry>
+    @Mock
+    private lateinit var onCanceled: Consumer<NotificationEntry>
+
+    @Before
+    fun setUp() {
+        initMocks(this)
+        extender = TestableSelfTrackingLifetimeExtender()
+        extender.setCallback(callback)
+        entry1 = NotificationEntryBuilder().setId(1).build()
+        entry2 = NotificationEntryBuilder().setId(2).build()
+    }
+
+    @Test
+    fun testName() {
+        assertThat(extender.name).isEqualTo("Testable")
+    }
+
+    @Test
+    fun testNoExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+        verify(onStarted, never()).accept(entry1)
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenCancelForRepost() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        verify(onCanceled, never()).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        extender.cancelLifetimeExtension(entry1)
+        verify(onCanceled).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenCancel_thenEndDoesNothing() {
+        testExtendThenCancelForRepost()
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+
+        extender.endLifetimeExtension(entry1.key)
+        extender.endLifetimeExtensionAfterDelay(entry1.key, 1000)
+        verify(callback, never()).onEndLifetimeExtension(any(), any())
+        verify(mainHandler, never()).postDelayed(any(), anyLong())
+    }
+
+    @Test
+    fun testExtendThenEnd() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        extender.endLifetimeExtension(entry1.key)
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenEndAfterDelay() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+
+        // Call the method and capture the posted runnable
+        extender.endLifetimeExtensionAfterDelay(entry1.key, 1234)
+        val runnable = withArgCaptor<Runnable> {
+            verify(mainHandler).postDelayed(capture(), eq(1234.toLong()))
+        }
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        verify(callback, never()).onEndLifetimeExtension(any(), any())
+
+        // now run the posted runnable and ensure it works as expected
+        runnable.run()
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+        verify(onCanceled, never()).accept(entry1)
+    }
+
+    @Test
+    fun testExtendThenEndAll() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        `when`(shouldExtend.test(entry2)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        assertThat(extender.isExtending(entry2.key)).isFalse()
+        assertThat(extender.shouldExtendLifetime(entry2, 0)).isTrue()
+        verify(onStarted).accept(entry2)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+        assertThat(extender.isExtending(entry2.key)).isTrue()
+        extender.endAllLifetimeExtensions()
+        verify(callback).onEndLifetimeExtension(extender, entry1)
+        verify(callback).onEndLifetimeExtension(extender, entry2)
+        verify(onCanceled, never()).accept(entry1)
+        verify(onCanceled, never()).accept(entry2)
+    }
+
+    @Test
+    fun testExtendWithinEndCanReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        }
+        extender.endLifetimeExtension(entry1.key)
+        verify(onStarted, times(2)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testExtendWithinEndCanNotReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true, false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        }
+        extender.endLifetimeExtension(entry1.key)
+        verify(onStarted, times(1)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+    }
+
+    @Test
+    fun testExtendWithinEndAllCanReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        }
+        extender.endAllLifetimeExtensions()
+        verify(onStarted, times(2)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isTrue()
+    }
+
+    @Test
+    fun testExtendWithinEndAllCanNotReExtend() {
+        `when`(shouldExtend.test(entry1)).thenReturn(true, false)
+        assertThat(extender.shouldExtendLifetime(entry1, 0)).isTrue()
+        verify(onStarted, times(1)).accept(entry1)
+
+        `when`(callback.onEndLifetimeExtension(extender, entry1)).thenAnswer {
+            assertThat(extender.shouldExtendLifetime(entry1, 0)).isFalse()
+        }
+        extender.endAllLifetimeExtensions()
+        verify(onStarted, times(1)).accept(entry1)
+        assertThat(extender.isExtending(entry1.key)).isFalse()
+    }
+
+    inner class TestableSelfTrackingLifetimeExtender(debug: Boolean = false) :
+            SelfTrackingLifetimeExtender("Test", "Testable", debug, mainHandler) {
+
+        override fun queryShouldExtendLifetime(entry: NotificationEntry) =
+                shouldExtend.test(entry)
+
+        override fun onStartedLifetimeExtension(entry: NotificationEntry) {
+            onStarted.accept(entry)
+        }
+
+        override fun onCanceledLifetimeExtension(entry: NotificationEntry) {
+            onCanceled.accept(entry)
+        }
+    }
+}