Merge "Avoid autoboxing."
diff --git a/core/api/current.txt b/core/api/current.txt
index 8e3faaa..d283e0a 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -18163,8 +18163,9 @@
     method @NonNull public java.util.List<android.util.Size> getExtensionSupportedSizes(int, int);
     method @NonNull public java.util.List<java.lang.Integer> getSupportedExtensions();
     field public static final int EXTENSION_AUTOMATIC = 0; // 0x0
-    field public static final int EXTENSION_BEAUTY = 1; // 0x1
+    field @Deprecated public static final int EXTENSION_BEAUTY = 1; // 0x1
     field public static final int EXTENSION_BOKEH = 2; // 0x2
+    field public static final int EXTENSION_FACE_RETOUCH = 1; // 0x1
     field public static final int EXTENSION_HDR = 3; // 0x3
     field public static final int EXTENSION_NIGHT = 4; // 0x4
   }
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index f695b5c4..0fb8bf4 100755
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -384,6 +384,7 @@
     field public static final int config_systemNotificationIntelligence = 17039413; // 0x1040035
     field public static final int config_systemShell = 17039402; // 0x104002a
     field public static final int config_systemSpeechRecognizer = 17039406; // 0x104002e
+    field public static final int config_systemSupervision;
     field public static final int config_systemTelevisionNotificationHandler = 17039409; // 0x1040031
     field public static final int config_systemTextIntelligence = 17039414; // 0x1040036
     field public static final int config_systemUi = 17039418; // 0x104003a
diff --git a/core/java/android/hardware/camera2/CameraExtensionCharacteristics.java b/core/java/android/hardware/camera2/CameraExtensionCharacteristics.java
index 395c655..5c636c7 100644
--- a/core/java/android/hardware/camera2/CameraExtensionCharacteristics.java
+++ b/core/java/android/hardware/camera2/CameraExtensionCharacteristics.java
@@ -15,12 +15,13 @@
  */
 package android.hardware.camera2;
 
+import android.annotation.IntDef;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.ServiceConnection;
-import android.content.pm.PackageManager;
 import android.graphics.ImageFormat;
 import android.hardware.camera2.extension.IAdvancedExtenderImpl;
 import android.hardware.camera2.extension.ICameraExtensionsProxyService;
@@ -35,24 +36,22 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.SystemProperties;
-import android.annotation.IntDef;
-import android.annotation.NonNull;
 import android.util.Log;
 import android.util.Pair;
 import android.util.Range;
 import android.util.Size;
 
-import java.util.HashSet;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.List;
-import java.util.Objects;
 
 /**
  * <p>Allows clients to query availability and supported resolutions of camera extensions.</p>
@@ -96,7 +95,15 @@
      * Device-specific extension implementation which tends to smooth the skin and apply other
      * cosmetic effects to people's faces.
      */
-    public static final int EXTENSION_BEAUTY = 1;
+    public static final int EXTENSION_FACE_RETOUCH = 1;
+
+    /**
+     * Device-specific extension implementation which tends to smooth the skin and apply other
+     * cosmetic effects to people's faces.
+     *
+     * @deprecated Use {@link #EXTENSION_FACE_RETOUCH} instead.
+     */
+    public @Deprecated static final int EXTENSION_BEAUTY = EXTENSION_FACE_RETOUCH;
 
     /**
      * Device-specific extension implementation which can blur certain regions of the final image
@@ -121,7 +128,7 @@
      */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(flag = true, value = {EXTENSION_AUTOMATIC,
-                EXTENSION_BEAUTY,
+                EXTENSION_FACE_RETOUCH,
                 EXTENSION_BOKEH,
                 EXTENSION_HDR,
                 EXTENSION_NIGHT})
@@ -145,7 +152,7 @@
     private static final @Extension
     int[] EXTENSION_LIST = new int[]{
             EXTENSION_AUTOMATIC,
-            EXTENSION_BEAUTY,
+            EXTENSION_FACE_RETOUCH,
             EXTENSION_BOKEH,
             EXTENSION_HDR,
             EXTENSION_NIGHT};
diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index 1853c65..ba9332d 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -757,6 +757,11 @@
         }
 
         @Nullable
+        public Key[] getKeys(@PowerComponent int componentId) {
+            return mData.getKeys(componentId);
+        }
+
+        @Nullable
         public Key getKey(@PowerComponent int componentId, @ProcessState int processState) {
             return mData.getKey(componentId, processState);
         }
diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index da968b3..be36027 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -660,6 +660,7 @@
             mapUidProcessStateToBatteryConsumerProcessState(int processState) {
         switch (processState) {
             case BatteryStats.Uid.PROCESS_STATE_TOP:
+            case BatteryStats.Uid.PROCESS_STATE_FOREGROUND:
                 return BatteryConsumer.PROCESS_STATE_FOREGROUND;
             case BatteryStats.Uid.PROCESS_STATE_BACKGROUND:
             case BatteryStats.Uid.PROCESS_STATE_TOP_SLEEPING:
@@ -1028,6 +1029,16 @@
         public abstract long getCpuMeasuredBatteryConsumptionUC();
 
         /**
+         * Returns the battery consumption (in microcoulombs) of the uid's cpu usage when in the
+         * specified process state.
+         * Will return {@link #POWER_DATA_UNAVAILABLE} if data is unavailable.
+         *
+         * {@hide}
+         */
+        public abstract long getCpuMeasuredBatteryConsumptionUC(
+                @BatteryConsumer.ProcessState int processState);
+
+        /**
          * Returns the battery consumption (in microcoulombs) of the uid's GNSS usage, derived from
          * on device power measurement data.
          * Will return {@link #POWER_DATA_UNAVAILABLE} if data is unavailable.
diff --git a/core/java/android/os/BatteryUsageStatsQuery.java b/core/java/android/os/BatteryUsageStatsQuery.java
index 187b64f..81e49e9 100644
--- a/core/java/android/os/BatteryUsageStatsQuery.java
+++ b/core/java/android/os/BatteryUsageStatsQuery.java
@@ -111,6 +111,10 @@
         return (mFlags & FLAG_BATTERY_USAGE_STATS_POWER_PROFILE_MODEL) != 0;
     }
 
+    public boolean isProcessStateDataNeeded() {
+        return (mFlags & FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0;
+    }
+
     /**
      * Returns the client's tolerance for stale battery stats. The data is allowed to be up to
      * this many milliseconds out-of-date.
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 1e20cfa..9bdf608 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -588,6 +588,10 @@
                     procState = u.mProcessState;
                 }
 
+                if (procState == ActivityManager.PROCESS_STATE_NONEXISTENT) {
+                    continue;
+                }
+
                 final long timestampMs = mClock.elapsedRealtime();
                 final LongArrayMultiStateCounter onBatteryCounter =
                         u.getProcStateTimeCounter().getCounter();
@@ -8599,6 +8603,23 @@
             return mUidMeasuredEnergyStats.getAccumulatedStandardBucketCharge(bucket);
         }
 
+        /**
+         * Returns the battery consumption (in microcoulombs) of this uid for a standard power
+         * bucket and a process state, such as Uid.PROCESS_STATE_TOP.
+         */
+        @GuardedBy("mBsi")
+        public long getMeasuredBatteryConsumptionUC(@StandardPowerBucket int bucket,
+                int processState) {
+            if (mBsi.mGlobalMeasuredEnergyStats == null
+                    || !mBsi.mGlobalMeasuredEnergyStats.isStandardBucketSupported(bucket)) {
+                return POWER_DATA_UNAVAILABLE;
+            }
+            if (mUidMeasuredEnergyStats == null) {
+                return 0L; // It is supported, but was never filled, so it must be 0
+            }
+            return mUidMeasuredEnergyStats.getAccumulatedStandardBucketCharge(bucket, processState);
+        }
+
         @GuardedBy("mBsi")
         @Override
         public long[] getCustomConsumerMeasuredBatteryConsumptionUC() {
@@ -8626,6 +8647,13 @@
 
         @GuardedBy("mBsi")
         @Override
+        public long getCpuMeasuredBatteryConsumptionUC(
+                @BatteryConsumer.ProcessState int processState) {
+            return getMeasuredBatteryConsumptionUC(MeasuredEnergyStats.POWER_BUCKET_CPU,
+                    processState);
+        }
+
+        @Override
         public long getGnssMeasuredBatteryConsumptionUC() {
             return getMeasuredBatteryConsumptionUC(MeasuredEnergyStats.POWER_BUCKET_GNSS);
         }
diff --git a/core/java/com/android/internal/os/CpuPowerCalculator.java b/core/java/com/android/internal/os/CpuPowerCalculator.java
index e693d9d..4599231 100644
--- a/core/java/com/android/internal/os/CpuPowerCalculator.java
+++ b/core/java/com/android/internal/os/CpuPowerCalculator.java
@@ -25,11 +25,13 @@
 import android.util.Log;
 import android.util.SparseArray;
 
+import java.util.Arrays;
 import java.util.List;
 
 public class CpuPowerCalculator extends PowerCalculator {
     private static final String TAG = "CpuPowerCalculator";
     private static final boolean DEBUG = BatteryStatsHelper.DEBUG;
+    private static final BatteryConsumer.Key[] UNINITIALIZED_KEYS = new BatteryConsumer.Key[0];
     private final int mNumCpuClusters;
 
     // Time-in-state based CPU power estimation model computes the estimated power
@@ -44,13 +46,16 @@
     private final UsageBasedPowerEstimator[] mPerClusterPowerEstimators;
     // Multiple estimators per cluster: one per available scaling frequency. Note that different
     // clusters have different sets of frequencies and corresponding power consumption averages.
-    private final UsageBasedPowerEstimator[][] mPerCpuFreqPowerEstimators;
+    private final UsageBasedPowerEstimator[][] mPerCpuFreqPowerEstimatorsByCluster;
+    // Flattened array of estimators across clusters
+    private final UsageBasedPowerEstimator[] mPerCpuFreqPowerEstimators;
 
     private static class Result {
         public long durationMs;
         public double powerMah;
         public long durationFgMs;
         public String packageWithHighestDrain;
+        public double[] perProcStatePowerMah;
     }
 
     public CpuPowerCalculator(PowerProfile profile) {
@@ -65,14 +70,23 @@
                     profile.getAveragePowerForCpuCluster(cluster));
         }
 
-        mPerCpuFreqPowerEstimators = new UsageBasedPowerEstimator[mNumCpuClusters][];
+        int freqCount = 0;
+        for (int cluster = 0; cluster < mNumCpuClusters; cluster++) {
+            freqCount += profile.getNumSpeedStepsInCpuCluster(cluster);
+        }
+
+        mPerCpuFreqPowerEstimatorsByCluster = new UsageBasedPowerEstimator[mNumCpuClusters][];
+        mPerCpuFreqPowerEstimators = new UsageBasedPowerEstimator[freqCount];
+        int index = 0;
         for (int cluster = 0; cluster < mNumCpuClusters; cluster++) {
             final int speedsForCluster = profile.getNumSpeedStepsInCpuCluster(cluster);
-            mPerCpuFreqPowerEstimators[cluster] = new UsageBasedPowerEstimator[speedsForCluster];
+            mPerCpuFreqPowerEstimatorsByCluster[cluster] =
+                    new UsageBasedPowerEstimator[speedsForCluster];
             for (int speed = 0; speed < speedsForCluster; speed++) {
-                mPerCpuFreqPowerEstimators[cluster][speed] =
-                        new UsageBasedPowerEstimator(
-                                profile.getAveragePowerForCpuCore(cluster, speed));
+                final UsageBasedPowerEstimator estimator = new UsageBasedPowerEstimator(
+                        profile.getAveragePowerForCpuCore(cluster, speed));
+                mPerCpuFreqPowerEstimatorsByCluster[cluster][speed] = estimator;
+                mPerCpuFreqPowerEstimators[index++] = estimator;
             }
         }
     }
@@ -82,12 +96,20 @@
             long rawRealtimeUs, long rawUptimeUs, BatteryUsageStatsQuery query) {
         double totalPowerMah = 0;
 
+        BatteryConsumer.Key[] keys = UNINITIALIZED_KEYS;
         Result result = new Result();
         final SparseArray<UidBatteryConsumer.Builder> uidBatteryConsumerBuilders =
                 builder.getUidBatteryConsumerBuilders();
         for (int i = uidBatteryConsumerBuilders.size() - 1; i >= 0; i--) {
             final UidBatteryConsumer.Builder app = uidBatteryConsumerBuilders.valueAt(i);
-            calculateApp(app, app.getBatteryStatsUid(), query, result);
+            if (keys == UNINITIALIZED_KEYS) {
+                if (query.isProcessStateDataNeeded()) {
+                    keys = app.getKeys(BatteryConsumer.POWER_COMPONENT_CPU);
+                } else {
+                    keys = null;
+                }
+            }
+            calculateApp(app, app.getBatteryStatsUid(), query, result, keys);
             totalPowerMah += result.powerMah;
         }
 
@@ -105,7 +127,7 @@
     }
 
     private void calculateApp(UidBatteryConsumer.Builder app, BatteryStats.Uid u,
-            BatteryUsageStatsQuery query, Result result) {
+            BatteryUsageStatsQuery query, Result result, BatteryConsumer.Key[] keys) {
         final long consumptionUC = u.getCpuMeasuredBatteryConsumptionUC();
         final int powerModel = getPowerModel(consumptionUC, query);
         calculatePowerAndDuration(u, powerModel, consumptionUC, BatteryStats.STATS_SINCE_CHARGED,
@@ -114,6 +136,75 @@
         app.setConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU, result.powerMah, powerModel)
                 .setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_CPU, result.durationMs)
                 .setPackageWithHighestDrain(result.packageWithHighestDrain);
+
+        if (query.isProcessStateDataNeeded() && keys != null) {
+            switch (powerModel) {
+                case BatteryConsumer.POWER_MODEL_MEASURED_ENERGY:
+                    calculateMeasuredPowerPerProcessState(app, u, keys);
+                    break;
+                case BatteryConsumer.POWER_MODEL_POWER_PROFILE:
+                    calculateModeledPowerPerProcessState(app, u, keys, result);
+                    break;
+            }
+        }
+    }
+
+    private void calculateMeasuredPowerPerProcessState(UidBatteryConsumer.Builder app,
+            BatteryStats.Uid u, BatteryConsumer.Key[] keys) {
+        for (BatteryConsumer.Key key : keys) {
+            // The key for "PROCESS_STATE_ANY" has already been populated with the
+            // full energy across all states.  We don't want to override it with
+            // the energy for "other" states, which excludes the tracked states like
+            // foreground, background etc.
+            if (key.processState == BatteryConsumer.PROCESS_STATE_ANY) {
+                continue;
+            }
+
+            final long consumptionUC = u.getCpuMeasuredBatteryConsumptionUC(key.processState);
+            if (consumptionUC != 0) {
+                app.setConsumedPower(key, uCtoMah(consumptionUC),
+                        BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+            }
+        }
+    }
+
+    private void calculateModeledPowerPerProcessState(UidBatteryConsumer.Builder app,
+            BatteryStats.Uid u, BatteryConsumer.Key[] keys, Result result) {
+        if (result.perProcStatePowerMah == null) {
+            result.perProcStatePowerMah = new double[BatteryConsumer.PROCESS_STATE_COUNT];
+        } else {
+            Arrays.fill(result.perProcStatePowerMah, 0);
+        }
+
+        for (int uidProcState = 0; uidProcState < BatteryStats.Uid.NUM_PROCESS_STATE;
+                uidProcState++) {
+            @BatteryConsumer.ProcessState int procState =
+                    BatteryStats.mapUidProcessStateToBatteryConsumerProcessState(uidProcState);
+            if (procState == BatteryConsumer.PROCESS_STATE_ANY) {
+                continue;
+            }
+
+            // TODO(b/191921016): use per-state CPU active time
+            final long cpuActiveTime = 0;
+            // TODO(b/191921016): use per-state CPU cluster times
+            final long[] cpuClusterTimes = null;
+
+            final long[] cpuFreqTimes = u.getCpuFreqTimes(BatteryStats.STATS_SINCE_CHARGED,
+                    uidProcState);
+            if (cpuActiveTime != 0 || cpuClusterTimes != null || cpuFreqTimes != null) {
+                result.perProcStatePowerMah[procState] += calculateUidModeledPowerMah(u,
+                        cpuActiveTime, cpuClusterTimes, cpuFreqTimes);
+            }
+        }
+
+        for (BatteryConsumer.Key key : keys) {
+            if (key.processState == BatteryConsumer.PROCESS_STATE_ANY) {
+                continue;
+            }
+
+            app.setConsumedPower(key, result.perProcStatePowerMah[key.processState],
+                    BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+        }
     }
 
     @Override
@@ -145,7 +236,7 @@
         long durationMs = (u.getUserCpuTimeUs(statsType) + u.getSystemCpuTimeUs(statsType)) / 1000;
 
         final double powerMah;
-        switch(powerModel) {
+        switch (powerModel) {
             case BatteryConsumer.POWER_MODEL_MEASURED_ENERGY:
                 powerMah = uCtoMah(consumptionUC);
                 break;
@@ -205,16 +296,21 @@
      * Calculates CPU power consumed by the specified app, using the PowerProfile model.
      */
     public double calculateUidModeledPowerMah(BatteryStats.Uid u, int statsType) {
+        return calculateUidModeledPowerMah(u, u.getCpuActiveTime(), u.getCpuClusterTimes(),
+                u.getCpuFreqTimes(statsType));
+    }
+
+    private double calculateUidModeledPowerMah(BatteryStats.Uid u, long cpuActiveTime,
+            long[] cpuClusterTimes, long[] cpuFreqTimes) {
         // Constant battery drain when CPU is active
-        double powerMah = calculateActiveCpuPowerMah(u.getCpuActiveTime());
+        double powerMah = calculateActiveCpuPowerMah(cpuActiveTime);
 
         // Additional per-cluster battery drain
-        long[] cpuClusterTimes = u.getCpuClusterTimes();
         if (cpuClusterTimes != null) {
             if (cpuClusterTimes.length == mNumCpuClusters) {
                 for (int cluster = 0; cluster < mNumCpuClusters; cluster++) {
-                    double power = calculatePerCpuClusterPowerMah(cluster,
-                            cpuClusterTimes[cluster]);
+                    final double power = mPerClusterPowerEstimators[cluster]
+                            .calculatePower(cpuClusterTimes[cluster]);
                     powerMah += power;
                     if (DEBUG) {
                         Log.d(TAG, "UID " + u.getUid() + ": CPU cluster #" + cluster
@@ -228,21 +324,17 @@
             }
         }
 
-        // Additional per-frequency battery drain
-        for (int cluster = 0; cluster < mNumCpuClusters; cluster++) {
-            final int speedsForCluster = mPerCpuFreqPowerEstimators[cluster].length;
-            for (int speed = 0; speed < speedsForCluster; speed++) {
-                final long timeUs = u.getTimeAtCpuSpeed(cluster, speed, statsType);
-                final double power = calculatePerCpuFreqPowerMah(cluster, speed,
-                        timeUs / 1000);
-                if (DEBUG) {
-                    Log.d(TAG, "UID " + u.getUid() + ": CPU cluster #" + cluster + " step #"
-                            + speed + " timeUs=" + timeUs + " power="
-                            + formatCharge(power));
+        if (cpuFreqTimes != null) {
+            if (cpuFreqTimes.length == mPerCpuFreqPowerEstimators.length) {
+                for (int i = 0; i < cpuFreqTimes.length; i++) {
+                    powerMah += mPerCpuFreqPowerEstimators[i].calculatePower(cpuFreqTimes[i]);
                 }
-                powerMah += power;
+            } else {
+                Log.w(TAG, "UID " + u.getUid() + " CPU freq # mismatch: Power Profile # "
+                        + mPerCpuFreqPowerEstimators.length + " actual # " + cpuFreqTimes.length);
             }
         }
+
         return powerMah;
     }
 
@@ -259,7 +351,7 @@
     /**
      * Calculates CPU cluster power consumption.
      *
-     * @param cluster CPU cluster used.
+     * @param cluster           CPU cluster used.
      * @param clusterDurationMs duration of CPU cluster usage.
      * @return a double in milliamp-hours of estimated CPU cluster power consumption.
      */
@@ -270,14 +362,14 @@
     /**
      * Calculates CPU cluster power consumption at a specific speedstep.
      *
-     * @param cluster CPU cluster used.
-     * @param speedStep which speedstep used.
+     * @param cluster                 CPU cluster used.
+     * @param speedStep               which speedstep used.
      * @param clusterSpeedDurationsMs duration of CPU cluster usage at the specified speed step.
      * @return a double in milliamp-hours of estimated CPU cluster-speed power consumption.
      */
     public double calculatePerCpuFreqPowerMah(int cluster, int speedStep,
             long clusterSpeedDurationsMs) {
-        return mPerCpuFreqPowerEstimators[cluster][speedStep].calculatePower(
+        return mPerCpuFreqPowerEstimatorsByCluster[cluster][speedStep].calculatePower(
                 clusterSpeedDurationsMs);
     }
 }
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index e4b3991..fceb951 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -4697,6 +4697,9 @@
     only. The component must be part of a system app. -->
     <string name="config_defaultSupervisionProfileOwnerComponent" translatable="false"></string>
 
+    <!-- The package name of the default supervision package. -->
+    <string name="config_systemSupervision" translatable="false"></string>
+
     <!-- Trigger a warning for notifications with RemoteView objects that are larger in bytes than
     this value (default 1MB)-->
     <integer name="config_notificationWarnRemoteViewSizeBytes">2000000</integer>
diff --git a/core/res/res/values/public.xml b/core/res/res/values/public.xml
index ed49fe4..366dccb 100644
--- a/core/res/res/values/public.xml
+++ b/core/res/res/values/public.xml
@@ -3315,7 +3315,9 @@
   <staging-public-group type="style" first-id="0x0dfd0000">
   </staging-public-group>
 
-  <staging-public-group type="string" first-id="0x0dfc0000">
+  <staging-public-group type="string" first-id="0x01dc0000">
+    <!-- @hide @SystemApi -->
+    <public name="config_systemSupervision" />
   </staging-public-group>
 
   <staging-public-group type="dimen" first-id="0x01db0000">
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
index 0c7218e..eb378b9 100644
--- a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
+++ b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
@@ -34,7 +34,9 @@
     enum EntryType {
         UID_TOTAL_POWER,
         UID_POWER_MODELED,
+        UID_POWER_MODELED_PROCESS_STATE,
         UID_POWER_MEASURED,
+        UID_POWER_MEASURED_PROCESS_STATE,
         UID_POWER_CUSTOM,
         UID_DURATION,
         DEVICE_TOTAL_POWER,
@@ -135,10 +137,16 @@
                         requestedBatteryConsumer.getConsumedPower(component),
                         totalPowerByComponentMah[component]
                 );
+                addProcessStateEntries(metricTitle, EntryType.UID_POWER_MEASURED_PROCESS_STATE,
+                        requestedBatteryConsumer, component
+                );
                 addEntry(metricTitle + " (modeled)", EntryType.UID_POWER_MODELED,
                         requestedModeledBatteryConsumer.getConsumedPower(component),
                         totalModeledPowerByComponentMah[component]
                 );
+                addProcessStateEntries(metricTitle, EntryType.UID_POWER_MODELED_PROCESS_STATE,
+                        requestedModeledBatteryConsumer, component
+                );
             }
         }
 
@@ -164,6 +172,33 @@
                 batteryConsumerId, context.getPackageManager());
     }
 
+    private void addProcessStateEntries(String metricTitle, EntryType entryType,
+            BatteryConsumer batteryConsumer, int component) {
+        final BatteryConsumer.Key[] keys = batteryConsumer.getKeys(component);
+        if (keys == null || keys.length <= 1) {
+            return;
+        }
+
+        for (BatteryConsumer.Key key : keys) {
+            String label;
+            switch (key.processState) {
+                case BatteryConsumer.PROCESS_STATE_FOREGROUND:
+                    label = "foreground";
+                    break;
+                case BatteryConsumer.PROCESS_STATE_BACKGROUND:
+                    label = "background";
+                    break;
+                case BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE:
+                    label = "FGS";
+                    break;
+                default:
+                    continue;
+            }
+            addEntry(metricTitle + " \u2022 " + label, entryType,
+                    batteryConsumer.getConsumedPower(key), 0);
+        }
+    }
+
     private void populateForAggregateBatteryConsumer(Context context,
             List<BatteryUsageStats> batteryUsageStatsList) {
         BatteryUsageStats batteryUsageStats = batteryUsageStatsList.get(0);
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryStatsViewerActivity.java b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryStatsViewerActivity.java
index 33ce6bf..0f45c3b 100644
--- a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryStatsViewerActivity.java
+++ b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryStatsViewerActivity.java
@@ -138,12 +138,14 @@
             final BatteryUsageStatsQuery queryDefault =
                     new BatteryUsageStatsQuery.Builder()
                             .includePowerModels()
+                            .includeProcessStateData()
                             .setMaxStatsAgeMs(maxStatsAgeMs)
                             .build();
             final BatteryUsageStatsQuery queryPowerProfileModeledOnly =
                     new BatteryUsageStatsQuery.Builder()
                             .powerProfileModeledOnly()
                             .includePowerModels()
+                            .includeProcessStateData()
                             .setMaxStatsAgeMs(maxStatsAgeMs)
                             .build();
             return mBatteryStatsManager.getBatteryUsageStats(
@@ -290,6 +292,13 @@
                     setPowerText(viewHolder.value1TextView, entry.value1);
                     setProportionText(viewHolder.value2TextView, entry);
                     break;
+                case UID_POWER_MODELED_PROCESS_STATE:
+                    setTitleIconAndBackground(viewHolder, "    " + entry.title,
+                            R.drawable.gm_calculate_24,
+                            R.color.battery_consumer_bg_power_profile);
+                    setPowerText(viewHolder.value1TextView, entry.value1);
+                    viewHolder.value2TextView.setVisibility(View.INVISIBLE);
+                    break;
                 case UID_POWER_MEASURED:
                     setTitleIconAndBackground(viewHolder, entry.title,
                             R.drawable.gm_amp_24,
@@ -297,6 +306,13 @@
                     setPowerText(viewHolder.value1TextView, entry.value1);
                     setProportionText(viewHolder.value2TextView, entry);
                     break;
+                case UID_POWER_MEASURED_PROCESS_STATE:
+                    setTitleIconAndBackground(viewHolder, "    " + entry.title,
+                            R.drawable.gm_amp_24,
+                            R.color.battery_consumer_bg_measured_energy);
+                    setPowerText(viewHolder.value1TextView, entry.value1);
+                    viewHolder.value2TextView.setVisibility(View.INVISIBLE);
+                    break;
                 case UID_POWER_CUSTOM:
                     setTitleIconAndBackground(viewHolder, entry.title,
                             R.drawable.gm_custom_24,
diff --git a/core/tests/coretests/src/com/android/internal/os/CpuPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/CpuPowerCalculatorTest.java
index 152d246..8540d91d 100644
--- a/core/tests/coretests/src/com/android/internal/os/CpuPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/CpuPowerCalculatorTest.java
@@ -21,13 +21,18 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.BatteryUsageStatsQuery;
 import android.os.Process;
 import android.os.UidBatteryConsumer;
+import android.util.SparseArray;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -43,12 +48,15 @@
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
+@SuppressWarnings("GuardedBy")
 public class CpuPowerCalculatorTest {
     private static final double PRECISION = 0.00001;
 
     private static final int APP_UID1 = Process.FIRST_APPLICATION_UID + 42;
     private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 272;
 
+    private static final int NUM_CPU_FREQS = 2 + 2;  // 2 clusters * 2 freqs each
+
     @Rule
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
             .setAveragePower(PowerProfile.POWER_CPU_ACTIVE, 720)
@@ -79,6 +87,8 @@
     private KernelCpuUidTimeReader.KernelCpuUidActiveTimeReader mMockKerneCpuUidActiveTimeReader;
     @Mock
     private SystemServerCpuThreadReader mMockSystemServerCpuThreadReader;
+    @Mock
+    private KernelSingleUidTimeReader mMockKernelSingleUidTimeReader;
 
     @Before
     public void setUp() {
@@ -95,6 +105,7 @@
                 .setKernelCpuUidClusterTimeReader(mMockKernelCpuUidClusterTimeReader)
                 .setKernelCpuUidUserSysTimeReader(mMockKernelCpuUidUserSysTimeReader)
                 .setKernelCpuUidActiveTimeReader(mMockKerneCpuUidActiveTimeReader)
+                .setKernelSingleUidTimeReader(mMockKernelSingleUidTimeReader)
                 .setSystemServerCpuThreadReader(mMockSystemServerCpuThreadReader)
                 .initMeasuredEnergyStatsLocked(supportedPowerBuckets, new String[0]);
     }
@@ -133,6 +144,14 @@
             return null;
         }).when(mMockKernelCpuUidClusterTimeReader).readDelta(anyBoolean(), any());
 
+        // Per-frequency CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            callback.onUidCpuTime(APP_UID1, new long[]{1100, 11, 2200, 22});
+            callback.onUidCpuTime(APP_UID2, new long[]{3300, 33, 4400, 44});
+            return null;
+        }).when(mMockCpuUidFreqTimeReader).readDelta(anyBoolean(), any());
+
         mStatsRule.getBatteryStats().updateCpuTimeLocked(true, true, null);
 
         mStatsRule.getUidStats(APP_UID1).getProcessStatsLocked("foo").addCpuTimeLocked(4321, 1234);
@@ -147,7 +166,7 @@
         assertThat(uidConsumer1.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(3333);
         assertThat(uidConsumer1.getConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU))
-                .isWithin(PRECISION).of(1.092233);
+                .isWithin(PRECISION).of(1.031677);
         assertThat(uidConsumer1.getPowerModel(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
         assertThat(uidConsumer1.getPackageWithHighestDrain()).isEqualTo("bar");
@@ -156,20 +175,20 @@
         assertThat(uidConsumer2.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(7777);
         assertThat(uidConsumer2.getConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU))
-                .isWithin(PRECISION).of(2.672322);
+                .isWithin(PRECISION).of(2.489544);
         assertThat(uidConsumer2.getPowerModel(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
         assertThat(uidConsumer2.getPackageWithHighestDrain()).isNull();
 
         final BatteryConsumer deviceBatteryConsumer = mStatsRule.getDeviceBatteryConsumer();
         assertThat(deviceBatteryConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU))
-                .isWithin(PRECISION).of(3.76455);
+                .isWithin(PRECISION).of(3.52122);
         assertThat(deviceBatteryConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
 
         final BatteryConsumer appsBatteryConsumer = mStatsRule.getAppsBatteryConsumer();
         assertThat(appsBatteryConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU))
-                .isWithin(PRECISION).of(3.76455);
+                .isWithin(PRECISION).of(3.52122);
         assertThat(appsBatteryConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
@@ -249,4 +268,189 @@
         assertThat(appsBatteryConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_CPU))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
     }
+
+    @Test
+    public void testTimerBasedModel_byProcessState() {
+        mStatsRule.getBatteryStats().setTrackingCpuByProcStateEnabled(true);
+
+        when(mMockUserInfoProvider.exists(anyInt())).thenReturn(true);
+
+        when(mMockCpuUidFreqTimeReader.allUidTimesAvailable()).thenReturn(true);
+        when(mMockCpuUidFreqTimeReader.readFreqs(any())).thenReturn(new long[]{100, 200, 300, 400});
+
+        when(mMockKernelSingleUidTimeReader.singleUidCpuTimesAvailable()).thenReturn(true);
+
+        SparseArray<long[]> allUidCpuFreqTimeMs = new SparseArray<>();
+        allUidCpuFreqTimeMs.put(APP_UID1, new long[0]);
+        allUidCpuFreqTimeMs.put(APP_UID2, new long[0]);
+        when(mMockCpuUidFreqTimeReader.getAllUidCpuFreqTimeMs()).thenReturn(allUidCpuFreqTimeMs);
+
+        mStatsRule.setTime(1000, 1000);
+
+        mStatsRule.getUidStats(APP_UID1).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_FOREGROUND, 1000);
+        mStatsRule.getUidStats(APP_UID2).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_BACKGROUND, 1000);
+
+        // Initialize time-in-state counts to 0
+        mockSingleUidTimeReader(APP_UID1, new long[NUM_CPU_FREQS]);
+        mockSingleUidTimeReader(APP_UID2, new long[NUM_CPU_FREQS]);
+
+        mStatsRule.getBatteryStats().copyFromAllUidsCpuTimes(true, true);
+
+        mockSingleUidTimeReader(APP_UID1, new long[]{1000, 2000, 3000, 4000});
+        mockSingleUidTimeReader(APP_UID2, new long[]{1111, 2222, 3333, 4444});
+
+        mStatsRule.setTime(2000, 2000);
+        mStatsRule.getBatteryStats().updateCpuTimeLocked(true, true, null);
+        mStatsRule.getBatteryStats().copyFromAllUidsCpuTimes(true, true);
+
+        mockSingleUidTimeReader(APP_UID1, new long[] {5000, 6000, 7000, 8000});
+        mockSingleUidTimeReader(APP_UID2, new long[]{5555, 6666, 7777, 8888});
+
+        mStatsRule.getUidStats(APP_UID1).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_FOREGROUND_SERVICE, 2000);
+        mStatsRule.getUidStats(APP_UID2).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_TOP, 2000);
+
+        mStatsRule.setTime(3000, 3000);
+        mStatsRule.getBatteryStats().updateCpuTimeLocked(true, true, null);
+        mStatsRule.getBatteryStats().copyFromAllUidsCpuTimes(true, true);
+
+        CpuPowerCalculator calculator =
+                new CpuPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(new BatteryUsageStatsQuery.Builder()
+                .powerProfileModeledOnly()
+                .includePowerModels()
+                .includeProcessStateData()
+                .build(), calculator);
+
+        UidBatteryConsumer uidConsumer1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        UidBatteryConsumer uidConsumer2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+
+        final BatteryConsumer.Key foreground = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND);
+        final BatteryConsumer.Key background = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND);
+        final BatteryConsumer.Key fgs = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE);
+
+        assertThat(uidConsumer1.getConsumedPower(foreground)).isWithin(PRECISION).of(1.388888);
+        assertThat(uidConsumer1.getConsumedPower(background)).isWithin(PRECISION).of(0);
+        assertThat(uidConsumer1.getConsumedPower(fgs)).isWithin(PRECISION).of(2.0);
+        assertThat(uidConsumer2.getConsumedPower(foreground)).isWithin(PRECISION).of(2.222);
+        assertThat(uidConsumer2.getConsumedPower(background)).isWithin(PRECISION).of(1.543055);
+        assertThat(uidConsumer2.getConsumedPower(fgs)).isWithin(PRECISION).of(0);
+    }
+
+    private void mockSingleUidTimeReader(int uid, long[] cpuTimes) {
+        doAnswer(invocation -> {
+            LongArrayMultiStateCounter counter = invocation.getArgument(1);
+            long timestampMs = invocation.getArgument(2);
+            LongArrayMultiStateCounter.LongArrayContainer container =
+                    new LongArrayMultiStateCounter.LongArrayContainer(NUM_CPU_FREQS);
+            container.setValues(cpuTimes);
+            counter.updateValues(container, timestampMs);
+            return null;
+        }).when(mMockKernelSingleUidTimeReader).addDelta(eq(uid),
+                any(LongArrayMultiStateCounter.class), anyLong());
+    }
+
+    @Test
+    public void testMeasuredEnergyBasedModel_perProcessState() {
+        when(mMockUserInfoProvider.exists(anyInt())).thenReturn(true);
+
+        when(mMockKernelCpuSpeedReaders[0].readDelta()).thenReturn(new long[]{1000, 2000});
+        when(mMockKernelCpuSpeedReaders[1].readDelta()).thenReturn(new long[]{3000, 4000});
+
+        when(mMockCpuUidFreqTimeReader.perClusterTimesAvailable()).thenReturn(false);
+
+        mStatsRule.setTime(1000, 1000);
+
+        mStatsRule.getUidStats(APP_UID1).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_FOREGROUND, 1000);
+        mStatsRule.getUidStats(APP_UID2).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_BACKGROUND, 1000);
+
+        // User/System CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            // User/system time in microseconds
+            callback.onUidCpuTime(APP_UID1, new long[]{1111000, 2222000});
+            callback.onUidCpuTime(APP_UID2, new long[]{3333000, 4444000});
+            return null;
+        }).when(mMockKernelCpuUidUserSysTimeReader).readDelta(anyBoolean(), any());
+
+        // Per-frequency CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            callback.onUidCpuTime(APP_UID1, new long[]{1100, 11, 2200, 22});
+            callback.onUidCpuTime(APP_UID2, new long[]{3300, 33, 4400, 44});
+            return null;
+        }).when(mMockCpuUidFreqTimeReader).readDelta(anyBoolean(), any());
+
+        mStatsRule.setTime(2000, 2000);
+        final long[] clusterChargesUC = new long[]{13577531, 24688642};
+        mStatsRule.getBatteryStats().updateCpuTimeLocked(true, true, clusterChargesUC);
+
+        mStatsRule.getUidStats(APP_UID1).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_FOREGROUND_SERVICE, 2000);
+        mStatsRule.getUidStats(APP_UID2).setProcessStateForTest(
+                BatteryStats.Uid.PROCESS_STATE_TOP, 2000);
+
+        // User/System CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            // User/system time in microseconds
+            callback.onUidCpuTime(APP_UID1, new long[]{5555000, 6666000});
+            callback.onUidCpuTime(APP_UID2, new long[]{7777000, 8888000});
+            return null;
+        }).when(mMockKernelCpuUidUserSysTimeReader).readDelta(anyBoolean(), any());
+
+        // Per-frequency CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            callback.onUidCpuTime(APP_UID1, new long[]{5500, 55, 6600, 66});
+            callback.onUidCpuTime(APP_UID2, new long[]{7700, 77, 8800, 88});
+            return null;
+        }).when(mMockCpuUidFreqTimeReader).readDelta(anyBoolean(), any());
+
+        mStatsRule.setTime(3000, 3000);
+
+        clusterChargesUC[0] += 10000000;
+        clusterChargesUC[1] += 20000000;
+        mStatsRule.getBatteryStats().updateCpuTimeLocked(true, true, clusterChargesUC);
+
+        CpuPowerCalculator calculator =
+                new CpuPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(new BatteryUsageStatsQuery.Builder()
+                .includePowerModels()
+                .includeProcessStateData()
+                .build(), calculator);
+
+        UidBatteryConsumer uidConsumer1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        UidBatteryConsumer uidConsumer2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+
+        final BatteryConsumer.Key foreground = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND);
+        final BatteryConsumer.Key background = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND);
+        final BatteryConsumer.Key fgs = uidConsumer1.getKey(
+                BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE);
+
+        assertThat(uidConsumer1.getConsumedPower(foreground)).isWithin(PRECISION).of(3.18884);
+        assertThat(uidConsumer1.getConsumedPower(background)).isWithin(PRECISION).of(0);
+        assertThat(uidConsumer1.getConsumedPower(fgs)).isWithin(PRECISION).of(8.02273);
+        assertThat(uidConsumer2.getConsumedPower(foreground)).isWithin(PRECISION).of(10.94009);
+        assertThat(uidConsumer2.getConsumedPower(background)).isWithin(PRECISION).of(7.44064);
+        assertThat(uidConsumer2.getConsumedPower(fgs)).isWithin(PRECISION).of(0);
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/SystemServicePowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/SystemServicePowerCalculatorTest.java
index a36d9fe..33222dd 100644
--- a/core/tests/coretests/src/com/android/internal/os/SystemServicePowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/SystemServicePowerCalculatorTest.java
@@ -47,6 +47,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
+@SuppressWarnings("GuardedBy")
 public class SystemServicePowerCalculatorTest {
 
     private static final double PRECISION = 0.000001;
@@ -78,6 +79,8 @@
     private KernelCpuUidTimeReader.KernelCpuUidActiveTimeReader mMockKerneCpuUidActiveTimeReader;
     @Mock
     private SystemServerCpuThreadReader mMockSystemServerCpuThreadReader;
+    @Mock
+    private KernelSingleUidTimeReader mMockKernelSingleUidTimeReader;
 
     private final KernelCpuSpeedReader[] mMockKernelCpuSpeedReaders = new KernelCpuSpeedReader[]{
             mock(KernelCpuSpeedReader.class),
@@ -96,6 +99,7 @@
                 .setKernelCpuUidClusterTimeReader(mMockKernelCpuUidClusterTimeReader)
                 .setKernelCpuUidUserSysTimeReader(mMockKernelCpuUidUserSysTimeReader)
                 .setKernelCpuUidActiveTimeReader(mMockKerneCpuUidActiveTimeReader)
+                .setKernelSingleUidTimeReader(mMockKernelSingleUidTimeReader)
                 .setSystemServerCpuThreadReader(mMockSystemServerCpuThreadReader);
     }
 
@@ -142,19 +146,19 @@
 
         assertThat(mStatsRule.getUidBatteryConsumer(APP_UID1)
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_SYSTEM_SERVICES))
-                .isWithin(PRECISION).of(1.979351);
+                .isWithin(PRECISION).of(2.105425);
         assertThat(mStatsRule.getUidBatteryConsumer(APP_UID2)
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_SYSTEM_SERVICES))
-                .isWithin(PRECISION).of(17.814165);
+                .isWithin(PRECISION).of(18.948825);
         assertThat(mStatsRule.getUidBatteryConsumer(Process.SYSTEM_UID)
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS))
-                .isWithin(PRECISION).of(-19.793517);
+                .isWithin(PRECISION).of(-21.054250);
         assertThat(mStatsRule.getDeviceBatteryConsumer()
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_SYSTEM_SERVICES))
-                .isWithin(PRECISION).of(19.793517);
+                .isWithin(PRECISION).of(21.054250);
         assertThat(mStatsRule.getAppsBatteryConsumer()
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_SYSTEM_SERVICES))
-                .isWithin(PRECISION).of(19.793517);
+                .isWithin(PRECISION).of(21.054250);
     }
 
     private void prepareBatteryStats(long[] clusterChargesUc) {
@@ -192,6 +196,17 @@
             return null;
         }).when(mMockKernelCpuUidClusterTimeReader).readDelta(anyBoolean(), any());
 
+        when(mMockKernelSingleUidTimeReader.singleUidCpuTimesAvailable()).thenReturn(true);
+
+        // Per-frequency CPU time
+        doAnswer(invocation -> {
+            final KernelCpuUidTimeReader.Callback<long[]> callback = invocation.getArgument(1);
+            callback.onUidCpuTime(APP_UID1, new long[]{1100, 11, 2200, 22});
+            callback.onUidCpuTime(APP_UID2, new long[]{3300, 33, 4400, 44});
+            callback.onUidCpuTime(Process.SYSTEM_UID, new long[]{20_000, 30_000, 40_000, 40_000});
+            return null;
+        }).when(mMockCpuUidFreqTimeReader).readDelta(anyBoolean(), any());
+
         // System service CPU time
         final SystemServerCpuThreadReader.SystemServiceCpuThreadTimes threadTimes =
                 new SystemServerCpuThreadReader.SystemServiceCpuThreadTimes();
diff --git a/media/java/android/media/tv/tuner/frontend/DtmbFrontendSettings.java b/media/java/android/media/tv/tuner/frontend/DtmbFrontendSettings.java
index 91102d4..6b5d6ca 100644
--- a/media/java/android/media/tv/tuner/frontend/DtmbFrontendSettings.java
+++ b/media/java/android/media/tv/tuner/frontend/DtmbFrontendSettings.java
@@ -269,7 +269,7 @@
     /**
      * Gets Code Rate.
      */
-    @Modulation
+    @CodeRate
     public int getCodeRate() {
         return mCodeRate;
     }
@@ -277,7 +277,7 @@
     /**
      * Gets Transmission Mode.
      */
-    @Modulation
+    @TransmissionMode
     public int getTransmissionMode() {
         return mTransmissionMode;
     }
@@ -285,7 +285,7 @@
     /**
      * Gets Bandwidth.
      */
-    @Modulation
+    @Bandwidth
     public int getBandwidth() {
         return mBandwidth;
     }
@@ -293,16 +293,15 @@
     /**
      * Gets Time Interleave Mode.
      */
-    @Modulation
+    @TimeInterleaveMode
     public int getTimeInterleaveMode() {
         return mTimeInterleaveMode;
     }
 
-
     /**
      * Gets Guard Interval.
      */
-    @Modulation
+    @GuardInterval
     public int getGuardInterval() {
         return mGuardInterval;
     }
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 9569cf9..ae8439f 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -52,12 +52,18 @@
 
 filegroup {
     name: "ReleaseJavaFiles",
-    srcs: ["src/com/android/systemui/flags/FeatureFlagManager.java"],
+    srcs: [
+        "src-release/**/*.kt",
+        "src-release/**/*.java",
+    ],
 }
 
 filegroup {
     name: "DebugJavaFiles",
-    srcs: ["src-debug/com/android/systemui/flags/FeatureFlagManager.java"],
+    srcs: [
+        "src-debug/**/*.kt",
+        "src-debug/**/*.java",
+    ],
 }
 
 android_library {
@@ -66,6 +72,8 @@
         "src/**/*.kt",
         "src/**/*.java",
         "src/**/I*.aidl",
+        "src-release/**/*.kt",
+        "src-release/**/*.java",
     ],
     product_variables: {
         debuggable: {
@@ -171,6 +179,8 @@
         "src/**/*.kt",
         "src/**/*.java",
         "src/**/I*.aidl",
+        "src-release/**/*.kt",
+        "src-release/**/*.java",
     ],
     static_libs: [
         "WifiTrackerLib",
diff --git a/packages/SystemUI/res/layout/ongoing_call_chip.xml b/packages/SystemUI/res/layout/ongoing_call_chip.xml
index 5389d9b..c949ba0 100644
--- a/packages/SystemUI/res/layout/ongoing_call_chip.xml
+++ b/packages/SystemUI/res/layout/ongoing_call_chip.xml
@@ -21,6 +21,7 @@
     android:layout_width="wrap_content"
     android:layout_height="match_parent"
     android:layout_gravity="center_vertical|start"
+    android:layout_marginStart="5dp"
 >
     <LinearLayout
         android:id="@+id/ongoing_call_chip_background"
diff --git a/packages/SystemUI/src-debug/com/android/systemui/flags/FeatureFlagManager.java b/packages/SystemUI/src-debug/com/android/systemui/flags/FeatureFlagManager.java
index 3a8ee29..1eeb516 100644
--- a/packages/SystemUI/src-debug/com/android/systemui/flags/FeatureFlagManager.java
+++ b/packages/SystemUI/src-debug/com/android/systemui/flags/FeatureFlagManager.java
@@ -75,6 +75,7 @@
     }
 
     /** Return a {@link BooleanFlag}'s value. */
+    @Override
     public boolean isEnabled(int id, boolean defaultValue) {
         if (!mBooleanFlagCache.containsKey(id)) {
             Boolean result = isEnabledInternal(id);
@@ -105,6 +106,7 @@
     }
 
     /** Set whether a given {@link BooleanFlag} is enabled or not. */
+    @Override
     public void setEnabled(int id, boolean value) {
         Boolean currentValue = isEnabledInternal(id);
         if (currentValue != null && currentValue == value) {
@@ -136,8 +138,10 @@
         Log.i(TAG, "Erase id " + id);
     }
 
+    @Override
     public void addListener(Listener run) {}
 
+    @Override
     public void removeListener(Listener run) {}
 
     private void restartSystemUI() {
@@ -198,6 +202,7 @@
 
     @Override
     public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw, @NonNull String[] args) {
+        pw.println("can override: true");
         ArrayList<String> flagStrings = new ArrayList<>(mBooleanFlagCache.size());
         for (Map.Entry<Integer, Boolean> entry : mBooleanFlagCache.entrySet()) {
             flagStrings.add("  sysui_flag_" + entry.getKey() + ": " + entry.getValue());
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagManager.java b/packages/SystemUI/src-release/com/android/systemui/flags/FeatureFlagManager.java
similarity index 89%
rename from packages/SystemUI/src/com/android/systemui/flags/FeatureFlagManager.java
rename to packages/SystemUI/src-release/com/android/systemui/flags/FeatureFlagManager.java
index 78f0b5f..e501a07 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FeatureFlagManager.java
+++ b/packages/SystemUI/src-release/com/android/systemui/flags/FeatureFlagManager.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.flags;
 
+import android.content.Context;
 import android.util.SparseBooleanArray;
 
 import androidx.annotation.NonNull;
@@ -39,21 +40,21 @@
 public class FeatureFlagManager implements FlagReader, FlagWriter, Dumpable {
     SparseBooleanArray mAccessedFlags = new SparseBooleanArray();
     @Inject
-    public FeatureFlagManager(DumpManager dumpManager) {
+    public FeatureFlagManager(SystemPropertiesHelper systemPropertiesHelper, Context context,
+            DumpManager dumpManager) {
         dumpManager.registerDumpable("SysUIFlags", this);
     }
-    public boolean isEnabled(String key, boolean defaultValue) {
-        return defaultValue;
-    }
+    @Override
     public boolean isEnabled(int key, boolean defaultValue) {
         mAccessedFlags.append(key, defaultValue);
         return defaultValue;
     }
-    public void setEnabled(String key, boolean value) {}
+    @Override
     public void setEnabled(int key, boolean value) {}
 
     @Override
     public void dump(@NonNull FileDescriptor fd, @NonNull PrintWriter pw, @NonNull String[] args) {
+        pw.println("can override: false");
         int size = mAccessedFlags.size();
         for (int i = 0; i < size; i++) {
             pw.println("  sysui_flag_" + mAccessedFlags.keyAt(i)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index 18a3d86..1ce7f03 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/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.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.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 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 @@
     private static final boolean DEBUG = false;
     private static final String TAG = "NotifRemoteInputManager";
 
-    /**
-     * 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<>();
+    private RemoteInputListener mRemoteInputListener;
 
     // Dependencies:
     private final NotificationLockscreenUserManager mLockscreenUserManager;
@@ -125,18 +102,17 @@
     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 @@
                 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 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 @@
             ActionClickLogger logger,
             DumpManager dumpManager) {
         mContext = context;
+        mFeatureFlags = featureFlags;
         mLockscreenUserManager = lockscreenUserManager;
         mSmartReplyController = smartReplyController;
         mEntryManager = notificationEntryManager;
@@ -303,7 +283,11 @@
         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 @@
         });
     }
 
+    /** 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 @@
         mRemoteInputController.addCallback(new RemoteInputController.Callback() {
             @Override
             public void onRemoteInputSent(NotificationEntry entry) {
-                if (FORCE_REMOTE_INPUT_HISTORY
-                        && isNotificationKeptForRemoteInputHistory(entry.getKey())) {
-                    mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-                } else if (mEntriesKeptForRemoteInputActive.contains(entry)) {
-                    // We're currently holding onto this notification, but from the apps point of
-                    // view it is already canceled, so we'll need to cancel it on the apps behalf
-                    // after sending - unless the app posts an update in the mean time, so wait a
-                    // bit.
-                    mMainHandler.postDelayed(() -> {
-                        if (mEntriesKeptForRemoteInputActive.remove(entry)) {
-                            mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-                        }
-                    }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
+                if (mRemoteInputListener != null) {
+                    mRemoteInputListener.onRemoteInputSent(entry);
                 }
                 try {
                     mBarService.onNotificationDirectReplied(entry.getSbn().getKey());
@@ -381,12 +379,12 @@
                 }
             }
         });
-        mSmartReplyController.setCallback((entry, reply) -> {
-            StatusBarNotification newSbn =
-                    rebuildNotificationWithRemoteInputInserted(entry, reply, true /* showSpinner */,
-                            null /* mimeType */, null /* uri */);
-            mEntryManager.updateNotification(newSbn, null /* ranking */);
-        });
+        if (!mFeatureFlags.isNewNotifPipelineRenderingEnabled()) {
+            mSmartReplyController.setCallback((entry, reply) -> {
+                StatusBarNotification newSbn = mRebuilder.rebuildForSendingSmartReply(entry, reply);
+                mEntryManager.updateNotification(newSbn, null /* ranking */);
+            });
+        }
     }
 
     public void addControllerCallback(RemoteInputController.Callback callback) {
@@ -574,51 +572,47 @@
         if (v == null) {
             return null;
         }
-        return (RemoteInputView) 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());
+        return v.findViewWithTag(RemoteInputView.VIEW_TAG);
     }
 
     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)) {
-            mKeysKeptForRemoteInputHistory.remove(key);
-        }
+        // OLD pipeline code ONLY; can assume implementation
+        ((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++) {
-            NotificationEntry entry = mEntriesKeptForRemoteInputActive.valueAt(i);
-            if (mRemoteInputController != null) {
-                mRemoteInputController.removeRemoteInput(entry, null);
-            }
-            if (mNotificationLifetimeFinishedCallback != null) {
-                mNotificationLifetimeFinishedCallback.onSafeToRemove(entry.getKey());
-            }
+        if (mRemoteInputListener != null) {
+            mRemoteInputListener.onPanelCollapsed();
         }
-        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 @@
         if (entry == null) {
             return;
         }
-        final String key = entry.getKey();
-        if (isNotificationKeptForRemoteInputHistory(key)) {
-            mMainHandler.postDelayed(() -> {
-                if (isNotificationKeptForRemoteInputHistory(key)) {
-                    mNotificationLifetimeFinishedCallback.onSafeToRemove(key);
-                }
-            }, REMOTE_INPUT_KEPT_ENTRY_AUTO_CANCEL_DELAY);
+        if (mRemoteInputListener != null) {
+            mRemoteInputListener.releaseNotificationIfKeptForRemoteInputHistory(entry);
         }
     }
 
+    /** 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 @@
         }
     }
 
-    @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:");
-        pw.print("  mKeysKeptForRemoteInputHistory: ");
-        pw.println(mKeysKeptForRemoteInputHistory);
-        pw.print("  mEntriesKeptForRemoteInputActive: ");
-        pw.println(mEntriesKeptForRemoteInputActive);
+        if (mRemoteInputListener instanceof Dumpable) {
+            ((Dumpable) mRemoteInputListener).dump(fd, pw, args);
+        }
     }
 
     public void bindRow(ExpandableNotificationRow row) {
@@ -734,11 +671,6 @@
         return mInteractionHandler;
     }
 
-    @VisibleForTesting
-    public Set<NotificationEntry> getEntriesKeptForRemoteInputActive() {
-        return mEntriesKeptForRemoteInputActive;
-    }
-
     public boolean isRemoteInputActive() {
         return mRemoteInputController != null && mRemoteInputController.isRemoteInputActive();
     }
@@ -758,131 +690,6 @@
     }
 
     /**
-     * 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 @@
          */
         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);
+                }
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputController.java b/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputController.java
index 83701a0..cde3b0e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputController.java
@@ -299,6 +299,9 @@
         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.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilder.java b/packages/SystemUI/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilder.java
new file mode 100644
index 0000000..90abec1
--- /dev/null
+++ b/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());
+    }
+
+
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
index 7fc18b7..e288b1530 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
@@ -19,35 +19,44 @@
 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 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 @@
         }
     }
 
+    @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.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
index 1c9174a..bb697c3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarDependenciesModule.java
@@ -43,6 +43,7 @@
 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 @@
     @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 @@
             DumpManager dumpManager) {
         return new NotificationRemoteInputManager(
                 context,
+                featureFlags,
                 lockscreenUserManager,
                 smartReplyController,
                 notificationEntryManager,
+                rebuilder,
                 statusBarOptionalLazy,
                 statusBarStateController,
                 mainHandler,
@@ -166,10 +171,11 @@
     @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);
     }
 
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
index 60f44a0d..8bc41c2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationEntryManager.java
@@ -689,8 +689,9 @@
         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()) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollection.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollection.java
index dfdc548..4440c3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollection.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollection.java
@@ -47,6 +47,7 @@
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.app.Notification;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.service.notification.NotificationListenerService;
@@ -61,6 +62,7 @@
 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;
@@ -75,6 +77,7 @@
 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;
@@ -130,6 +133,7 @@
     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<>();
@@ -153,6 +157,7 @@
             SystemClock clock,
             FeatureFlags featureFlags,
             NotifCollectionLogger logger,
+            @Main Handler mainHandler,
             LogBufferEulogizer logBufferEulogizer,
             DumpManager dumpManager) {
         Assert.isMainThread();
@@ -160,6 +165,7 @@
         mClock = clock;
         mFeatureFlags = featureFlags;
         mLogger = logger;
+        mMainHandler = mainHandler;
         mEulogizer = logBufferEulogizer;
 
         dumpManager.registerDumpable(TAG, this);
@@ -441,7 +447,7 @@
             mEventQueue.add(new BindEntryEvent(entry, sbn));
 
             mLogger.logNotifUpdated(sbn.getKey());
-            mEventQueue.add(new EntryUpdatedEvent(entry));
+            mEventQueue.add(new EntryUpdatedEvent(entry, true /* fromSystem */));
         }
     }
 
@@ -788,6 +794,51 @@
 
     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,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifPipeline.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifPipeline.java
index 5777925..27ba4c2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifPipeline.java
+++ b/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.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;
@@ -223,6 +226,17 @@
     }
 
     /**
+     * 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
      * will return the current state of the list, which will likely be only partially-generated.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
index 33d9036..bf3e712 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
@@ -48,6 +48,7 @@
     conversationCoordinator: ConversationCoordinator,
     preparationCoordinator: PreparationCoordinator,
     mediaCoordinator: MediaCoordinator,
+    remoteInputCoordinator: RemoteInputCoordinator,
     shadeEventCoordinator: ShadeEventCoordinator,
     smartspaceDedupingCoordinator: SmartspaceDedupingCoordinator,
     viewConfigCoordinator: ViewConfigCoordinator,
@@ -73,6 +74,7 @@
         mCoordinators.add(bubbleCoordinator)
         mCoordinators.add(conversationCoordinator)
         mCoordinators.add(mediaCoordinator)
+        mCoordinators.add(remoteInputCoordinator)
         mCoordinators.add(shadeEventCoordinator)
         mCoordinators.add(viewConfigCoordinator)
         mCoordinators.add(visualStabilityCoordinator)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
new file mode 100644
index 0000000..3397815
--- /dev/null
+++ b/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)
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/InternalNotifUpdater.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/InternalNotifUpdater.java
new file mode 100644
index 0000000..5692fb2
--- /dev/null
+++ b/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);
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionListener.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionListener.java
index db0c174..68a346f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionListener.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionListener.java
@@ -56,6 +56,17 @@
     /**
      * 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.
      */
     default void onEntryUpdated(@NonNull NotificationEntry entry) {
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionLogger.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionLogger.kt
index f8a778d..1ebc66e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionLogger.kt
@@ -121,6 +121,26 @@
         })
     }
 
+    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
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifEvent.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifEvent.kt
index 2810b89..179e953 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifEvent.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifEvent.kt
@@ -64,10 +64,11 @@
 }
 
 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)
     }
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtender.kt
new file mode 100644
index 0000000..145c1e5
--- /dev/null
+++ b/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}") }
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallController.kt
index 3806d9a..31cc823 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallController.kt
@@ -86,7 +86,7 @@
         //
         // 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) {
diff --git a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java
index e639313a..5568f64 100644
--- a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java
@@ -17,6 +17,7 @@
 package com.android.systemui.util.sensors;
 
 import android.hardware.SensorManager;
+import android.os.Build;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -56,7 +57,7 @@
  */
 class ProximitySensorImpl implements ProximitySensor {
     private static final String TAG = "ProxSensor";
-    private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
+    private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG) || Build.IS_DEBUGGABLE;
     private static final long SECONDARY_PING_INTERVAL_MS = 5000;
 
     ThresholdSensor mPrimaryThresholdSensor;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagManagerTest.java
index b3c098c..8243be8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagManagerTest.java
@@ -20,13 +20,21 @@
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.verifyZeroInteractions;
+
+import android.content.Context;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -36,17 +44,33 @@
 import java.io.PrintWriter;
 import java.io.StringWriter;
 
+/**
+ * NOTE: This test is for the version of FeatureFlagManager in src-release, which should not allow
+ * overriding, and should never return any value other than the one provided as the default.
+ */
 @SmallTest
 public class FeatureFlagManagerTest extends SysuiTestCase {
     FeatureFlagManager mFeatureFlagManager;
 
+    @Mock private SystemPropertiesHelper mProps;
+    @Mock private Context mContext;
     @Mock private DumpManager mDumpManager;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
 
-        mFeatureFlagManager = new FeatureFlagManager(mDumpManager);
+        mFeatureFlagManager = new FeatureFlagManager(mProps, mContext, mDumpManager);
+    }
+
+    @After
+    public void onFinished() {
+        // SystemPropertiesHelper and Context are provided for constructor consistency with the
+        // debug version of the FeatureFlagManager, but should never be used.
+        verifyZeroInteractions(mProps, mContext);
+        // The dump manager should be registered with even for the release version, but that's it.
+        verify(mDumpManager).registerDumpable(anyString(), any());
+        verifyNoMoreInteractions(mDumpManager);
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
index 5944e9c..4ed7224 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
+++ b/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.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.RemoteInputHistoryExtender;
-import com.android.systemui.statusbar.NotificationRemoteInputManager.SmartReplyHistoryExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.LegacyRemoteInputLifetimeExtender.RemoteInputActiveExtender;
+import com.android.systemui.statusbar.NotificationRemoteInputManager.LegacyRemoteInputLifetimeExtender.RemoteInputHistoryExtender;
+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 @@
     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 void testShouldExtendLifetime_remoteInputActive() {
         when(mController.isRemoteInputActive(mEntry)).thenReturn(true);
 
+        assertTrue(mRemoteInputManager.isRemoteInputActive(mEntry));
         assertTrue(mRemoteInputActiveExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -128,6 +149,7 @@
         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 @@
         NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY = true;
         mEntry.lastRemoteInputSent = SystemClock.elapsedRealtime();
 
+        assertTrue(mRemoteInputManager.shouldKeepForRemoteInputHistory(mEntry));
         assertTrue(mRemoteInputHistoryExtender.shouldExtendLifetime(mEntry));
     }
 
@@ -144,6 +167,7 @@
         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 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 @@
                 DumpManager dumpManager) {
             super(
                     context,
+                    featureFlags,
                     lockscreenUserManager,
                     smartReplyController,
                     notificationEntryManager,
+                    rebuilder,
                     statusBarOptionalLazy,
                     statusBarStateController,
                     mainHandler,
@@ -297,14 +223,28 @@
             mRemoteInputController = controller;
         }
 
+        @NonNull
         @Override
-        protected void addLifetimeExtenders() {
-            mRemoteInputActiveExtender = new RemoteInputActiveExtender();
-            mRemoteInputHistoryExtender = new RemoteInputHistoryExtender();
-            mSmartReplyHistoryExtender = new SmartReplyHistoryExtender();
-            mLifetimeExtenders.add(mRemoteInputHistoryExtender);
-            mLifetimeExtenders.add(mSmartReplyHistoryExtender);
-            mLifetimeExtenders.add(mRemoteInputActiveExtender);
+        protected LegacyRemoteInputLifetimeExtender createLegacyRemoteInputLifetimeExtender(
+                Handler mainHandler,
+                NotificationEntryManager notificationEntryManager,
+                SmartReplyController smartReplyController) {
+            mLegacyRemoteInputLifetimeExtender = new FakeLegacyRemoteInputLifetimeExtender();
+            return mLegacyRemoteInputLifetimeExtender;
         }
+
+        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);
+            }
+        }
+
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java
new file mode 100644
index 0000000..ce11d6a
--- /dev/null
+++ b/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));
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
index 837d71f..99c965a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
@@ -39,6 +39,7 @@
 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 @@
         mDependency.injectTestDependency(NotificationEntryManager.class,
                 mNotificationEntryManager);
 
-        mSmartReplyController = new SmartReplyController(mNotificationEntryManager,
-                mIStatusBarService, mClickNotifier);
+        mSmartReplyController = new SmartReplyController(
+                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,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionTest.java
index ebeb591..f08a74a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionTest.java
@@ -35,6 +35,7 @@
 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 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.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 @@
     @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 @@
                 mClock,
                 mFeatureFlags,
                 mLogger,
+                mMainHandler,
                 mEulogizer,
                 mock(DumpManager.class));
         mCollection.attach(mGroupCoalescer);
@@ -1322,6 +1327,78 @@
         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)
@@ -1372,6 +1449,11 @@
         }
 
         @Override
+        public void onEntryUpdated(NotificationEntry entry, boolean fromSystem) {
+            onEntryUpdated(entry);
+        }
+
+        @Override
         public void onEntryRemoved(NotificationEntry entry, int reason) {
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
new file mode 100644
index 0000000..0ce6ada
--- /dev/null
+++ b/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()
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
new file mode 100644
index 0000000..37ad835
--- /dev/null
+++ b/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)
+        }
+    }
+}
diff --git a/services/companion/Android.bp b/services/companion/Android.bp
index 4ae9365..e3926b4 100644
--- a/services/companion/Android.bp
+++ b/services/companion/Android.bp
@@ -9,7 +9,10 @@
 
 filegroup {
     name: "services.companion-sources",
-    srcs: ["java/**/*.java"],
+    srcs: [
+        "java/**/*.java",
+        "java/**/*.proto",
+    ],
     path: "java",
     visibility: ["//frameworks/base/services"],
 }
@@ -17,6 +20,9 @@
 java_library_static {
     name: "services.companion",
     defaults: ["platform_service_defaults"],
+    proto: {
+        type: "stream",
+    },
     srcs: [":services.companion-sources"],
     libs: ["services.core"],
 }
diff --git a/services/companion/java/com/android/server/companion/proto/companion_apps_permissions.proto b/services/companion/java/com/android/server/companion/proto/companion_apps_permissions.proto
new file mode 100644
index 0000000..b786bcc
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/proto/companion_apps_permissions.proto
@@ -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.
+ */
+
+syntax = "proto3";
+
+option java_multiple_files = true;
+
+/* Represents granted permissions of a list of apps */
+message CompanionAppsPermissions {
+  // granted permissions of apps
+  repeated AppPermissions appPermissions = 1;
+
+  /* Represents the granted permissions of an app */
+  message AppPermissions {
+    // package name of the app
+    string packageName = 1;
+
+    // signing certificates used to sign the APK contents of this app
+    bytes certificates = 2;
+
+    // granted permissions
+    repeated string permission = 3;
+  }
+}
diff --git a/services/companion/java/com/android/server/companion/proto/companion_message.proto b/services/companion/java/com/android/server/companion/proto/companion_message.proto
new file mode 100644
index 0000000..2309be3
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/proto/companion_message.proto
@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+syntax = "proto3";
+
+option java_multiple_files = true;
+
+/* Represents a message between companion devices */
+message CompanionMessage {
+  // id of the message
+  int32 messageId = 1;
+
+  // type of the message
+  CompanionMessageType type = 2;
+
+  // data contained in the message
+  bytes data = 3;
+
+  // types of CompanionMessage
+  enum CompanionMessageType {
+    // default value for proto3
+    UNKNOWN = 0;
+
+    // handshake message to establish secure channel
+    SECURE_CHANNEL_HANDSHAKE = 1;
+
+    // permission sync
+    PERMISSION_SYNC = 2;
+  }
+}
diff --git a/services/core/java/com/android/server/apphibernation/AppHibernationManagerInternal.java b/services/core/java/com/android/server/apphibernation/AppHibernationManagerInternal.java
index b0335fe..a3c9612 100644
--- a/services/core/java/com/android/server/apphibernation/AppHibernationManagerInternal.java
+++ b/services/core/java/com/android/server/apphibernation/AppHibernationManagerInternal.java
@@ -43,4 +43,9 @@
      * @see AppHibernationService#setHibernatingGlobally
      */
     public abstract void setHibernatingGlobally(String packageName, boolean isHibernating);
+
+    /**
+     * @see AppHibernationService#isOatArtifactDeletionEnabled
+     */
+    public abstract boolean isOatArtifactDeletionEnabled();
 }
diff --git a/services/core/java/com/android/server/apphibernation/AppHibernationService.java b/services/core/java/com/android/server/apphibernation/AppHibernationService.java
index bd066ff..4d025c9 100644
--- a/services/core/java/com/android/server/apphibernation/AppHibernationService.java
+++ b/services/core/java/com/android/server/apphibernation/AppHibernationService.java
@@ -200,6 +200,14 @@
     }
 
     /**
+     * Whether global hibernation should delete ART ahead-of-time compilation artifacts and prevent
+     * package manager from re-optimizing the APK.
+     */
+    private boolean isOatArtifactDeletionEnabled() {
+        return mOatArtifactDeletionEnabled;
+    }
+
+    /**
      * Whether a package is hibernating for a given user.
      *
      * @param packageName the package to check
@@ -730,6 +738,11 @@
         public boolean isHibernatingGlobally(String packageName) {
             return mService.isHibernatingGlobally(packageName);
         }
+
+        @Override
+        public boolean isOatArtifactDeletionEnabled() {
+            return mService.isOatArtifactDeletionEnabled();
+        }
     }
 
     private final AppHibernationServiceStub mServiceStub = new AppHibernationServiceStub(this);
diff --git a/services/core/java/com/android/server/pm/DexOptHelper.java b/services/core/java/com/android/server/pm/DexOptHelper.java
index 9390284..5c3a991 100644
--- a/services/core/java/com/android/server/pm/DexOptHelper.java
+++ b/services/core/java/com/android/server/pm/DexOptHelper.java
@@ -49,8 +49,6 @@
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.logging.MetricsLogger;
-import com.android.server.apphibernation.AppHibernationManagerInternal;
-import com.android.server.apphibernation.AppHibernationService;
 import com.android.server.pm.dex.DexManager;
 import com.android.server.pm.dex.DexoptOptions;
 import com.android.server.pm.parsing.pkg.AndroidPackage;
@@ -171,7 +169,7 @@
                 }
             }
 
-            if (!PackageDexOptimizer.canOptimizePackage(pkg)) {
+            if (!mPm.mPackageDexOptimizer.canOptimizePackage(pkg)) {
                 if (DEBUG_DEXOPT) {
                     Log.i(TAG, "Skipping update of non-optimizable app " + pkg.getPackageName());
                 }
@@ -291,16 +289,11 @@
         ArraySet<String> pkgs = new ArraySet<>();
         synchronized (mPm.mLock) {
             for (AndroidPackage p : mPm.mPackages.values()) {
-                if (PackageDexOptimizer.canOptimizePackage(p)) {
+                if (mPm.mPackageDexOptimizer.canOptimizePackage(p)) {
                     pkgs.add(p.getPackageName());
                 }
             }
         }
-        if (AppHibernationService.isAppHibernationEnabled()) {
-            AppHibernationManagerInternal appHibernationManager =
-                    mPm.mInjector.getLocalService(AppHibernationManagerInternal.class);
-            pkgs.removeIf(pkgName -> appHibernationManager.isHibernatingGlobally(pkgName));
-        }
         return pkgs;
     }
 
diff --git a/services/core/java/com/android/server/pm/OtaDexoptService.java b/services/core/java/com/android/server/pm/OtaDexoptService.java
index 68801d6..9122221f 100644
--- a/services/core/java/com/android/server/pm/OtaDexoptService.java
+++ b/services/core/java/com/android/server/pm/OtaDexoptService.java
@@ -387,7 +387,7 @@
             }
 
             // Does the package have code? If not, there won't be any artifacts.
-            if (!PackageDexOptimizer.canOptimizePackage(pkg)) {
+            if (!mPackageManagerService.mPackageDexOptimizer.canOptimizePackage(pkg)) {
                 continue;
             }
             if (pkg.getPath() == null) {
diff --git a/services/core/java/com/android/server/pm/PackageDexOptimizer.java b/services/core/java/com/android/server/pm/PackageDexOptimizer.java
index 7739f2f..cac1978 100644
--- a/services/core/java/com/android/server/pm/PackageDexOptimizer.java
+++ b/services/core/java/com/android/server/pm/PackageDexOptimizer.java
@@ -64,7 +64,10 @@
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.IndentingPrintWriter;
+import com.android.server.LocalServices;
+import com.android.server.apphibernation.AppHibernationManagerInternal;
 import com.android.server.pm.Installer.InstallerException;
 import com.android.server.pm.dex.ArtManagerService;
 import com.android.server.pm.dex.ArtStatsLogUtils;
@@ -134,16 +137,24 @@
     private volatile boolean mSystemReady;
 
     private final ArtStatsLogger mArtStatsLogger = new ArtStatsLogger();
+    private final Injector mInjector;
+
 
     private static final Random sRandom = new Random();
 
     PackageDexOptimizer(Installer installer, Object installLock, Context context,
             String wakeLockTag) {
-        this.mInstaller = installer;
-        this.mInstallLock = installLock;
+        this(new Injector() {
+            @Override
+            public AppHibernationManagerInternal getAppHibernationManagerInternal() {
+                return LocalServices.getService(AppHibernationManagerInternal.class);
+            }
 
-        PowerManager powerManager = (PowerManager)context.getSystemService(Context.POWER_SERVICE);
-        mDexoptWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, wakeLockTag);
+            @Override
+            public PowerManager getPowerManager(Context context) {
+                return context.getSystemService(PowerManager.class);
+            }
+        }, installer, installLock, context, wakeLockTag);
     }
 
     protected PackageDexOptimizer(PackageDexOptimizer from) {
@@ -151,9 +162,21 @@
         this.mInstallLock = from.mInstallLock;
         this.mDexoptWakeLock = from.mDexoptWakeLock;
         this.mSystemReady = from.mSystemReady;
+        this.mInjector = from.mInjector;
     }
 
-    static boolean canOptimizePackage(AndroidPackage pkg) {
+    @VisibleForTesting
+    PackageDexOptimizer(@NonNull Injector injector, Installer installer, Object installLock,
+            Context context, String wakeLockTag) {
+        this.mInstaller = installer;
+        this.mInstallLock = installLock;
+
+        PowerManager powerManager = injector.getPowerManager(context);
+        mDexoptWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, wakeLockTag);
+        mInjector = injector;
+    }
+
+    boolean canOptimizePackage(AndroidPackage pkg) {
         // We do not dexopt a package with no code.
         // Note that the system package is marked as having no code, however we can
         // still optimize it via dexoptSystemServerPath.
@@ -161,6 +184,13 @@
             return false;
         }
 
+        // We do not dexopt unused packages.
+        AppHibernationManagerInternal ahm = mInjector.getAppHibernationManagerInternal();
+        if (ahm.isHibernatingGlobally(pkg.getPackageName())
+                && ahm.isOatArtifactDeletionEnabled()) {
+            return false;
+        }
+
         return true;
     }
 
@@ -1000,4 +1030,13 @@
     private Installer getInstallerWithoutLock() {
         return mInstaller;
     }
+
+    /**
+     * Injector for {@link PackageDexOptimizer} dependencies
+     */
+    interface Injector {
+        AppHibernationManagerInternal getAppHibernationManagerInternal();
+
+        PowerManager getPowerManager(Context context);
+    }
 }
diff --git a/services/core/java/com/android/server/policy/WindowManagerPolicy.java b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
index 0992d00..5ec2f83 100644
--- a/services/core/java/com/android/server/policy/WindowManagerPolicy.java
+++ b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
@@ -523,43 +523,44 @@
             case TYPE_PHONE:
                 return  3;
             case TYPE_SEARCH_BAR:
-            case TYPE_VOICE_INTERACTION_STARTING:
                 return  4;
-            case TYPE_VOICE_INTERACTION:
-                // voice interaction layer is almost immediately above apps.
-                return  5;
             case TYPE_INPUT_CONSUMER:
-                return  6;
+                return  5;
             case TYPE_SYSTEM_DIALOG:
-                return  7;
+                return  6;
             case TYPE_TOAST:
                 // toasts and the plugged-in battery thing
-                return  8;
+                return  7;
             case TYPE_PRIORITY_PHONE:
                 // SIM errors and unlock.  Not sure if this really should be in a high layer.
-                return  9;
+                return  8;
             case TYPE_SYSTEM_ALERT:
                 // like the ANR / app crashed dialogs
                 // Type is deprecated for non-system apps. For system apps, this type should be
                 // in a higher layer than TYPE_APPLICATION_OVERLAY.
-                return  canAddInternalSystemWindow ? 13 : 10;
+                return  canAddInternalSystemWindow ? 12 : 9;
             case TYPE_APPLICATION_OVERLAY:
-                return  12;
+                return  11;
             case TYPE_INPUT_METHOD:
                 // on-screen keyboards and other such input method user interfaces go here.
-                return  15;
+                return  13;
             case TYPE_INPUT_METHOD_DIALOG:
                 // on-screen keyboards and other such input method user interfaces go here.
-                return  16;
+                return  14;
             case TYPE_STATUS_BAR:
-                return  17;
+                return  15;
             case TYPE_STATUS_BAR_ADDITIONAL:
-                return  18;
+                return  16;
             case TYPE_NOTIFICATION_SHADE:
-                return  19;
+                return  17;
             case TYPE_STATUS_BAR_SUB_PANEL:
-                return  20;
+                return  18;
             case TYPE_KEYGUARD_DIALOG:
+                return  19;
+            case TYPE_VOICE_INTERACTION_STARTING:
+                return  20;
+            case TYPE_VOICE_INTERACTION:
+                // voice interaction layer should show above the lock screen.
                 return  21;
             case TYPE_VOLUME_OVERLAY:
                 // the on-screen volume indicator and controller shown when the user
@@ -568,7 +569,7 @@
             case TYPE_SYSTEM_OVERLAY:
                 // the on-screen volume indicator and controller shown when the user
                 // changes the device volume
-                return  canAddInternalSystemWindow ? 23 : 11;
+                return  canAddInternalSystemWindow ? 23 : 10;
             case TYPE_NAVIGATION_BAR:
                 // the navigation bar, if available, shows atop most things
                 return  24;
@@ -581,7 +582,7 @@
                 return  26;
             case TYPE_SYSTEM_ERROR:
                 // system-level error dialogs
-                return  canAddInternalSystemWindow ? 27 : 10;
+                return  canAddInternalSystemWindow ? 27 : 9;
             case TYPE_MAGNIFICATION_OVERLAY:
                 // used to highlight the magnified portion of a display
                 return  28;
diff --git a/services/core/java/com/android/server/wm/TaskDisplayArea.java b/services/core/java/com/android/server/wm/TaskDisplayArea.java
index 901f14f..d4035b5 100644
--- a/services/core/java/com/android/server/wm/TaskDisplayArea.java
+++ b/services/core/java/com/android/server/wm/TaskDisplayArea.java
@@ -42,8 +42,6 @@
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ROOT_TASK;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
-import static java.lang.Integer.MIN_VALUE;
-
 import android.annotation.ColorInt;
 import android.annotation.Nullable;
 import android.app.ActivityOptions;
@@ -907,15 +905,11 @@
         float r = ((color >> 16) & 0xff) / 255.0f;
         float g = ((color >>  8) & 0xff) / 255.0f;
         float b = ((color >>  0) & 0xff) / 255.0f;
-        float a = ((color >> 24) & 0xff) / 255.0f;
 
         mColorLayerCounter++;
 
-        getPendingTransaction().setLayer(mColorBackgroundLayer, MIN_VALUE)
+        getPendingTransaction()
                 .setColor(mColorBackgroundLayer, new float[]{r, g, b})
-                .setAlpha(mColorBackgroundLayer, a)
-                .setWindowCrop(mColorBackgroundLayer, getSurfaceWidth(), getSurfaceHeight())
-                .setPosition(mColorBackgroundLayer, 0, 0)
                 .show(mColorBackgroundLayer);
 
         scheduleAnimation();
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/PackageManagerServiceHibernationTests.kt b/services/tests/mockingservicestests/src/com/android/server/pm/PackageManagerServiceHibernationTests.kt
index 3ee2348..15acdac 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/PackageManagerServiceHibernationTests.kt
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/PackageManagerServiceHibernationTests.kt
@@ -16,8 +16,10 @@
 
 package com.android.server.pm
 
+import android.content.Context
 import android.os.Build
 import android.os.Handler
+import android.os.PowerManager
 import android.provider.DeviceConfig
 import android.provider.DeviceConfig.NAMESPACE_APP_HIBERNATION
 import android.testing.AndroidTestingRunner
@@ -55,6 +57,8 @@
 
     @Mock
     lateinit var appHibernationManager: AppHibernationManagerInternal
+    @Mock
+    lateinit var powerManager: PowerManager
 
     @Before
     @Throws(Exception::class)
@@ -68,6 +72,24 @@
             .thenReturn(appHibernationManager)
         whenever(rule.mocks().injector.handler)
             .thenReturn(Handler(TestableLooper.get(this).looper))
+        val injector = object : PackageDexOptimizer.Injector {
+            override fun getAppHibernationManagerInternal(): AppHibernationManagerInternal {
+                return appHibernationManager
+            }
+
+            override fun getPowerManager(context: Context?): PowerManager {
+                return powerManager
+            }
+        }
+        val packageDexOptimizer = PackageDexOptimizer(
+            injector,
+            rule.mocks().installer,
+            rule.mocks().installLock,
+            rule.mocks().context,
+            "*dexopt*")
+        whenever(rule.mocks().injector.packageDexOptimizer)
+            .thenReturn(packageDexOptimizer)
+        whenever(appHibernationManager.isOatArtifactDeletionEnabled).thenReturn(true)
     }
 
     @Test